libpipe 0.7__tar.gz → 0.8__tar.gz

libpipe-0.8/PKG-INFO

@@ -0,0 +1,58 @@
+Metadata-Version: 2.3
+Name: libpipe
+Version: 0.8
+Summary: Pipeline orchestration and node management utilities
+License: GPL-3.0-or-later
+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
+Provides-Extra: dev
+Requires-Dist: asyncssh (>=2.0)
+Requires-Dist: click (>=8.0)
+Requires-Dist: progressbar2 (>=3.40)
+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: toml (>=0.10)
+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
+Description-Content-Type: text/markdown
+
+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.
+

libpipe-0.8/README.md

@@ -0,0 +1,30 @@
+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.

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

@@ -15,6 +15,7 @@ import signal
 import socket
 import sys
 import time
+import warnings
 from enum import Enum
 
 import asyncssh
@@ -23,6 +24,26 @@ from click import secho, style
 
 localhost_shortname = socket.gethostname().split('.', 1)[0]
 
+# Host name selecting the Slurm client: each task runs through srun
+SLURM_HOST = 'slurm'
+SLURM_OPTIONS_ENV = 'LIBPIPE_SLURM_OPTIONS'
+SLURM_MARKER = '[libpipe] slurm job '
+
+
+def get_event_loop():
+    """The current event loop, or a new one when there is none or it is closed (e.g. after
+    an asyncio.run() elsewhere in the process)."""
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', DeprecationWarning)
+            loop = asyncio.get_event_loop()
+    except RuntimeError:
+        loop = None
+    if loop is None or loop.is_closed():
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+    return loop
+
 
 def kill_remote(host, user, pid):
     os.system(f'ssh {user}@{host} "kill -TERM -{pid}"')
@@ -53,10 +74,11 @@ 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):
+                    debug=False, dry_run=False, slurm_options=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)
+                      env_source_file=env_file, max_time=max_time, debug=debug, dry_run=dry_run,
+                      slurm_options=slurm_options)
 
 
 class TaskStatus(Enum):
@@ -108,6 +130,10 @@ class Task:
         self.timed_out = False
         self.pid_filename = pid_filename
         self.keep_pid_file_on_error = keep_pid_file_on_error
+        # Set when the process starts (timeouts count from then), and when it was signalled
+        self.start_time = None
+        self.sigterm_time = None
+        self.slurm_job_id = None
 
     def done_callback(self):
         err_msg = ''
@@ -126,12 +152,14 @@ class Task:
 
     def init_log(self):
         if self.output_file is not None:
-            self.fd = open(self.output_file, 'w')  # noqa: SIM115 - closed in close_log()
+            # Line-buffered: the log follows the task, and survives a crash of the pool
+            self.fd = open(self.output_file, 'w', buffering=1)  # noqa: SIM115 - see close_log()
             self.fd.write(f'# Logging starting at {datetime.datetime.now()}\n')
             self.fd.write(f'# Input command: {self.command}\n')
 
     def set_process(self, process):
         self.process = process
+        self.start_time = time.time()
         if hasattr(self.process.process, 'pid'):
             self.log(f'Starting on host {self.process.client.host} with PID '
                      f'{self.process.process.pid}\n\n', self.process.client.host)
@@ -145,24 +173,14 @@ class Task:
                     self.log(f'Error creating PID file {self.pid_filename}',
                              self.process.client.host, err=True)
 
-    async def terminate(self):
-        if self.process is not None and hasattr(self.process.process, 'pid'):
-            if self.process.client.host == localhost_shortname:
-                self.log(f'Killing local process PID {self.process.process.pid}\n',
-                         'local', err=True)
-                try:
-                    os.killpg(os.getpgid(self.process.process.pid), signal.SIGTERM)
-                except ProcessLookupError:
-                    print(f'Process {self.process.process.pid} not found')
-            else:
-                self.log(f'Killing remote process PID {self.process.process.pid}\n',
-                         self.process.client.host, err=True)
-                kill_remote(self.process.client.host, self.process.client.user,
-                            self.process.process.pid)
+    async def terminate(self, sig=signal.SIGTERM):
+        if self.process is not None:
+            await self.process.client.terminate(self, sig)
 
     def log(self, line, host, err=False):
         if isinstance(line, bytes):
