libpipe 0.10__tar.gz → 0.11__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: libpipe
-Version: 0.10
+Version: 0.11
 Summary: Pipeline orchestration and node management utilities
 License-Expression: GPL-3.0-or-later
 License-File: LICENSE

{libpipe-0.10 → libpipe-0.11}/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
@@ -47,6 +50,9 @@ THREAD_VARS = ('OMP_NUM_THREADS', 'OPENBLAS_NUM_THREADS', 'MKL_NUM_THREADS')
 ALLOC_IDLE_MINUTES = 10
 # Allocations requested in a row before giving up, when they end before running
 ALLOC_MAX_REQUESTS = 3
+# 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
 
@@ -55,6 +61,15 @@ 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 +210,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 +285,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 +296,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 +353,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,6 +550,14 @@ 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
@@ -552,6 +582,7 @@ class SlurmClient:
         self.steps = False
         self.alloc_nodes = alloc_nodes
         self.alloc_id = None
+        self.alloc_refused = None
         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 +623,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 +658,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,13 +682,18 @@ 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, 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."""
         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
+                    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')
@@ -662,7 +703,7 @@ class SlurmClient:
                 if state == 'RUNNING':
                     self.alloc_checked = time.monotonic()
                     await self._set_cpus(self.alloc_id)
-                    return
+                    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')
@@ -680,11 +721,17 @@ 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={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 '
@@ -705,20 +752,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 +789,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 +864,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 +965,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:
@@ -1237,6 +1321,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.11}/pyproject.toml

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