libpipe 0.8__tar.gz → 0.9__tar.gz

{libpipe-0.8 → libpipe-0.9}/PKG-INFO

@@ -1,25 +1,32 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: libpipe
-Version: 0.8
+Version: 0.9
 Summary: Pipeline orchestration and node management utilities
-License: GPL-3.0-or-later
+License-Expression: GPL-3.0-or-later
+License-File: LICENSE
 Author: Florent Mertens
 Author-email: florent.mertens@gmail.com
 Requires-Python: >=3.10,<4.0
-Classifier: License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Provides-Extra: dev
+Provides-Extra: docs
 Requires-Dist: asyncssh (>=2.0)
 Requires-Dist: click (>=8.0)
+Requires-Dist: myst-parser (>=2) ; extra == "docs"
 Requires-Dist: progressbar2 (>=3.40)
+Requires-Dist: pydata-sphinx-theme (>=0.15) ; extra == "docs"
 Requires-Dist: pytest (>=8.0) ; extra == "dev"
 Requires-Dist: pytest-cov (>=5.0) ; extra == "dev"
 Requires-Dist: ruff (==0.16.1) ; extra == "dev"
+Requires-Dist: sphinx (>=7) ; extra == "docs"
+Requires-Dist: sphinx-click (>=6) ; extra == "docs"
 Requires-Dist: toml (>=0.10)
+Project-URL: Documentation, https://libpipe.readthedocs.io
 Project-URL: Homepage, https://gitlab.com/flomertens/libpipe
 Project-URL: Issues, https://gitlab.com/flomertens/libpipe/-/issues
 Project-URL: Repository, https://gitlab.com/flomertens/libpipe
@@ -53,6 +60,26 @@ on a login node, as a job step inside an allocation (`salloc`, `sbatch`).
 account, resources) come from `slurm_options`, or from the `LIBPIPE_SLURM_OPTIONS`
 environment variable; `max_time` becomes the job time limit when the options set none.
 Host constraints (`run_on_host`) do not apply: Slurm places the tasks, so the data must be
-on a file system shared with the compute nodes. Give the memory (`--mem`) in the options:
-inside an allocation, a step without it can claim all the memory of its node.
+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
+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:
+
+    pool = worker.WorkerPool(['slurm'], name='Cubes', slurm_nodes=4,
+                             resources={'heavy': {'cpus': 12, 'mem': '40GB'}})
+    pool.add('pstool gen_vis_cube ...', category='small', label='gen_vis_cube')
+
+The categories (`DEFAULT_RESOURCES`) are `small` (1 CPU, 2 GB), `medium` (4, 8 GB),
+`high_mem` (2, 16 GB), `cpu` (12, 8 GB) and `heavy` (12, 40 GB, the default), the others capped
+at `heavy`; a step also gets the matching `OMP_NUM_THREADS`. The options apply to the jobs and
+the allocation, not to the steps. Inside an allocation started by hand, the tasks are job
+steps with these resources as well.
+
+Slurm does not necessarily limit the memory of a job step: the memory of a category must stay
+above the peak memory of its tasks. The local and Slurm pools measure each task (wall time,
+CPU time, peak RSS, in its log file), and `WorkerPoolReport.usage_lines()` summarises them per
+task type (`label`).
 

libpipe-0.9/README.md