-            line = line.decode()
+            # Tools can print any byte: a decoding error must not interrupt the task
+            line = line.decode(errors='replace')
         if self.fd is not None:
             self.fd.write(line)
         line = line.strip()
@@ -194,33 +212,63 @@ class Task:
                 print(f'Error closing log file for task {self.name}. Ignoring.')
 
     async def run(self, client_pool):
-        self.init_log()
-
-        for i in range(self.n_try):
-            self.timed_out = False
-
-            worker = await client_pool.get_worker(self.run_on_host)
-            self.status = TaskStatus.RUNNING
-
-            ret_status, err_msg = await worker.run(self)
-
-            if not ret_status:
-                # mark host temporarily bad so next retry likely uses a different host
-                host = worker.client.host
-                client_pool.bad_hosts_until[host] = time.time() + client_pool.host_cooldown_s
+        try:
+            self.init_log()
+        except OSError as exc:
+            self.status = TaskStatus.ERROR
+            print(style(f'[{self.name}] Error opening log file {self.output_file}: {exc}',
+                        fg='red'))
+            return
 
-            if err_msg:
-                self.log(err_msg, worker.client.host, err=True)
+        try:
+            for i in range(self.n_try):
+                self.timed_out = False
+                self.process = None
+                self.start_time = None
+                self.sigterm_time = None
 
-            if ret_status and not self.timed_out:
-                self.status = TaskStatus.SUCCESS
-                break
+                worker = await client_pool.get_worker(self.run_on_host)
+                self.status = TaskStatus.RUNNING
 
+                try:
+                    ret_status, err_msg = await worker.run(self)
+                except BaseException:
+                    client_pool.release_worker()
+                    raise
+
+                if not ret_status:
+                    # mark host temporarily bad so next retry likely uses a different host
+                    host = worker.client.host
+                    client_pool.bad_hosts_until[host] = time.time() + client_pool.host_cooldown_s
+                client_pool.release_worker()
+
+                if err_msg:
+                    self.log(err_msg, worker.client.host, err=True)
+
+                if ret_status and not self.timed_out:
+                    self.status = TaskStatus.SUCCESS
+                    break
+
+                self.status = TaskStatus.ERROR
+                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
+                    self.log('Not retrying: the process had started\n', worker.client.host,
+                             err=True)
+                    break
+                if i + 1 < self.n_try:
+                    self.log(f'Will retry {i + 1} / {self.n_try} ...\n', worker.client.host,
+                             err=True)
+        except asyncio.CancelledError:
+            self.status = TaskStatus.CANCELED
+            raise
+        except Exception as exc:
+            # An error of this task must not bring down the pool and its other tasks
             self.status = TaskStatus.ERROR
-            if i + 1 < self.n_try:
-                self.log(f'Will retry {i + 1} / {self.n_try} ...\n', worker.client.host, err=True)
-
-        self.close_log()
+            with contextlib.suppress(Exception):
+                self.log(f'Unexpected error running task: {exc!r}\n', 'pool', err=True)
+        finally:
+            self.close_log()
 
 
 class SSHClient(asyncssh.SSHClient):
@@ -299,7 +347,8 @@ class Client:
             raise ConnectionError('Client not connected')
 
         async with self.creating_session:
-            process = await self.conn.create_process('echo $$;' + task.command)
+            process = await self.conn.create_process('echo $$;' + task.command,
+                                                     encoding='utf-8', errors='replace')
 
         try:
             process.pid = int(await process.stdout.readline())
@@ -320,6 +369,31 @@ class Client:
             await self.conn.run('sync')
             await asyncio.sleep(1)
 
