pivot 0.3.7__tar.gz → 0.3.8__tar.gz

{pivot-0.3.7 → pivot-0.3.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pivot
-Version: 0.3.7
+Version: 0.3.8
 Summary: High-performance Python pipeline tool with automatic code change detection
 Author: Sami Jawhar
 Author-email: Sami Jawhar <sami@metr.org>
@@ -23,6 +23,7 @@ Requires-Dist: loky>=3.4
 Requires-Dist: networkx>=3.0
 Requires-Dist: pathspec>=0.12.0
 Requires-Dist: pandas>=2.0
+Requires-Dist: psutil>=7.0
 Requires-Dist: pydantic>=2.13
 Requires-Dist: pygtrie>=2.5.0
 Requires-Dist: pyyaml>=6.0
@@ -55,6 +56,7 @@ Requires-Dist: pytest-xdist>=3.8.0 ; extra == 'dev'
 Requires-Dist: pytest>=7.0 ; extra == 'dev'
 Requires-Dist: ruff>=0.8.0 ; extra == 'dev'
 Requires-Dist: types-aioboto3[all]>=13.0.0 ; extra == 'dev'
+Requires-Dist: types-psutil>=7.0 ; extra == 'dev'
 Requires-Dist: mkdocs>=1.6.0 ; extra == 'docs'
 Requires-Dist: mkdocs-material>=9.5.0 ; extra == 'docs'
 Requires-Dist: mkdocstrings[python]>=0.26.0 ; extra == 'docs'

{pivot-0.3.7 → pivot-0.3.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pivot"
-version = "0.3.7"
+version = "0.3.8"
 description = "High-performance Python pipeline tool with automatic code change detection"
 requires-python = ">=3.13,<3.14"
 readme = "README.md"
@@ -27,6 +27,7 @@ dependencies = [
   "networkx>=3.0",
   "pathspec>=0.12.0",
   "pandas>=2.0",
+  "psutil>=7.0",
   "pydantic>=2.13",
   "pygtrie>=2.5.0",
   "pyyaml>=6.0",
@@ -76,6 +77,7 @@ dev = [
   "pytest>=7.0",
   "ruff>=0.8.0",
   "types-aioboto3[all]>=13.0.0",
+  "types-psutil>=7.0",
 ]
 
 [project.scripts]

{pivot-0.3.7 → pivot-0.3.8}/src/pivot/engine/worker_pool.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
+import contextlib
 import errno
 import logging
 import os
 import pathlib
+import signal
 import threading
+import time
 import typing
 from dataclasses import dataclass, field
 
@@ -17,6 +20,9 @@ _T = typing.TypeVar("_T")
 
 _logger = logging.getLogger(__name__)
 
+# How long hard_cancel waits for a SIGKILLed worker to be reaped; it never blocks pivot's exit.
+_KILLED_WORKER_JOIN_S = 5.0
+
 _FileStamp = tuple[int, int, int, int, int]
 
 
@@ -194,16 +200,57 @@ class WorkerPool:
         self._shutdown_event = None
 
     def hard_cancel(self) -> None:
-        """Forcefully shut down workers and signal drain thread to stop."""
+        """Stop workers mid-stage and signal the drain thread to stop.
+
+        Each worker is sent SIGTERM, which unwinds its running stage
+        (``worker.unwind_when_stopped``) so the stage's cleanup runs, and is given
+        ``worker.TERMINATION_GRACE_S`` to exit; a worker still alive then is killed with its
+        descendants.
+
+        The workers are first taken out of loky's executor. Its manager thread takes a worker
+        that exits without being asked to for a broken pool and SIGKILLs every worker the
+        executor still lists, so the first worker to exit would cut the others' cleanup off.
+        """
         self._accepting = False
         self._in_run = False
-        if self._executor is not None:
-            self._executor.shutdown(wait=False, kill_workers=True)
-            self._executor = None
-        if self._manager is not None:
-            self._manager.shutdown()
-            self._manager = None
-        if self._shutdown_event is not None:
-            self._shutdown_event.set()
-        self._output_queue = None
-        self._shutdown_event = None
+        try:
+            if self._executor is not None:
+                executor = self._executor
+                self._executor = None
+                try:
+                    self._stop_workers(executor)
+                finally:
+                    # Only shutdown wakes loky's manager thread once its workers are gone.
+                    executor.shutdown(wait=False, kill_workers=True)
+        finally:
+            if self._manager is not None:
+                self._manager.shutdown()
+                self._manager = None
+            if self._shutdown_event is not None:
+                self._shutdown_event.set()
+            self._output_queue = None
+            self._shutdown_event = None
+
+    @staticmethod
+    def _stop_workers(executor: loky.ProcessPoolExecutor) -> None:
+        """Take ``executor``'s workers out of it, send each SIGTERM, and SIGKILL with its
+        descendants any worker still alive after ``worker.TERMINATION_GRACE_S``."""
+        # loky exposes no public list of its worker processes. A watch reload may have
+        # replaced this executor already, which leaves it no lock and no workers.
+        with executor._processes_management_lock or contextlib.nullcontext():
+            processes = list(executor._processes.values())
+            executor._processes.clear()
+        for process in processes:
+            if process.pid is not None:
+                with contextlib.suppress(ProcessLookupError):
+                    os.kill(process.pid, signal.SIGTERM)
+        deadline = time.monotonic() + worker.TERMINATION_GRACE_S
+        for process in processes:
+            process.join(max(0.0, deadline - time.monotonic()))
+        for process in processes:
+            if process.pid is not None and process.exitcode is None:
+                worker.kill_descendants(process.pid)
+                # loky's process object has no kill(); signal the pid directly.
+                with contextlib.suppress(ProcessLookupError):
+                    os.kill(process.pid, signal.SIGKILL)
+                process.join(_KILLED_WORKER_JOIN_S)

{pivot-0.3.7 → pivot-0.3.8}/src/pivot/executor/AGENTS.md

@@ -3,8 +3,10 @@
 Workers execute in separate processes via `loky.get_reusable_executor()`. An Engine keeps its
 workers and manager alive across its runs; each run receives a fresh manager queue for output,
 and a manager that cannot be reached when a run starts is replaced then.
-The loky initializer and the manager's initializer start the parent-death watcher, so a
-worker or manager exits when its Pivot process dies. Keep `timeout=None` and
+The loky initializer (`worker.unwind_when_stopped`) and the manager's initializer
+(`core._exit_when_parent_dies`) start the parent-death watcher, so a worker or manager exits
+when its Pivot process dies; a worker first unwinds its running stage so its `finally` blocks
+run, within `worker.TERMINATION_GRACE_S`. Keep `timeout=None` and
 `reuse=False if kill_workers else "auto"` in the shared executor policy.
 
 ## WorkerStageInfo Contract

{pivot-0.3.7 → pivot-0.3.8}/src/pivot/executor/core.py

@@ -37,7 +37,10 @@ def _noop() -> None:
 
 
 def _exit_when_parent_dies(parent_pid: int) -> None:
-    """Exit a long-lived child (worker or manager) after its owning Pivot process dies."""
+    """Exit the output-queue manager after its owning Pivot process dies.
+
+    Workers use ``worker.unwind_when_stopped`` instead, which lets the running stage unwind.
+    """
 
     def _watch_parent() -> None:
         poll = threading.Event()
@@ -82,7 +85,7 @@ def _get_executor(max_workers: int, *, kill_workers: bool = False) -> loky.Proce
         timeout=_WORKER_IDLE_TIMEOUT_S,
         kill_workers=kill_workers,
         reuse=False if kill_workers else "auto",
-        initializer=_exit_when_parent_dies,
+        initializer=worker.unwind_when_stopped,
         initargs=(os.getpid(),),
     )
 

{pivot-0.3.7 → pivot-0.3.8}/src/pivot/executor/worker.py

@@ -10,17 +10,31 @@ import codecs
 import collections
 import collections.abc as collections_abc
 import contextlib
+import dataclasses
 import logging
 import os
 import pathlib
 import queue
 import random
+import signal
 import sys
 import threading
 import traceback
 import warnings
-from typing import IO, TYPE_CHECKING, Any, Literal, NamedTuple, TypedDict, cast, overload, override
+from typing import (
+    IO,
+    TYPE_CHECKING,
+    Any,
+    Literal,
+    NamedTuple,
+    NoReturn,
+    TypedDict,
+    cast,
+    overload,
+    override,
+)
 
+import psutil
 import pydantic
 
 from pivot import (
@@ -58,8 +72,142 @@ if TYPE_CHECKING:
     from multiprocessing import Queue
     from types import TracebackType
 
+
 logger = logging.getLogger(__name__)
 
+# How long a stopped worker's running stage may take to unwind before the worker exits.
+TERMINATION_GRACE_S = 10.0
+
+
+class WorkerTerminated(KeyboardInterrupt):  # noqa: N818 - names the event, as KeyboardInterrupt does
+    """Raised in a worker's running stage when its run is cancelled or its Pivot process dies.
+
+    It is a KeyboardInterrupt, as Ctrl-C raises at a terminal: a stage's ``except Exception``
+    does not stop it, and its ``except KeyboardInterrupt`` cleanup, ``finally`` blocks and
+    context managers run. A stage that catches it and carries on has its result discarded.
+    """
+
+
+@dataclasses.dataclass(slots=True)
+class _StopState:
+    """Process-wide stop state of a worker, set by its signal handlers and its watchdog.
+
+    Signals and the parent watchdog act on the whole process, so this state is too.
+    """
+
+    stage_running: bool = False
+    # The running stage's run was cancelled: its result must not be committed.
+    stop_requested: bool = False
+    # WorkerTerminated was raised into the running stage; never raise it again mid-cleanup.
+    interrupted: bool = False
+    # The worker exits once the running stage has unwound (SIGTERM or parent death).
+    exit_requested: bool = False
+    exit_timer_started: bool = False
+    lock: threading.Lock = dataclasses.field(default_factory=threading.Lock)
+
+
+_stop = _StopState()
+
+
+def kill_descendants(pid: int) -> None:
+    """SIGKILL every live descendant of process ``pid``. A process that detached from that
+    tree (double-forked, or orphaned when its parent exited) is not one."""
+    try:
+        descendants = psutil.Process(pid).children(recursive=True)
+    except psutil.Error:
+        return
+    for descendant in descendants:
+        with contextlib.suppress(psutil.Error):
+            descendant.kill()
+
+
+def _exit_worker(code: int) -> NoReturn:
+    """Exit this worker, first killing the processes still descended from it, which would
+    otherwise be left to init (a stage's ``Popen`` with no cleanup, or one whose cleanup was
+    cut off)."""
+    kill_descendants(os.getpid())
+    os._exit(code)
+
+
+def _exit_after_grace() -> None:
+    threading.Event().wait(TERMINATION_GRACE_S)
+    _exit_worker(1)
+
+
+def _start_exit_timer() -> None:
+    """Exit this worker ``TERMINATION_GRACE_S`` from now, whatever its running stage does."""
+    with _stop.lock:
+        if _stop.exit_timer_started:
+            return
+        _stop.exit_timer_started = True
+    threading.Thread(target=_exit_after_grace, daemon=True, name="pivot-stop-grace").start()
+
+
+def _interrupt_stage(_signum: int, _frame: object) -> None:
+    """SIGINT handler of a worker: Ctrl-C at a terminal, which reaches Pivot and its workers at
+    once. Interrupt the running stage once, so its cleanup runs, and keep the worker up until
+    Pivot stops it: loky takes a worker that exits on its own for a broken pool and SIGKILLs
+    every other worker, cutting their cleanup off."""
+    if not _stop.stage_running or _stop.interrupted:
+        return
+    _stop.stop_requested = True
+    _stop.interrupted = True
+    raise WorkerTerminated
+
+
+def _stop_worker(_signum: int, _frame: object) -> None:
+    """SIGTERM handler of a worker: interrupt the running stage once, exit once it has unwound,
+    and exit regardless after ``TERMINATION_GRACE_S``. An idle worker exits at once."""
+    _stop.stop_requested = True
+    _stop.exit_requested = True
+    _start_exit_timer()
+    if not _stop.stage_running:
+        _exit_worker(0)
+    if _stop.interrupted:
+        return
+    _stop.interrupted = True
+    raise WorkerTerminated
+
+
+def _raise_if_stopped() -> None:
+    """Refuse to commit the result of a stage whose run was cancelled, which can only be here
+    because the stage caught ``WorkerTerminated`` and carried on."""
+    if _stop.stop_requested:
+        raise WorkerTerminated
+
+
+def unwind_when_stopped(parent_pid: int) -> None:
+    """Loky worker initializer: stop this worker when it is told to or its parent dies.
+
+    SIGTERM (``WorkerPool.hard_cancel``) or SIGINT (Ctrl-C at a terminal) raises
+    ``WorkerTerminated`` in the running stage, so its cleanup runs (stopping the processes or
+    containers it started); after SIGTERM the worker then exits. When the Pivot process that
+    started the worker dies, as when it is SIGKILLed, a watchdog thread sends this worker's
+    main thread SIGTERM (a signal to that thread interrupts a blocking call there) and exits
+    the worker after ``TERMINATION_GRACE_S`` itself, in case the stage replaced the SIGTERM
+    handler or is blocked where no signal reaches it. Every exit first kills the processes
+    still descended from the worker.
+
+    A SIGINT the worker inherited as ignored (pivot run as a background job of a
+    non-interactive shell) stays ignored, as it does for pivot itself.
+    """
+    signal.signal(signal.SIGTERM, _stop_worker)
+    if signal.getsignal(signal.SIGINT) is not signal.SIG_IGN:
+        signal.signal(signal.SIGINT, _interrupt_stage)
+    main_thread = threading.main_thread().ident
+    assert main_thread is not None
+
+    def _watch_parent() -> None:
+        poll = threading.Event()
+        while os.getppid() == parent_pid:
+            poll.wait(1)
+        _stop.exit_requested = True  # first: a stage starting now keeps stop_requested set
+        _stop.stop_requested = True
+        _start_exit_timer()
+        signal.pthread_kill(main_thread, signal.SIGTERM)
+
+    threading.Thread(target=_watch_parent, daemon=True, name="pivot-parent-watch").start()
+
 
 class _QueueLoggingHandler(logging.Handler):
     """Logging handler that sends records to the output queue.
@@ -375,6 +523,7 @@ def _execute_stage(
                                 ),
                                 output_paths=lock_output_paths(stage_outs, store, project_root),
                             )
+                            _raise_if_stopped()
                             deferred = _commit_lock_and_build_deferred(
                                 stage_info,
                                 new_lock_data,
@@ -447,6 +596,7 @@ def _execute_stage(
 
                 # Single StateDB open for post-execution work
                 with state.StateDB(state_dir, readonly=True) as state_db:
+                    _raise_if_stopped()
                     deferred = _commit_lock_and_build_deferred(
                         stage_info,
                         new_lock_data,
@@ -509,14 +659,39 @@ def user_code_files() -> list[str]:
     return sorted(files)
 
 
+def _begin_stage() -> None:
+    """Mark a stage running, clearing a Ctrl-C that interrupted the previous one. The running
+    flag is set last, so a stop signal before it is handled as one to an idle worker."""
+    _stop.interrupted = False
+    if not _stop.exit_requested:
+        _stop.stop_requested = False
+    _stop.stage_running = True
+
+
 def execute_stage(
     stage_name: str,
     stage_info: WorkerStageInfo,
     cache_dir: pathlib.Path,
     output_queue: Queue[OutputMessage],
 ) -> StageResult:
-    """Execute a stage and report every user-code module resident in this worker."""
-    result = _execute_stage(stage_name, stage_info, cache_dir, output_queue)
+    """Execute a stage and report every user-code module resident in this worker.
+
+    A stage interrupted by a stop signal (``unwind_when_stopped``) fails without committing
+    anything. After SIGTERM or the death of its Pivot process, the worker exits once the stage
+    has unwound instead of returning.
+    """
+    try:
+        _begin_stage()
+        result = _execute_stage(stage_name, stage_info, cache_dir, output_queue)
+    except WorkerTerminated:
+        # Raised outside the stage's own error handling, as while its logging is set up.
+        result = StageResult(
+            status=StageStatus.FAILED, reason="KeyboardInterrupt", input_hash=None, output_lines=[]
+        )
+    finally:
+        _stop.stage_running = False
+    if _stop.exit_requested:
+        _exit_worker(1)
     result["user_code_files"] = user_code_files()
     return result
 
@@ -1116,6 +1291,7 @@ def _run_stage_function_with_store(
         _set_deterministic_seeds()
 
         result = _execute_with_joblib_protection(func, kwargs)
+        _raise_if_stopped()
 
         output_values = _resolve_output_values(result, outs, stage_name)
         output_paths = dict[types.ArtifactRef, pathlib.Path]()