@@ -0,0 +1,50 @@
+Libpipe
+=======
+
+Common pipeline framework and library
+
+Installation
+------------
+
+libpipe can be installed via pip:
+
+    $ pip install libpipe
+
+and requires Python 3.10 or higher.
+
+Worker pools on a Slurm cluster
+-------------------------------
+
+A `WorkerPool` runs its tasks locally, on other hosts over SSH, or through Slurm with the
+host `slurm`: each task is then started with `srun`, as a job of its own when the pool runs
+on a login node, as a job step inside an allocation (`salloc`, `sbatch`).
+
+    pool = worker.WorkerPool(['slurm'], name='Imaging', max_tasks_per_worker=8,
+                             slurm_options='-p regular --mem=200G --time=24:00:00')
+
+`max_tasks_per_worker` is the number of jobs in flight. The srun options (partition,
+account, resources) come from `slurm_options`, or from the `LIBPIPE_SLURM_OPTIONS`
+environment variable; `max_time` becomes the job time limit when the options set none.
+Host constraints (`run_on_host`) do not apply: Slurm places the tasks, so the data must be
+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
+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:
+
+    pool = worker.WorkerPool(['slurm'], name='Cubes', slurm_nodes=4,
+                             resources={'heavy': {'cpus': 12, 'mem': '40GB'}})
+    pool.add('pstool gen_vis_cube ...', category='small', label='gen_vis_cube')
+
+The categories (`DEFAULT_RESOURCES`) are `small` (1 CPU, 2 GB), `medium` (4, 8 GB),
+`high_mem` (2, 16 GB), `cpu` (12, 8 GB) and `heavy` (12, 40 GB, the default), the others capped
+at `heavy`; a step also gets the matching `OMP_NUM_THREADS`. The options apply to the jobs and
+the allocation, not to the steps. Inside an allocation started by hand, the tasks are job
+steps with these resources as well.
+
+Slurm does not necessarily limit the memory of a job step: the memory of a category must stay
+above the peak memory of its tasks. The local and Slurm pools measure each task (wall time,
+CPU time, peak RSS, in its log file), and `WorkerPoolReport.usage_lines()` summarises them per
+task type (`label`).

libpipe-0.9/libpipe/taskrun.py

@@ -0,0 +1,25 @@
+'''Run a command, then report its wall time, CPU time and peak memory on stderr, for the task
+accounting of WorkerPool.
+
+Usage: python taskrun.py COMMAND [ARGS...]'''
+
+import resource
+import subprocess
+import sys
+import time
+
+MARKER = '[libpipe] usage '
+
+
+def main():
+    start = time.monotonic()
+    returncode = subprocess.call(sys.argv[1:])
+    # Of the waited descendants: peak RSS of the largest one, CPU time of all
+    usage = resource.getrusage(resource.RUSAGE_CHILDREN)
+    print(f'{MARKER}wall={time.monotonic() - start:.1f} cpu={usage.ru_utime + usage.ru_stime:.1f} '
+          f'maxrss_kb={usage.ru_maxrss}', file=sys.stderr, flush=True)
+    sys.exit(returncode if returncode >= 0 else 128 - returncode)
+
+
+if __name__ == '__main__':
+    main()

{libpipe-0.8 → libpipe-0.9}/libpipe/worker.py

@@ -13,6 +13,7 @@ import re
 import shlex
 import signal
 import socket
+import subprocess
 import sys
 import time
 import warnings
@@ -22,6 +23,9 @@ import asyncssh
 import progressbar
 from click import secho, style
 
+from libpipe import taskrun
+from libpipe.taskrun import MARKER as USAGE_MARKER
+
 localhost_shortname = socket.gethostname().split('.', 1)[0]
 
 # Host name selecting the Slurm client: each task runs through srun
@@ -29,6 +33,57 @@ SLURM_HOST = 'slurm'
 SLURM_OPTIONS_ENV = 'LIBPIPE_SLURM_OPTIONS'
 SLURM_MARKER = '[libpipe] slurm job '
 
