libpipe 0.10__tar.gz → 0.12__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: libpipe
-Version: 0.10
+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.10 → 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.10 → libpipe-0.12}/libpipe/worker.py

@@ -32,6 +32,9 @@ localhost_shortname = socket.gethostname().split('.', 1)[0]
 SLURM_HOST = 'slurm'
 SLURM_OPTIONS_ENV = 'LIBPIPE_SLURM_OPTIONS'
 SLURM_MARKER = '[libpipe] slurm job '
+# Lines of srun and slurmstepd naming a node that failed, or the node of a step it ended
+NODE_FAILURE_RE = re.compile(r'Node failure on (\S+)'
+                             r'|\bON (\S+) CANCELLED AT \S+ DUE TO NODE FAILURE')
 
 # CPUs and memory of the task categories, used when the tasks run as job steps of an allocation;
 # the categories other than heavy are capped at its values
@@ -45,16 +48,54 @@ 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
 # Workers per node of an allocation: Slurm holds the steps that do not fit yet
 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')
 
 
+def note_node_failure(task, line):
+    '''Record on task the node named by a node-failure line (bytes) of srun or slurmstepd.'''
+    if b'Node failure' not in line and b'NODE FAILURE' not in line:
+        return
+    m = NODE_FAILURE_RE.search(line.decode(errors='replace'))
+    if m:
+        task.failed_nodes.add(m.group(1) or m.group(2))
+
+
 def mem_mb(text):
     '''Memory such as 40GB, 40G or 512MB in MB (GB is 1024 MB, as in Slurm).'''
     m = re.fullmatch(r'\s*(\d+(?:\.\d+)?)\s*([KMGT]?)(?:i?B)?\s*', str(text), re.IGNORECASE)
@@ -195,6 +236,8 @@ class Task:
         self.start_time = None
         self.sigterm_time = None
         self.slurm_job_id = None
+        # Slurm nodes whose failure stopped this task
+        self.failed_nodes = set()
         # Resources of a job step, and the usage reported by libpipe.taskrun
         self.category = category
         self.label = label or category
@@ -268,8 +311,9 @@ class Task:
         if self.fd is not None:
             self.fd.write(f'[libpipe] {now_str()} {text}\n')
 
-    async def log_stream(self, stream, host, err=False, stamp_prefix=None):
-        """Log the lines of stream; those starting with stamp_prefix get the current time."""
+    async def log_stream(self, stream, host, err=False, stamp_prefix=None, on_line=None):
+        """Log the lines of stream; those starting with stamp_prefix get the current time.
+        on_line is called with each line (bytes)."""
         while True:
             try:
                 line = await stream.readline()
@@ -278,6 +322,8 @@ class Task:
                 continue
             if not line:
                 break
+            if on_line is not None:
+                on_line(line)
             if stamp_prefix and line.startswith(stamp_prefix):
                 line = f'{now_str()} '.encode() + line
             self.log(line, host, err=err)
@@ -333,6 +379,8 @@ class Task:
                     break
 
                 self.status = TaskStatus.ERROR
+                if getattr(worker.client, 'alloc_refused', None):
+                    break
                 if not ret_status and self.process is not None:
                     # The process had started and may still be running: running the task
                     # again could have two copies write the same outputs
@@ -528,12 +576,20 @@ class LocalClient:
         pass
 
 
+class AllocationRefused(RuntimeError):
+    """Slurm refused the allocation request (account, partition, QOS, ...): not retried."""
+
+
+# sbatch errors worth another request; any other rejection is a configuration error
+SBATCH_TRANSIENT = ('timed out', 'temporarily', 'Unable to contact', 'Connection refused')
+
+
 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.
@@ -550,8 +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
@@ -592,16 +650,21 @@ class SlurmClient:
             self.cpus_changed.notify_all()
 
     async def _acquire_cpus(self, cpus):
+        '''CPUs taken, and the job whose CPUs they are (None without CPU count).'''
         if self.cpus_total is None:
-            return 0
-        cpus = min(cpus, self.cpus_total)
+            return 0, None
         async with self.cpus_changed:
-            await self.cpus_changed.wait_for(lambda: self.cpus_free >= cpus)
+            await self.cpus_changed.wait_for(
+                lambda: self.cpus_total is None or self.cpus_free >= min(cpus, self.cpus_total))
+            if self.cpus_total is None:
+                return 0, None
+            cpus = min(cpus, self.cpus_total)
             self.cpus_free -= cpus
-        return cpus
+            return cpus, self.cpus_of
 
-    async def _release_cpus(self, cpus):
-        if cpus and self.cpus_total is not None:
+    async def _release_cpus(self, cpus, job_id):
+        # CPUs of an allocation that has been replaced are not those of the new one
+        if cpus and self.cpus_total is not None and job_id == self.cpus_of:
             async with self.cpus_changed:
                 self.cpus_free = min(self.cpus_total, self.cpus_free + cpus)
                 self.cpus_changed.notify_all()
@@ -622,13 +685,13 @@ class SlurmClient:
         has_time = any(o == '-t' or o.startswith(('-t', '--time')) for o in options)
         return [f'--time={int(-(-self.max_time // 60))}'] if self.max_time and not has_time else []
 
-    def srun_command(self, task):
+    def srun_command(self, task, alloc_id=None):
         args = ['srun', f'--job-name={self.job_name}:{task.name}', '--ntasks=1', '--nodes=1',
                 '--unbuffered']
         command, job_id, options = task.command, '${SLURM_JOB_ID}', self.options
         if self.steps:
-            if self.alloc_id:
-                args.append(f'--jobid={self.alloc_id}')
+            if alloc_id:
+                args.append(f'--jobid={alloc_id}')
             if self.exact:
                 args.append('--exact')
             if task.cpus:
@@ -646,29 +709,48 @@ class SlurmClient:
         out, err = await process.communicate()
         return process.returncode, out.decode().strip(), err.decode().strip()
 