+    async def terminate(self, task, sig=signal.SIGTERM):
+        pid = getattr(task.process.process, 'pid', None)
+        if pid is None:
+            task.log('Cannot kill the remote process: PID unknown\n', self.host, err=True)
+            return
+        task.log(f'Killing remote process PID {pid} ({sig.name})\n', self.host, err=True)
+        cmd = f'kill -{sig.name.removeprefix("SIG")} -{pid}'
+        # On the open connection, or on a new one when it is lost: never blocking the pool
+        # on a password prompt or an unreachable host
+        try:
+            if self.conn is not None:
+                await asyncio.wait_for(self.conn.run(cmd, check=False), timeout=15)
+                return
+        except Exception:  # fall back to a new connection
+            pass
+        try:
+            process = await asyncio.create_subprocess_exec(
+                'ssh', '-o', 'BatchMode=yes', '-o', 'ConnectTimeout=10',
+                f'{self.user}@{self.host}', cmd,
+                stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.DEVNULL)
+            await asyncio.wait_for(process.wait(), timeout=20)
+        except Exception as exc:
+            task.log(f'Could not kill remote process PID {pid}: {exc!r}\n', self.host,
+                     err=True)
+
     def close(self):
         if self.conn is not None and not self.closing:
             self.closing = True
@@ -356,6 +430,13 @@ class LocalClient:
             os.sync()
             await asyncio.sleep(1)
 
+    async def terminate(self, task, sig=signal.SIGTERM):
+        pid = task.process.process.pid
+        task.log(f'Killing local process PID {pid} ({sig.name})\n', self.host, err=True)
+        # The process leads its own group (setpgrp): this reaches its children as well
+        with contextlib.suppress(ProcessLookupError, PermissionError):
+            os.killpg(pid, sig)
+
     def connected(self):
         return self.started.is_set()
 
@@ -364,6 +445,106 @@ class LocalClient:
         pass
 
 