+# 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
+DEFAULT_RESOURCES = {
+    'small': {'cpus': 1, 'mem': '2GB'},
+    'medium': {'cpus': 4, 'mem': '8GB'},
+    'high_mem': {'cpus': 2, 'mem': '16GB'},
+    'cpu': {'cpus': 12, 'mem': '8GB'},
+    'heavy': {'cpus': 12, 'mem': '40GB'},
+}
+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
+# Workers per node of an allocation: Slurm holds the steps that do not fit yet
+ALLOC_WORKERS_PER_NODE = 32
+
+
+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)
+    if not m:
+        raise ValueError(f'{text!r}: expected a memory such as 40GB')
+    unit = 'KMGT'.index((m.group(2) or 'M').upper()) - 1
+    return max(1, int(float(m.group(1)) * 1024 ** unit))
+
+
+def resolve_resources(resources=None):
+    '''{category: (cpus, mem_mb)}: DEFAULT_RESOURCES updated with RESOURCES ({category: {'cpus':
+    .., 'mem': ..}}), the categories other than heavy capped at its values.'''
+    merged = {k: dict(v) for k, v in DEFAULT_RESOURCES.items()}
+    for category, res in (resources or {}).items():
+        merged.setdefault(category, {}).update(res)
+    resolved = {}
+    for category, res in merged.items():
+        if 'cpus' not in res or 'mem' not in res:
+            raise ValueError(f'Resources of the category {category}: give cpus and mem')
+        resolved[category] = (int(res['cpus']), mem_mb(res['mem']))
+    heavy = resolved['heavy']
+    return {c: r if c == 'heavy' else (min(r[0], heavy[0]), min(r[1], heavy[1]))
+            for c, r in resolved.items()}
+
+
+def slurm_cluster(options):
+    '''The cluster of --clusters/-M in a list of Slurm options, or None.'''
+    for i, option in enumerate(options):
+        m = re.fullmatch(r'(?:--clusters?|-M)=?(.*)', option)
+        if m:
+            return m.group(1) or (options[i + 1] if i + 1 < len(options) else None)
+    return None
+
 
 def get_event_loop():
     """The current event loop, or a new one when there is none or it is closed (e.g. after
@@ -74,11 +129,12 @@ def get_hosts(host_string):
 
 
 def get_worker_pool(name, nodes='localhost', max_concurrent=4, env_file=None, max_time=None,
-                    debug=False, dry_run=False, slurm_options=None):
+                    debug=False, dry_run=False, slurm_options=None, slurm_nodes=None,
+                    resources=None):
     hosts = get_hosts(nodes)
     return WorkerPool(hosts, name=name, max_tasks_per_worker=max_concurrent,
                       env_source_file=env_file, max_time=max_time, debug=debug, dry_run=dry_run,
-                      slurm_options=slurm_options)
+                      slurm_options=slurm_options, slurm_nodes=slurm_nodes, resources=resources)
 
 
 class TaskStatus(Enum):
@@ -114,7 +170,8 @@ class Task:
     """
 
     def __init__(self, name, command, output_file=None, done_callback=None, task_done_callback=None,
-                 run_on_host=None, n_try=3, pid_filename=None, keep_pid_file_on_error=False):
+                 run_on_host=None, n_try=3, pid_filename=None, keep_pid_file_on_error=False,
+                 category=None, label=None, cpus=None, mem_mb=None):
         self.name = name
         self.command = command
         self.output_file = output_file
@@ -134,6 +191,12 @@ class Task:
         self.start_time = None
         self.sigterm_time = None
         self.slurm_job_id = None
+        # Resources of a job step, and the usage reported by libpipe.taskrun
+        self.category = category
+        self.label = label or category
+        self.cpus = cpus
+        self.mem_mb = mem_mb
+        self.usage = None
 
     def done_callback(self):
         err_msg = ''
@@ -181,9 +244,17 @@ class Task:
         if isinstance(line, bytes):
             # Tools can print any byte: a decoding error must not interrupt the task
             line = line.decode(errors='replace')
+        if line.startswith(USAGE_MARKER):
+            with contextlib.suppress(ValueError):
+                self.usage = {k: float(v) for k, v in
+                              (f.split('=', 1) for f in line[len(USAGE_MARKER):].split())}
+            if self.fd is not None:
+                self.fd.write(line)
+            return
         if self.fd is not None:
             self.fd.write(line)
-        line = line.strip()
+        # Keep the indentation of the tool's output
+        line = line.rstrip()
         if err:
             line = style(line, fg='red')
         print(f'[{self.name}:{host}] {line}')
@@ -447,13 +518,17 @@ class LocalClient:
 
 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).