-    async def ensure_allocation(self):
-        """Request the allocation, or a new one when it has ended, and wait for it to run."""
+    async def ensure_allocation(self, check=False):
+        """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:
-            n_requested = 0
+            if self.alloc_refused:
+                raise AllocationRefused(self.alloc_refused)
+            if check:
+                self.alloc_checked = 0
             while True:
                 if self.alloc_id and time.monotonic() - self.alloc_checked < 30:
-                    return
+                    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
-                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
+                    return self.alloc_id
+                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)
@@ -680,15 +762,22 @@ class SlurmClient:
         if not time_option and not any(o == '-t' or o.startswith(('-t', '--time'))
                                        for o in self.options):
             time_option = ['--time=24:00:00']
-        args = ['sbatch', '--parsable', f'--nodes={self.alloc_nodes}', '--exclusive',
+        # --no-kill: a node that fails ends only the steps on it, not the allocation
+        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)
         if returncode:
+            if not any(t in err for t in SBATCH_TRANSIENT):
+                # Every task would ask again and be refused the same way: say it once
+                self.alloc_refused = f'Slurm refused the allocation: {err}'
+                print(style(f'[{self.job_name}] {self.alloc_refused}; no task is run', fg='red'))
+                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):
@@ -705,20 +794,35 @@ class SlurmClient:
             if task.slurm_job_id is None and text.startswith(SLURM_MARKER):
                 task.slurm_job_id = text[len(SLURM_MARKER):].split()[0]
                 task.start_time = time.time()
+            note_node_failure(task, line)
             task.log(line, self.host)
 
     async def execute(self, task):
+        for launch in range(1, ALLOC_MAX_LAUNCHES + 1):
+            await self._launch(task, check=launch > 1)
+            # Without the marker line, the command did not start: in an allocation, srun ended
+            # because the allocation did (or a new one is not running yet), and the task can run
+            # on the next one
+            if (not self.alloc_nodes or task.slurm_job_id is not None or task.returncode == 0
+                    or task.sigterm_time is not None or launch == ALLOC_MAX_LAUNCHES):
+                return
+            task.log_event(f'the command did not start: launching it again '
+                           f'({launch + 1}/{ALLOC_MAX_LAUNCHES})')
+
+    async def _launch(self, task, check=False):
         task.slurm_job_id = None
-        if self.alloc_nodes:
-            await self.ensure_allocation()
-        cpus = 0
+        alloc_id = await self.ensure_allocation(check) if self.alloc_nodes else None
+        cpus, cpus_of = 0, None
         if self.steps and self.cpus_total is not None:
             task.log_event(f'waiting for {min(task.cpus or 1, self.cpus_total)} CPUs '
                            f'({self.cpus_free} of {self.cpus_total} free)')
-            cpus = await self._acquire_cpus(task.cpus or 1)
+            cpus, cpus_of = await self._acquire_cpus(task.cpus or 1)
+            if self.alloc_nodes:
+                # The allocation may have been replaced while waiting: the CPUs are of the new one
+                alloc_id = cpus_of or alloc_id
         try:
             process = await asyncio.create_subprocess_shell(
-                self.srun_command(task), stdout=asyncio.subprocess.PIPE,
+                self.srun_command(task, alloc_id), stdout=asyncio.subprocess.PIPE,
                 stderr=asyncio.subprocess.PIPE, preexec_fn=os.setpgrp, env=self.env)
             task.log_event(f'srun started{self._cpus_state()}')
             task.set_process(Process(self, process))
@@ -727,12 +831,13 @@ class SlurmClient:
 
             await asyncio.gather(
                 self._log_stdout(task, process.stdout),
-                task.log_stream(process.stderr, self.host, err=True, stamp_prefix=b'srun: '),
+                task.log_stream(process.stderr, self.host, err=True, stamp_prefix=b'srun: ',
+                                on_line=lambda line: note_node_failure(task, line)),
             )
             task.returncode = await process.wait()
             task.log_event(f'srun exited with code {task.returncode}')
         finally:
-            await self._release_cpus(cpus)
+            await self._release_cpus(cpus, cpus_of)
             if cpus:
                 task.log_event(f'{cpus} CPUs released{self._cpus_state()}')
 
@@ -801,6 +906,11 @@ class Worker:
         except asyncio.CancelledError:
             self.status = WorkerStatus.IDLE
             raise
+        except AllocationRefused as exc:
+            # Said once on the console by the client: in the task log only
+            task.log_event(f'not run: {exc}')
+            self.status = WorkerStatus.IDLE
+            return False, None
         except Exception as exc:
             if isinstance(exc, ConnectionError):
                 err_msg = f'Error connecting client: {exc!s}\n'
@@ -897,6 +1007,22 @@ class WorkerPoolReport:
                 for label, peak in sorted(peaks.items())
                 for category, mem in [requests[label]] if peak > mem]
 
+    def node_failure_lines(self):
+        '''One line per Slurm node whose failure stopped tasks, with these tasks.'''
+        tasks = collections.defaultdict(list)
+        for t in self.all_tasks:
+            for node in sorted(getattr(t, 'failed_nodes', ())):
+                tasks[node].append(t.name)
+        return [f'{node}: {len(names)} task(s) stopped ({", ".join(names)})'
+                for node, names in sorted(tasks.items())]
+
+    def print_node_failures(self):
+        lines = self.node_failure_lines()
+        if lines:
+            secho('\nSlurm node failures (run these tasks again):', fg='red')
+            print('\n'.join(lines))
+            print()
+
     def print_usage(self):
         usage = self.usage_lines()
         if not usage:
@@ -932,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.
         """
@@ -960,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)
@@ -971,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)
@@ -1192,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()))
@@ -1237,6 +1365,7 @@ class WorkerPool:
         report = self.last_report
         # At the end of the pool, before the next tasks fill the console
         report.print_usage()
+        report.print_node_failures()
 
         if print_summary:
             print("\n=== WorkerPool summary ===")

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

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