libpipe 0.11__tar.gz → 0.12__tar.gz

{libpipe-0.11 → libpipe-0.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: libpipe
-Version: 0.11
+Version: 0.12
 Summary: Pipeline orchestration and node management utilities
 License-Expression: GPL-3.0-or-later
 License-File: LICENSE
@@ -64,7 +64,8 @@ on a file system shared with the compute nodes.
 
 With `slurm_nodes`, a pool on a login node requests one allocation of up to that many nodes
 when its first task starts, runs each task as a job step in it (`srun --jobid --exact`), and
-releases it at the end, or after 10 minutes without step. Slurm packs the steps on the nodes
+releases it at the end, or after 10 minutes without step. `slurm_nodes='2-10'` starts it as
+soon as 2 nodes are free, with up to 10 of them, as many as free at that moment. Slurm packs the steps on the nodes
 from their CPUs and memory, given by the category of the task; a step starts when the
 allocation has its CPUs free, so that no srun waits polling the controller:
 

{libpipe-0.11 → libpipe-0.12}/README.md

@@ -30,7 +30,8 @@ on a file system shared with the compute nodes.
 
 With `slurm_nodes`, a pool on a login node requests one allocation of up to that many nodes
 when its first task starts, runs each task as a job step in it (`srun --jobid --exact`), and
-releases it at the end, or after 10 minutes without step. Slurm packs the steps on the nodes
+releases it at the end, or after 10 minutes without step. `slurm_nodes='2-10'` starts it as
+soon as 2 nodes are free, with up to 10 of them, as many as free at that moment. Slurm packs the steps on the nodes
 from their CPUs and memory, given by the category of the task; a step starts when the
 allocation has its CPUs free, so that no srun waits polling the controller:
 

{libpipe-0.11 → libpipe-0.12}/libpipe/worker.py

@@ -48,8 +48,9 @@ DEFAULT_RESOURCES = {
 THREAD_VARS = ('OMP_NUM_THREADS', 'OPENBLAS_NUM_THREADS', 'MKL_NUM_THREADS')
 # The holding job of an allocation ends after this many minutes without job step
 ALLOC_IDLE_MINUTES = 10
-# Allocations requested in a row before giving up, when they end before running
-ALLOC_MAX_REQUESTS = 3
+# Final states of an allocation that stop the pool instead of requesting another one: each new
+# allocation would land on other nodes, and suspend or end other jobs, for nothing
+ALLOC_FAILED_STATES = ('NODE_FAIL', 'BOOT_FAIL', 'FAILED', 'OUT_OF_MEMORY')
 # Launches of a task in an allocation, when srun ends before the command starts (the allocation
 # ended, or is being replaced)
 ALLOC_MAX_LAUNCHES = 3
@@ -57,6 +58,31 @@ ALLOC_MAX_LAUNCHES = 3
 ALLOC_WORKERS_PER_NODE = 32
 
 
+def node_range(slurm_nodes):
+    """(min, max) nodes of an allocation from SLURM_NODES: N, or 'MIN-MAX' to start it once MIN
+    nodes are free, with up to MAX of them; None without (None, 0)."""
+    if slurm_nodes is None or slurm_nodes == 0:
+        return None
+    if isinstance(slurm_nodes, tuple):
+        lo, hi = slurm_nodes
+    elif isinstance(slurm_nodes, int) and not isinstance(slurm_nodes, bool):
+        lo = hi = slurm_nodes
+    elif isinstance(slurm_nodes, str) and re.fullmatch(r'\d+(-\d+)?', slurm_nodes.strip()):
+        lo, _, hi = slurm_nodes.strip().partition('-')
+        lo, hi = int(lo), int(hi or lo)
+    else:
+        raise ValueError(f'slurm_nodes: {slurm_nodes!r} is neither a number of nodes nor MIN-MAX')
+    if not 1 <= lo <= hi:
+        raise ValueError(f'slurm_nodes: {slurm_nodes!r}, MIN-MAX needs 1 <= MIN <= MAX')
+    return lo, hi
+
+
+def nodes_option(nodes):
+    """The --nodes value of an allocation of (min, max) nodes."""
+    lo, hi = nodes
+    return f'{lo}-{hi}' if lo < hi else str(hi)
+
+
 def now_str():
     return datetime.datetime.now().isoformat(sep=' ', timespec='milliseconds')
 
@@ -560,10 +586,10 @@ SBATCH_TRANSIENT = ('timed out', 'temporarily', 'Unable to contact', 'Connection
 
 class SlurmClient:
     """Run each task with srun: a Slurm job of its own when the pool runs on a login node, a
-    job step when it runs inside an allocation (salloc, sbatch). With alloc_nodes, the client
-    requests an allocation of that many nodes and runs the tasks as job steps in it, packed on
-    its nodes from their CPUs and memory; the allocation ends with the pool, or after
-    ALLOC_IDLE_MINUTES without step.
+    job step when it runs inside an allocation (salloc, sbatch). With alloc_nodes (N or
+    'MIN-MAX', see node_range), the client requests an allocation of that many nodes and runs
+    the tasks as job steps in it, packed on its nodes from their CPUs and memory; the allocation
+    ends with the pool, or after ALLOC_IDLE_MINUTES without step.
 
     Options (partition, account, ...) come from slurm_options, or from the LIBPIPE_SLURM_OPTIONS
     environment variable; they apply to the jobs and the allocation, not to the job steps.
@@ -580,9 +606,10 @@ class SlurmClient:
         self.max_time = max_time
         self.exact = False
         self.steps = False
-        self.alloc_nodes = alloc_nodes
+        self.alloc_nodes = node_range(alloc_nodes)
         self.alloc_id = None
         self.alloc_refused = None
+        self.alloc_ran = False
         self.alloc_checked = 0
         self.alloc_lock = asyncio.Lock()
         # CPUs of the allocation and those free: a step starts when its CPUs are free, instead of
@@ -683,33 +710,47 @@ class SlurmClient:
         return process.returncode, out.decode().strip(), err.decode().strip()
 
     async def ensure_allocation(self, check=False):
-        """Request the allocation, or a new one when it has ended, wait for it to run and return
-        its id. check: ask Slurm for its state even if it was checked in the last 30 s."""
+        """Request the allocation, or a new one when it has ended after running normally (time
+        limit, idle), wait for it to run and return its id. An allocation that ends before
+        running, or in a failed state (node failure, ...), stops the pool. check: ask Slurm for
+        its state even if it was checked in the last 30 s."""
         async with self.alloc_lock:
             if self.alloc_refused:
                 raise AllocationRefused(self.alloc_refused)
             if check:
                 self.alloc_checked = 0
-            n_requested = 0
             while True:
                 if self.alloc_id and time.monotonic() - self.alloc_checked < 30:
                     return self.alloc_id
                 if self.alloc_id is None:
-                    if n_requested == ALLOC_MAX_REQUESTS:
-                        raise RuntimeError(f'{n_requested} allocations ended before running')
                     self.alloc_id = await self._submit_allocation()
-                    n_requested += 1
+                    self.alloc_ran = False
                 _, state, _ = await self._run('squeue', '-h', '-j', self.alloc_id, '-o', '%T')
                 if state == 'RUNNING':
                     self.alloc_checked = time.monotonic()
+                    self.alloc_ran = True
                     await self._set_cpus(self.alloc_id)
                     return self.alloc_id
-                if state not in ('PENDING', 'CONFIGURING'):
-                    print(f'[{self.job_name}] allocation {self.alloc_id} ended '
-                          f'({state or "gone"}): requesting a new one')
-                    self.alloc_id = None
+                if state in ('PENDING', 'CONFIGURING'):
+                    await asyncio.sleep(5)
                     continue
-                await asyncio.sleep(5)
+                final = await self._final_state(self.alloc_id) or state or 'gone'
+                if final.split()[0] in ALLOC_FAILED_STATES or not self.alloc_ran:
+                    when = '' if self.alloc_ran else ' before running'
+                    self.alloc_refused = (f'allocation {self.alloc_id} ended{when} ({final}): no '
+                                          f'new allocation is requested; check the nodes and the '
+                                          f'partition (sacct -j {self.alloc_id})')
+                    print(style(f'[{self.job_name}] {self.alloc_refused}', fg='red'))
+                    raise AllocationRefused(self.alloc_refused)
+                print(f'[{self.job_name}] allocation {self.alloc_id} ended ({final}): '
+                      'requesting a new one')
+                self.alloc_id = None
+
+    async def _final_state(self, job_id):
+        """State of an ended job from sacct, e.g. NODE_FAIL or COMPLETED; '' if unknown."""
+        returncode, out, _ = await self._run('sacct', '-n', '-X', '-P', '-j', job_id, '-o',
+                                             'State')
+        return out.splitlines()[0].strip() if returncode == 0 and out.strip() else ''
 
     async def _submit_allocation(self):
         # The holding job: ends after ALLOC_IDLE_MINUTES without step (other than its own)
@@ -722,7 +763,8 @@ class SlurmClient:
                                        for o in self.options):
             time_option = ['--time=24:00:00']
         # --no-kill: a node that fails ends only the steps on it, not the allocation
-        args = ['sbatch', '--parsable', f'--nodes={self.alloc_nodes}', '--exclusive', '--no-kill',
+        args = ['sbatch', '--parsable', f'--nodes={nodes_option(self.alloc_nodes)}', '--exclusive',
+                '--no-kill',
                 f'--job-name={self.job_name}:alloc', '--output=/dev/null', *time_option,
                 *self.options, '--wrap', holder]
         returncode, out, err = await self._run(*args)
@@ -734,8 +776,8 @@ class SlurmClient:
                 raise AllocationRefused(self.alloc_refused)
             raise RuntimeError(f'Requesting the allocation failed: {err}')
         job_id = out.split(';')[0]
-        print(f'[{self.job_name}] allocation {job_id} of {self.alloc_nodes} node(s): the tasks '
-              'run as job steps')
+        print(f'[{self.job_name}] allocation {job_id} of {nodes_option(self.alloc_nodes)} '
+              'node(s): the tasks run as job steps')
         return job_id
 
     async def _log_stdout(self, task, stream):
@@ -1016,8 +1058,9 @@ class WorkerPool:
                 '-p cpu -c 16 --mem=64G' (default: $LIBPIPE_SLURM_OPTIONS).
             kill_grace_s (int, optional): Time a task past max_time gets to stop after SIGTERM
                 before SIGKILL.
-            slurm_nodes (int, optional): For 'slurm' from a login node: run the tasks as job
-                steps of an allocation of up to that many nodes, instead of one job each.
+            slurm_nodes (int or str, optional): For 'slurm' from a login node: run the tasks as
+                job steps of an allocation of up to that many nodes, instead of one job each;
+                'MIN-MAX' starts it once MIN nodes are free, with up to MAX of them.
             resources (dict, optional): CPUs and memory of the task categories as job steps,
                 {category: {'cpus': .., 'mem': ..}}, updating DEFAULT_RESOURCES.
         """
@@ -1044,7 +1087,7 @@ class WorkerPool:
         self._waiters = collections.deque()
         self._cooldown_timer = None
         self.resources = resolve_resources(resources)
-        self.slurm_nodes = slurm_nodes
+        self.slurm_nodes = node_range(slurm_nodes)
         # Wall time, CPU time and peak memory of each task (libpipe.taskrun), with the Python
         # of the pool: on this machine or on the nodes of the cluster
         self.accounting = all(h in (localhost_shortname, 'localhost', SLURM_HOST) for h in hosts)
@@ -1055,9 +1098,9 @@ class WorkerPool:
                 client = LocalClient(localhost_shortname, force_sync=force_sync)
             elif host == SLURM_HOST:
                 client = SlurmClient(slurm_options, job_name=name, max_time=max_time,
-                                     alloc_nodes=slurm_nodes)
-                if slurm_nodes:
-                    n_workers = max(n_workers, ALLOC_WORKERS_PER_NODE * slurm_nodes)
+                                     alloc_nodes=self.slurm_nodes)
+                if self.slurm_nodes:
+                    n_workers = max(n_workers, ALLOC_WORKERS_PER_NODE * self.slurm_nodes[1])
             else:
                 client = Client(host, user, force_sync=force_sync, password=password,
                                 connect_timeout=connect_timeout)
@@ -1276,7 +1319,8 @@ class WorkerPool:
         # An allocation no larger than the number of tasks
         for client in {w.client for w in self.workers}:
             if isinstance(client, SlurmClient) and client.alloc_nodes:
-                client.alloc_nodes = min(self.slurm_nodes, len(self.tasks))
+                hi = min(self.slurm_nodes[1], len(self.tasks))
+                client.alloc_nodes = (min(self.slurm_nodes[0], hi), hi) if hi else None
 
         self.futures = [asyncio.ensure_future(task.run(self)) for task in self.tasks]
         self.futures.append(asyncio.ensure_future(self._monitor_queue()))

{libpipe-0.11 → libpipe-0.12}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "libpipe"
-version = "0.11"
+version = "0.12"
 description = "Pipeline orchestration and node management utilities"
 readme = "README.md"
 license = "GPL-3.0-or-later"