+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.
+    """
+
+    def __init__(self, options=None, job_name='libpipe', max_time=None):
+        self.host = SLURM_HOST
+        self.started = asyncio.Event()
+        if options is None:
+            options = os.environ.get(SLURM_OPTIONS_ENV, '')
+        self.options = shlex.split(options) if isinstance(options, str) else list(options)
+        self.job_name = re.sub(r'\s+', '_', job_name)
+        self.max_time = max_time
+        self.exact = False
+
+    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)
+        self.started.set()
+
+    @staticmethod
+    async def _slurm_version():
+        try:
+            process = await asyncio.create_subprocess_exec(
+                'srun', '--version', stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.DEVNULL)
+            out, _ = await process.communicate()
+        except FileNotFoundError:
+            return (0, 0)
+        m = re.search(r'(\d+)\.(\d+)', out.decode())
+        return (int(m.group(1)), int(m.group(2))) if m else (0, 0)
+
+    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}")}'
+
+    async def _log_stdout(self, task, stream):
+        """Log the task output, recording the job id from the marker line."""
+        while True:
+            try:
+                line = await stream.readline()
+            except ValueError:
+                task.log('(line too long)', self.host, err=True)
+                continue
+            if not line:
+                break
+            text = line.decode(errors='replace')
+            if task.slurm_job_id is None and text.startswith(SLURM_MARKER):
+                task.slurm_job_id = text[len(SLURM_MARKER):].split()[0]
+            task.log(line, self.host)
+
+    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()
+
+    async def terminate(self, task, sig=signal.SIGTERM):
+        job_id = getattr(task, 'slurm_job_id', None)
+        if job_id:
+            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)
+                await asyncio.wait_for(process.wait(), timeout=30)
+        else:
+            # Still pending: stopping srun withdraws the job request
+            task.log(f'Stopping srun PID {task.process.process.pid}\n', self.host, err=True)
+            with contextlib.suppress(ProcessLookupError, PermissionError):
+                os.killpg(task.process.process.pid, sig)
+
+    def connected(self):
+        return self.started.is_set()
+
+    def close(self):
+        pass
+
+
 class Worker:
 
     def __init__(self, client, name):
@@ -381,6 +562,10 @@ class Worker:
 
         try:
             await self.client.start()
+            await self.client.started.wait()
+        except asyncio.CancelledError:
+            self.status = WorkerStatus.IDLE
+            raise
         except Exception as exc:
             err_msg = f'Error starting client: {exc!s}\n'
             # transient failure: allow task retry to pick a worker again
@@ -389,32 +574,32 @@ class Worker:
                 self.client.close()  # ensure clean slate
             return False, err_msg
 
-        await self.client.started.wait()
-
         try:
             await self.client.execute(task)
-        except ConnectionError as exc:
-            err_msg = f'Error connecting client: {exc!s}\n'
-            self.status = WorkerStatus.IDLE
-            with contextlib.suppress(Exception):
-                self.client.close()
-            return False, err_msg
         except asyncio.CancelledError:
-            err_msg = f'Task {self.name} has been canceled\n'
-            return True, err_msg
+            self.status = WorkerStatus.IDLE
+            raise
         except Exception as exc:
-            err_msg = f'Error executing task {self.name}: {exc!s}\n'
+            if isinstance(exc, ConnectionError):
+                err_msg = f'Error connecting client: {exc!s}\n'
+            else:
+                err_msg = f'Error executing task {self.name}: {exc!r}\n'
+            if task.process is not None:
+                # The process may still run: stop it, the task will not be run again
+                with contextlib.suppress(Exception):
+                    await task.terminate()
             self.status = WorkerStatus.IDLE
+            if isinstance(exc, ConnectionError):
+                with contextlib.suppress(Exception):
+                    self.client.close()
             return False, err_msg
 
-        err_msg = task.done_callback()
-        # err_msg = ''
-        # try:
-        #     loop = asyncio.get_running_loop()
-        #     err_msg = await loop.run_in_executor(None, task.done_callback)
-        # except Exception as exc:
-        #     err_msg = f'Error executing callback of task {self.name}: {exc}\n'
+        if task.returncode is None:
+            # e.g. the SSH connection dropped before the exit status arrived
+            self.status = WorkerStatus.IDLE
+            return False, 'No exit status received: the outcome of the task is unknown\n'
 
+        err_msg = task.done_callback()
         self.status = WorkerStatus.IDLE
 
         return True, err_msg
@@ -446,7 +631,9 @@ class WorkerPoolReport:
         self.n_success_nonzero = len(self.success_nonzero)
         self.n_success_unknown_rc = len(self.success_unknown_rc)
 
-        self.ok = (self.n_failed == 0 and self.n_success_nonzero == 0)
+        # A task without exit status did not demonstrably succeed
+        self.ok = (self.n_failed == 0 and self.n_success_nonzero == 0
+                   and self.n_success_unknown_rc == 0)
 
     def host_of(self, task):
         if task.process and task.process.client:
@@ -459,22 +646,26 @@ 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):
+                 host_cooldown_s=10, slurm_options=None, kill_grace_s=30):
         """Initiate a worker pool.
 
         Args:
             name (str): Name of the worker pool.
-            hosts (list): List of host names, reachable via ssh (if not local host).
+            hosts (list): List of host names, reachable via ssh (if not local host). The host
+                'slurm' runs the tasks through srun (see SlurmClient).
             max_tasks_per_worker (int, optional): Maximum number of tasks to execute
-                concurrently on an host.
+                concurrently on an host (for 'slurm', the number of jobs in flight).
             env_source_file (str, optional): Name of a file to source before executing a task.
             user (str, optional): User to connect
             connect_timeout (int, optional): Per-attempt SSH connection timeout, in seconds.
             host_cooldown_s (int, optional): Time a just-failed host is avoided while
                 a healthy one exists.