-
-    Options (resources, partition, account, ...) come from slurm_options, or from the
-    LIBPIPE_SLURM_OPTIONS environment variable. max_time becomes the job time limit.
+    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.
+
+    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.
+    max_time becomes the time limit of the jobs and job steps.
     """
 
-    def __init__(self, options=None, job_name='libpipe', max_time=None):
+    def __init__(self, options=None, job_name='libpipe', max_time=None, alloc_nodes=None):
         self.host = SLURM_HOST
         self.started = asyncio.Event()
         if options is None:
@@ -462,16 +537,63 @@ class SlurmClient:
         self.job_name = re.sub(r'\s+', '_', job_name)
         self.max_time = max_time
         self.exact = False
+        self.steps = False
+        self.alloc_nodes = alloc_nodes
+        self.alloc_id = 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
+        # waiting in srun, which polls the controller
+        self.cpus_total = None
+        self.cpus_free = 0
+        self.cpus_of = None
+        self.cpus_changed = asyncio.Condition()
+        # squeue, scancel and steps of a job of another cluster need SLURM_CLUSTERS
+        cluster = slurm_cluster(self.options)
+        self.env = dict(os.environ, SLURM_CLUSTERS=cluster) if cluster else None
 
     async def start(self):
         if self.started.is_set():
             return
-        # Inside an allocation, --exact (Slurm >= 21.08) gives each step only its resources,
-        # so that steps run side by side instead of one after the other
-        if os.environ.get('SLURM_JOB_ID'):
-            self.exact = await self._slurm_version() >= (21, 8)
+        # As job steps, --exact (Slurm >= 20.11) gives each step only its resources, so that
+        # steps run side by side instead of one after the other
+        in_allocation = bool(os.environ.get('SLURM_JOB_ID'))
+        if in_allocation or self.alloc_nodes:
+            self.exact = await self._slurm_version() >= (20, 11)
+        if self.alloc_nodes and (in_allocation or not self.exact):
+            if not in_allocation:
+                print(f'[{self.job_name}] Slurm < 20.11: one job per task instead of job steps')
+            self.alloc_nodes = None
+        self.steps = in_allocation or bool(self.alloc_nodes)
+        if in_allocation and self.exact:
+            await self._set_cpus(os.environ['SLURM_JOB_ID'])
         self.started.set()
 
+    async def _set_cpus(self, job_id):
+        if self.cpus_of == job_id:
+            return
+        _, out, _ = await self._run('squeue', '-h', '-j', job_id, '-o', '%C')
+        async with self.cpus_changed:
+            self.cpus_of = job_id
+            self.cpus_total = int(out) if out.isdigit() else None
+            self.cpus_free = self.cpus_total or 0
+            self.cpus_changed.notify_all()
+
+    async def _acquire_cpus(self, cpus):
+        if self.cpus_total is None:
+            return 0
+        cpus = min(cpus, self.cpus_total)
+        async with self.cpus_changed:
+            await self.cpus_changed.wait_for(lambda: self.cpus_free >= cpus)
+            self.cpus_free -= cpus
+        return cpus
+
+    async def _release_cpus(self, cpus):
+        if cpus and self.cpus_total is not None:
+            async with self.cpus_changed:
+                self.cpus_free = min(self.cpus_total, self.cpus_free + cpus)
+                self.cpus_changed.notify_all()
+
     @staticmethod
     async def _slurm_version():
         try:
@@ -484,17 +606,79 @@ class SlurmClient:
         m = re.search(r'(\d+)\.(\d+)', out.decode())
         return (int(m.group(1)), int(m.group(2))) if m else (0, 0)
 
+    def _time_option(self, options):
+        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):
         args = ['srun', f'--job-name={self.job_name}:{task.name}', '--ntasks=1', '--nodes=1',
                 '--unbuffered']
-        if self.exact:
-            args.append('--exact')
-        has_time = any(o == '-t' or o.startswith(('-t', '--time')) for o in self.options)
-        if self.max_time and not has_time:
-            args.append(f'--time={int(-(-self.max_time // 60))}')
-        args += self.options
-        marker = f'echo "{SLURM_MARKER}${{SLURM_JOB_ID}} on $(hostname -s)"'
-        return f'{shlex.join(args)} bash -c {shlex.quote(f"{marker}; {task.command}")}'
+        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 self.exact:
+                # --quiet: no "step creation disabled, retrying" while a step waits for its node
+                args += ['--exact', '--quiet']
+            if task.cpus:
+                args += [f'--cpus-per-task={task.cpus}', f'--mem={task.mem_mb}M']
+                threads = ' '.join(f'{v}={task.cpus}' for v in THREAD_VARS)
+                command = f'export {threads}; {command}'
+            job_id, options = '${SLURM_JOB_ID}.${SLURM_STEP_ID}', []
+        args += self._time_option(options) + options
+        marker = f'echo "{SLURM_MARKER}{job_id} on $(hostname -s)"'
+        return f'{shlex.join(args)} bash -c {shlex.quote(f"{marker}; {command}")}'
+
+    async def _run(self, *args):
+        process = await asyncio.create_subprocess_exec(
+            *args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=self.env)
+        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 with self.alloc_lock:
+            n_requested = 0
+            while True:
+                if self.alloc_id and time.monotonic() - self.alloc_checked < 30:
+                    return
+                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
+                _, state, _ = await self._run('squeue', '-h', '-j', self.alloc_id, '-o', '%T')
+                if state == 'RUNNING':
+                    self.alloc_checked = time.monotonic()
+                    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
+                    continue
+                await asyncio.sleep(5)
+
+    async def _submit_allocation(self):
+        # The holding job: ends after ALLOC_IDLE_MINUTES without step (other than its own)
+        holder = (f'unset SLURM_CLUSTERS; idle=0; while [ $idle -lt {ALLOC_IDLE_MINUTES} ]; do '
+                  'sleep 60; if squeue -h -s -j "$SLURM_JOB_ID" -o %i '
+                  '| grep -q -v -E "\\.(batch|extern)$"; then idle=0; else idle=$((idle + 1)); '
+                  'fi; done')
+        time_option = self._time_option(self.options) if self.max_time else []
+        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',
+                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:
+            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')
+        return job_id
 
     async def _log_stdout(self, task, stream):
         """Log the task output, recording the job id from the marker line."""
@@ -513,16 +697,22 @@ class SlurmClient:
 
     async def execute(self, task):
         task.slurm_job_id = None
-        process = await asyncio.create_subprocess_shell(
-            self.srun_command(task), stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.PIPE, preexec_fn=os.setpgrp)
-        task.set_process(Process(self, process))
-
-        await asyncio.gather(
-            self._log_stdout(task, process.stdout),
-            task.log_stream(process.stderr, self.host, err=True),
-        )
-        task.returncode = await process.wait()
+        if self.alloc_nodes:
+            await self.ensure_allocation()
+        cpus = await self._acquire_cpus(task.cpus or 1) if self.steps else 0
+        try:
+            process = await asyncio.create_subprocess_shell(
+                self.srun_command(task), stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE, preexec_fn=os.setpgrp, env=self.env)
+            task.set_process(Process(self, process))
+
+            await asyncio.gather(
+                self._log_stdout(task, process.stdout),
+                task.log_stream(process.stderr, self.host, err=True),
+            )
+            task.returncode = await process.wait()
+        finally:
+            await self._release_cpus(cpus)
 
     async def terminate(self, task, sig=signal.SIGTERM):
         job_id = getattr(task, 'slurm_job_id', None)
@@ -530,7 +720,7 @@ class SlurmClient:
             task.log(f'Cancelling Slurm job {job_id} ({sig.name})\n', self.host, err=True)
             args = ['scancel'] + (['-s', 'KILL'] if sig == signal.SIGKILL else []) + [job_id]
             with contextlib.suppress(Exception):
-                process = await asyncio.create_subprocess_exec(*args)
+                process = await asyncio.create_subprocess_exec(*args, env=self.env)
                 await asyncio.wait_for(process.wait(), timeout=30)
         else:
             # Still pending: stopping srun withdraws the job request
@@ -542,7 +732,12 @@ class SlurmClient:
         return self.started.is_set()
 
     def close(self):
-        pass
+        if self.alloc_id:
+            job_id, self.alloc_id = self.alloc_id, None
+            print(f'[{self.job_name}] releasing the allocation {job_id}')
+            with contextlib.suppress(Exception):
+                subprocess.run(['scancel', job_id], env=self.env, capture_output=True,
+                               timeout=30)
 
 
 class Worker:
@@ -640,13 +835,59 @@ class WorkerPoolReport:
             return getattr(task.process.client, "host", "?")
         return task.run_on_host or "?"
 
+    def usage_lines(self):
+        '''Per task type (label): tasks, longest wall time, total CPU time, mean cores used,
+        peak RSS, and the CPUs and memory of its category.'''
+        groups = collections.defaultdict(list)
+        for t in self.all_tasks:
+            if getattr(t, 'usage', None):
+                groups[t.label or '-'].append(t)
+        if not groups:
+            return []
+        lines = [f'{"task":24s} {"n":>4s} {"wall max":>9s} {"CPU total":>10s} {"cores":>6s} '
+                 f'{"peak RSS":>9s}   category (request)']
+        for label, tasks in sorted(groups.items()):
+            wall = [t.usage.get('wall', 0) for t in tasks]
+            cpu = sum(t.usage.get('cpu', 0) for t in tasks)
+            rss = max(t.usage.get('maxrss_kb', 0) for t in tasks) / 2 ** 20
+            t = tasks[0]
+            lines.append(f'{label:24s} {len(tasks):4d} {max(wall) / 60:7.1f} m {cpu / 3600:8.2f} h '
+                         f'{cpu / max(sum(wall), 1e-9):6.1f} {rss:5.1f} GiB   {t.category} '
+                         f'({t.cpus} CPUs, {t.mem_mb / 1024:g} GB)')
+        return lines
+
+    def memory_warnings(self):
+        '''Task types whose peak memory exceeds the memory of their category.'''
+        peaks = collections.defaultdict(float)
+        requests = {}
+        for t in self.all_tasks:
+            if getattr(t, 'usage', None) and getattr(t, 'mem_mb', None):
+                label = t.label or '-'
+                peaks[label] = max(peaks[label], t.usage.get('maxrss_kb', 0) / 1024)
+                requests[label] = (t.category, t.mem_mb)
+        return [f'{label} used up to {peak / 1024:.1f} GiB, above the {mem / 1024:g} GB of its '
+                f'category {category}: raise its memory or limit the task'
+                for label, peak in sorted(peaks.items())
+                for category, mem in [requests[label]] if peak > mem]
+
+    def print_usage(self):
+        usage = self.usage_lines()
+        if not usage:
+            return
+        print('\nUsage per task type (the memory of a category must stay above its peak RSS):')
+        print('\n'.join(usage))
+        for warning in self.memory_warnings():
+            secho(f'Warning: {warning}', fg='red')
+        print()
+
 
 class WorkerPool:
 
     def __init__(self, hosts, name='Worker', max_tasks_per_worker=4,
                  env_source_file=None, user=None, max_time=None, force_sync=False,
                  debug=False, dry_run=False, password=None, connect_timeout=15,
-                 host_cooldown_s=10, slurm_options=None, kill_grace_s=30):
+                 host_cooldown_s=10, slurm_options=None, kill_grace_s=30, slurm_nodes=None,
+                 resources=None):
         """Initiate a worker pool.
 
         Args:
@@ -664,6 +905,10 @@ 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.
+            resources (dict, optional): CPUs and memory of the task categories as job steps,
+                {category: {'cpus': .., 'mem': ..}}, updating DEFAULT_RESOURCES.
         """
         get_event_loop()
 
@@ -687,17 +932,26 @@ class WorkerPool:
         # Tasks waiting for a worker, in arrival order: (requested host, future)
         self._waiters = collections.deque()
         self._cooldown_timer = None
+        self.resources = resolve_resources(resources)
+        self.slurm_nodes = 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)
 
         for host in self.hosts:
+            n_workers = max_tasks_per_worker
             if host == localhost_shortname or host == 'localhost':
                 client = LocalClient(localhost_shortname, force_sync=force_sync)
             elif host == SLURM_HOST:
-                client = SlurmClient(slurm_options, job_name=name, max_time=max_time)
+                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)
             else:
                 client = Client(host, user, force_sync=force_sync, password=password,
                                 connect_timeout=connect_timeout)
 
-            for i in range(max_tasks_per_worker):
+            for i in range(n_workers):
                 worker = Worker(client, f'{client.host}:{i}')
                 self.workers.append(worker)
 
@@ -783,7 +1037,8 @@ class WorkerPool:
                 min(ends) - now + 0.01, self._dispatch)
 
     def add(self, command, name=None, output_file=None, task_done_callback=None, done_callback=None,
-            run_on_host=None, n_try=3, pid_filename=None, keep_pid_file_on_error=False):
+            run_on_host=None, n_try=3, pid_filename=None, keep_pid_file_on_error=False,
+            category=None, label=None):
         """Add a command to execute by a worker in the pool. Optionally output result
         in output_file.
 