+            slurm_options (str, optional): srun options for the 'slurm' host, e.g.
+                '-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.
         """
-        if asyncio.get_event_loop().is_closed():
-            asyncio.set_event_loop(asyncio.new_event_loop())
+        get_event_loop()
 
         if user is None:
             user = getpass.getuser()
@@ -492,10 +683,16 @@ class WorkerPool:
         # host -> epoch time until which a just-failed host is avoided (steering only)
         self.host_cooldown_s = host_cooldown_s
         self.bad_hosts_until = {}
+        self.kill_grace_s = kill_grace_s
+        # Tasks waiting for a worker, in arrival order: (requested host, future)
+        self._waiters = collections.deque()
+        self._cooldown_timer = None
 
         for host in self.hosts:
             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)
             else:
                 client = Client(host, user, force_sync=force_sync, password=password,
                                 connect_timeout=connect_timeout)
@@ -514,29 +711,76 @@ class WorkerPool:
 
         self.pbar = progressbar.ProgressBar(redirect_stdout=True, widgets=widgets)
 
-    async def get_worker(self, host=None):
-        while True:
-            now = time.time()
+    def _find_idle_worker(self, host):
+        now = time.time()
 
-            # hosts not currently in cooldown
-            good_hosts = {h for h in self.hosts if self.bad_hosts_until.get(h, 0) <= now}
+        # hosts not currently in cooldown
+        good_hosts = {h for h in self.hosts if self.bad_hosts_until.get(h, 0) <= now}
 
-            for worker in self.workers:
-                w_host = worker.client.host
+        for worker in self.workers:
+            w_host = worker.client.host
 
-                # respect run_on_host if requested
-                if host is not None and w_host != host:
-                    continue
+            # respect run_on_host if requested
+            if host is not None and w_host != host:
+                continue
 
-                # if we have at least one good host, avoid bad ones
-                if host is None and good_hosts and w_host not in good_hosts:
-                    continue
+            # if we have at least one good host, avoid bad ones
+            if host is None and good_hosts and w_host not in good_hosts:
+                continue
 
-                if worker.status is WorkerStatus.IDLE:
-                    worker.status = WorkerStatus.COMMITED
-                    return worker
+            if worker.status is WorkerStatus.IDLE:
+                return worker
+        return None
 
-            await asyncio.sleep(0.1)
+    async def get_worker(self, host=None):
+        fut = asyncio.get_running_loop().create_future()
+        self._waiters.append((host, fut))
+        self._dispatch()
+        try:
+            return await fut
+        except asyncio.CancelledError:
+            if fut.done() and not fut.cancelled():
+                # A worker was handed over just as the task was cancelled: give it back
+                fut.result().status = WorkerStatus.IDLE
+                self._dispatch()
+            with contextlib.suppress(ValueError):
+                self._waiters.remove((host, fut))
+            raise
+
+    def release_worker(self):
+        """Hand the worker just freed to the next task waiting for it."""
+        self._dispatch()
+
+    def _dispatch(self):
+        """Hand the idle workers to the waiting tasks, in arrival order."""
+        i = 0
+        while i < len(self._waiters) and any(w.status is WorkerStatus.IDLE
+                                             for w in self.workers):
+            host, fut = self._waiters[i]
+            if fut.done():
+                del self._waiters[i]
+                continue
+            worker = self._find_idle_worker(host)
+            if worker is None:
+                i += 1
+                continue
+            worker.status = WorkerStatus.COMMITED
+            fut.set_result(worker)
+            del self._waiters[i]
+        self._schedule_cooldown_dispatch()
+
+    def _schedule_cooldown_dispatch(self):
+        """The end of a host cooldown can make an idle worker usable: dispatch again then."""
+        if self._cooldown_timer is not None:
+            self._cooldown_timer.cancel()
+            self._cooldown_timer = None
+        if not self._waiters:
+            return
+        now = time.time()
+        ends = [t for t in self.bad_hosts_until.values() if t > now]
+        if ends:
+            self._cooldown_timer = asyncio.get_running_loop().call_later(
+                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):
@@ -549,8 +793,10 @@ class WorkerPool:
             output_file (str, optional): Optional filename to log output into.
         """
         if run_on_host is not None and run_on_host not in self.hosts:
-            print(f'Requested run_on_host "{run_on_host}" is not in the hosts list. '
-                  'Releasing constraint ...')
+            # Slurm places the tasks itself: a host constraint is expected not to apply
+            if SLURM_HOST not in self.hosts:
+                print(f'Requested run_on_host "{run_on_host}" is not in the hosts list. '
+                      'Releasing constraint ...')
             run_on_host = None
 
         script = f". {shlex.quote(self.env_file)}; {command}" if self.env_file else command
@@ -587,21 +833,37 @@ class WorkerPool:
                 self.pbar.update(n_done)
 
                 if self.max_time is not None:
-                    for worker in self.workers:
-                        if worker.status == WorkerStatus.RUNNING:
-                            running_time = time.time() - worker.execute_start_time
-                            if running_time > (self.max_time + update_interval):
-                                task = worker.running_task
-                                task.timed_out = True
-                                task.log(f'Task timeout: waited for {running_time:.2f} s',
-                                         worker.client.host, err=True)
-                                task.log(f'Command was: {task.command}',
-                                         worker.client.host, err=True)
-                                await task.terminate()
+                    await self._enforce_max_time()
 
         except asyncio.CancelledError:
             pass
 
+    async def _enforce_max_time(self):
+        now = time.time()
+        for worker in self.workers:
+            # Slurm enforces max_time as the job time limit, from the job start
+            if isinstance(worker.client, SlurmClient) or worker.status != WorkerStatus.RUNNING:
+                continue
+            task = worker.running_task
+            # Count from the process start: connecting to the host is not running the task
+            if task is None or task.start_time is None or task.status != TaskStatus.RUNNING:
+                continue
+            running_time = now - task.start_time
+            if running_time <= self.max_time:
+                continue
+            if task.sigterm_time is None:
+                task.timed_out = True
+                task.sigterm_time = now
+                task.log(f'Task timeout: running for {running_time:.2f} s', worker.client.host,
+                         err=True)
+                task.log(f'Command was: {task.command}', worker.client.host, err=True)
+                await task.terminate(signal.SIGTERM)
+            elif now - task.sigterm_time > self.kill_grace_s and task.sigterm_time > 0:
+                task.log(f'Still running {self.kill_grace_s} s after SIGTERM: sending SIGKILL',
+                         worker.client.host, err=True)
+                task.sigterm_time = -1  # SIGKILL sent, nothing more to do
+                await task.terminate(signal.SIGKILL)
+
     async def terminate_all(self):
         for worker in self.workers:
             worker.close()
@@ -615,7 +877,7 @@ class WorkerPool:
     async def shutdown(self, exit_code=0):
         print('Shutting down')
         await self.terminate_all()
-        asyncio.get_event_loop().call_soon_threadsafe(sys.exit, exit_code)
+        asyncio.get_running_loop().call_soon_threadsafe(sys.exit, exit_code)
 
     async def process_tasks(self):
         # Sort tasks so that those with run_on_host are started first
@@ -636,7 +898,7 @@ class WorkerPool:
             self.pbar.finish()
             return [], []
 
-        loop = asyncio.get_event_loop()
+        loop = get_event_loop()
 
         loop.add_signal_handler(signal.SIGTERM,
                                 lambda: asyncio.ensure_future(self.shutdown(exit_code=1)))
@@ -648,6 +910,10 @@ class WorkerPool:
             raise
         except asyncio.CancelledError:
             secho('Error: one (or more) task(s) have been canceled', fg='red')
+        except Exception:
+            # Do not leave the running tasks behind
+            loop.run_until_complete(self.terminate_all())
+            raise
         finally:
             for worker in self.workers:
                 worker.close()

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

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

libpipe-0.7/PKG-INFO

@@ -1,42 +0,0 @@
-Metadata-Version: 2.4
-Name: libpipe
-Version: 0.7
-Summary: Pipeline orchestration and node management utilities
-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: 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
-Requires-Dist: asyncssh (>=2.0)
-Requires-Dist: click (>=8.0)
-Requires-Dist: progressbar2 (>=3.40)
-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: toml (>=0.10)
-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
-Description-Content-Type: text/markdown
-
-Libpipe
-=======
-
-Common pipeline framework and library
-
-Installation
-------------
-
-libpipe can be installed via pip:
-
-    $ pip install libpipe
-
-and requires Python 3.6.0 or higher.
-

libpipe-0.7/README.md

@@ -1,13 +0,0 @@
-Libpipe
-=======
-
-Common pipeline framework and library
-
-Installation
-------------
-
-libpipe can be installed via pip:
-
-    $ pip install libpipe
-
-and requires Python 3.6.0 or higher.