@@ -791,7 +1046,14 @@ class WorkerPool:
             command (str): Command to execute.
             name (str, optional): Name of the task. If not set it will be set by the worker pool
             output_file (str, optional): Optional filename to log output into.
+            category (str, optional): Resource category of the task as a job step (see
+                DEFAULT_RESOURCES): small, medium, high_mem, cpu or heavy (the default).
+            label (str, optional): Task type, grouping the tasks in the usage summary
+                (default: the category).
         """
+        if category is not None and category not in self.resources:
+            raise ValueError(f'Unknown resource category {category!r}: '
+                             f'{", ".join(self.resources)}')
         if run_on_host is not None and run_on_host not in self.hosts:
             # Slurm places the tasks itself: a host constraint is expected not to apply
             if SLURM_HOST not in self.hosts:
@@ -802,6 +1064,8 @@ class WorkerPool:
         script = f". {shlex.quote(self.env_file)}; {command}" if self.env_file else command
 
         command = f"bash -lc {shlex.quote(script)}"
+        if self.accounting:
+            command = f'{shlex.quote(sys.executable)} {shlex.quote(taskrun.__file__)} {command}'
 
         if name is None:
             name = f'T{len(self.tasks)}'
@@ -811,10 +1075,12 @@ class WorkerPool:
         if self.dry_run:
             return
 
+        cpus, mem = self.resources[category or 'heavy']
         self.tasks.append(Task(name, command, output_file=output_file, done_callback=done_callback,
                                task_done_callback=task_done_callback, run_on_host=run_on_host,
                                n_try=n_try, pid_filename=pid_filename,
-                               keep_pid_file_on_error=keep_pid_file_on_error))
+                               keep_pid_file_on_error=keep_pid_file_on_error,
+                               category=category or 'heavy', label=label, cpus=cpus, mem_mb=mem))
 
     def get_all_tasks(self):
         tasks = collections.defaultdict(list)
@@ -885,6 +1151,11 @@ class WorkerPool:
 
         self.pbar.start(max_value=len(self.tasks))
 
+        # 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))
+
         self.futures = [asyncio.ensure_future(task.run(self)) for task in self.tasks]
         self.futures.append(asyncio.ensure_future(self._monitor_queue()))
 
@@ -926,6 +1197,8 @@ class WorkerPool:
 
         self.last_report = WorkerPoolReport(self.tasks)
         report = self.last_report
+        # At the end of the pool, before the next tasks fill the console
+        report.print_usage()
 
         if print_summary:
             print("\n=== WorkerPool summary ===")

{libpipe-0.8 → libpipe-0.9}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "libpipe"
-version = "0.8"
+version = "0.9"
 description = "Pipeline orchestration and node management utilities"
 readme = "README.md"
 license = "GPL-3.0-or-later"
@@ -26,11 +26,18 @@ dev = [
     "pytest-cov>=5.0",
     "ruff==0.16.1",
 ]
+docs = [
+    "sphinx>=7",
+    "myst-parser>=2",
+    "pydata-sphinx-theme>=0.15",
+    "sphinx-click>=6",
+]
 
 [project.urls]
 Homepage = "https://gitlab.com/flomertens/libpipe"
 Repository = "https://gitlab.com/flomertens/libpipe"
 Issues = "https://gitlab.com/flomertens/libpipe/-/issues"
+Documentation = "https://libpipe.readthedocs.io"
 
 [project.scripts]
 nodetool = "libpipe.tools.nodetool:main"

libpipe-0.8/README.md

@@ -1,30 +0,0 @@
-Libpipe
-=======
-
-Common pipeline framework and library
-
-Installation
-------------
-
-libpipe can be installed via pip:
-
-    $ pip install libpipe
-
-and requires Python 3.10 or higher.
-
-Worker pools on a Slurm cluster
--------------------------------
-
-A `WorkerPool` runs its tasks locally, on other hosts over SSH, or through Slurm with the
-host `slurm`: each task is then started with `srun`, as a job of its own when the pool runs
-on a login node, as a job step inside an allocation (`salloc`, `sbatch`).
-
-    pool = worker.WorkerPool(['slurm'], name='Imaging', max_tasks_per_worker=8,
-                             slurm_options='-p regular --mem=200G --time=24:00:00')
-
-`max_tasks_per_worker` is the number of jobs in flight. The srun options (partition,
-account, resources) come from `slurm_options`, or from the `LIBPIPE_SLURM_OPTIONS`
-environment variable; `max_time` becomes the job time limit when the options set none.
-Host constraints (`run_on_host`) do not apply: Slurm places the tasks, so the data must be
-on a file system shared with the compute nodes. Give the memory (`--mem`) in the options:
-inside an allocation, a step without it can claim all the memory of its node.