slide2vec 4.5.1__tar.gz → 4.5.2__tar.gz

{slide2vec-4.5.1 → slide2vec-4.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: slide2vec
-Version: 4.5.1
+Version: 4.5.2
 Summary: Embedding of whole slide images with Foundation Models
 Author-email: Clément Grisi <clement.grisi@radboudumc.nl>
 License-Expression: Apache-2.0

{slide2vec-4.5.1 → slide2vec-4.5.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "slide2vec"
-version = "4.5.1"
+version = "4.5.2"
 description = "Embedding of whole slide images with Foundation Models"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -164,7 +164,7 @@ no_implicit_reexport = true
 max-line-length = 160
 
 [tool.bumpver]
-current_version = "4.5.1"
+current_version = "4.5.2"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit = false       # We do version bumping in CI, not as a commit
 tag = false          # Git tag already exists — we don't auto-tag

{slide2vec-4.5.1 → slide2vec-4.5.2}/slide2vec/__init__.py

@@ -11,7 +11,7 @@ from slide2vec.api import (
 from slide2vec.artifacts import HierarchicalEmbeddingArtifact, SlideEmbeddingArtifact, TileEmbeddingArtifact
 
 
-__version__ = "4.5.1"
+__version__ = "4.5.2"
 
 __all__ = [
     "Model",

{slide2vec-4.5.1 → slide2vec-4.5.2}/slide2vec/runtime/distributed.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import heapq
+import os
 import shutil
+import signal
 import subprocess
 import sys
 import tempfile
@@ -55,6 +57,42 @@ def write_worker_logs(module: str, output_dir: Path, stdout_text: str, stderr_te
     return stdout_log_path, stderr_log_path
 
 
+def terminate_process_group(process, *, grace_seconds: float = 10.0) -> None:
+    """SIGTERM then SIGKILL the worker's whole process group.
+
+    The torchrun *agent* and the GPU worker processes it spawns share the session
+    we start the agent in (``start_new_session=True``), so signalling the group id
+    reaps the agent and every worker at once — including the elastic agent, which
+    would otherwise respawn workers if we only killed them individually. A no-op if
+    the process already exited or the platform lacks process groups.
+    """
+    if process.poll() is not None:
+        return
+    pid = getattr(process, "pid", None)
+    if pid is None or not hasattr(os, "killpg"):
+        process.terminate()
+        return
+    try:
+        pgid = os.getpgid(pid)
+    except (ProcessLookupError, OSError):
+        return
+    try:
+        os.killpg(pgid, signal.SIGTERM)
+    except (ProcessLookupError, OSError):
+        return
+    try:
+        process.wait(timeout=grace_seconds)
+    except subprocess.TimeoutExpired:
+        try:
+            os.killpg(pgid, signal.SIGKILL)
+        except (ProcessLookupError, OSError):
+            pass
+        try:
+            process.wait(timeout=5.0)
+        except subprocess.TimeoutExpired:
+            pass
+
+
 def run_torchrun_worker(
     *,
     module: str,
@@ -79,6 +117,19 @@ def run_torchrun_worker(
         "--request-path",
         str(request_path),
     ]
+    # Run the agent in its own session so a single killpg reaps agent + workers
+    # (see terminate_process_group). A bare SIGTERM to *this* process would skip
+    # the finally block, so while the agent is alive we convert SIGTERM into a
+    # KeyboardInterrupt — but only from the main thread, where signal.signal is
+    # allowed; the original handler is restored in finally.
+    previous_sigterm = None
+    if threading.current_thread() is threading.main_thread():
+        def _raise_on_sigterm(signum, frame):  # noqa: ANN001
+            raise KeyboardInterrupt
+        try:
+            previous_sigterm = signal.signal(signal.SIGTERM, _raise_on_sigterm)
+        except (ValueError, OSError):
+            previous_sigterm = None
     process = popen_factory(
         command,
         cwd=str(Path(__file__).resolve().parents[2]),
@@ -86,43 +137,52 @@ def run_torchrun_worker(
         stderr=subprocess.PIPE,
         text=True,
         bufsize=1,
+        start_new_session=True,
     )
-    stdout_chunks: list[str] = []
-    stderr_chunks: list[str] = []
-    stdout_thread = threading.Thread(target=drain_stream_to_buffer, args=(process.stdout, stdout_chunks), daemon=True)
-    stderr_thread = threading.Thread(target=drain_stream_to_buffer, args=(process.stderr, stderr_chunks), daemon=True)
-    stdout_thread.start()
-    stderr_thread.start()
-    offsets: dict[Path, int] = {}
-    while process.poll() is None:
+    try:
+        stdout_chunks: list[str] = []
+        stderr_chunks: list[str] = []
+        stdout_thread = threading.Thread(target=drain_stream_to_buffer, args=(process.stdout, stdout_chunks), daemon=True)
+        stderr_thread = threading.Thread(target=drain_stream_to_buffer, args=(process.stderr, stderr_chunks), daemon=True)
+        stdout_thread.start()
+        stderr_thread.start()
+        offsets: dict[Path, int] = {}
+        while process.poll() is None:
+            if progress_events_path is not None:
+                events, offsets = read_progress_events(progress_events_path, offsets=offsets)
+                for event in events:
+                    emit_progress_event(event)
+                    if progress_event_callback is not None:
+                        progress_event_callback(event)
+            time.sleep(0.1)
         if progress_events_path is not None:
             events, offsets = read_progress_events(progress_events_path, offsets=offsets)
             for event in events:
                 emit_progress_event(event)
                 if progress_event_callback is not None:
                     progress_event_callback(event)
-        time.sleep(0.1)
-    if progress_events_path is not None:
-        events, offsets = read_progress_events(progress_events_path, offsets=offsets)
-        for event in events:
-            emit_progress_event(event)
-            if progress_event_callback is not None:
-                progress_event_callback(event)
-    returncode = process.wait()
-    stdout_thread.join(timeout=1.0)
-    stderr_thread.join(timeout=1.0)
-    stdout_text = "".join(stdout_chunks)
-    stderr_text = "".join(stderr_chunks)
-    stdout_log_path, stderr_log_path = write_worker_logs(module, output_dir, stdout_text, stderr_text)
-    if returncode != 0:
-        raise RuntimeError(
-            f"{failure_title}.\n"
-            f"See logs:\n"
-            f"stdout: {stdout_log_path}\n"
-            f"stderr: {stderr_log_path}\n"
-            f"stdout:\n{stdout_text}\n"
-            f"stderr:\n{stderr_text}"
-        )
+        returncode = process.wait()
+        stdout_thread.join(timeout=1.0)
+        stderr_thread.join(timeout=1.0)
+        stdout_text = "".join(stdout_chunks)
+        stderr_text = "".join(stderr_chunks)
+        stdout_log_path, stderr_log_path = write_worker_logs(module, output_dir, stdout_text, stderr_text)
+        if returncode != 0:
+            raise RuntimeError(
+                f"{failure_title}.\n"
+                f"See logs:\n"
+                f"stdout: {stdout_log_path}\n"
+                f"stderr: {stderr_log_path}\n"
+                f"stdout:\n{stdout_text}\n"
+                f"stderr:\n{stderr_text}"
+            )
+    finally:
+        # On any early exit (Ctrl-C, converted SIGTERM, RuntimeError) reap the
+        # whole worker group so no orphaned agent/workers keep holding the GPUs.
+        # No-op on the normal path: the agent has already exited.
+        terminate_process_group(process)
+        if previous_sigterm is not None:
+            signal.signal(signal.SIGTERM, previous_sigterm)
 
 
 def assign_slides_to_ranks(

{slide2vec-4.5.1 → slide2vec-4.5.2}/slide2vec.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: slide2vec
-Version: 4.5.1
+Version: 4.5.2
 Summary: Embedding of whole slide images with Foundation Models
 Author-email: Clément Grisi <clement.grisi@radboudumc.nl>
 License-Expression: Apache-2.0

{slide2vec-4.5.1 → slide2vec-4.5.2}/tests/test_progress.py

@@ -684,6 +684,64 @@ def test_run_torchrun_worker_uses_standalone_rendezvous(monkeypatch, tmp_path: P
     assert "--rdzv-endpoint" not in " ".join(command)
 
 
+def test_run_torchrun_worker_starts_new_session(monkeypatch, tmp_path: Path):
+    # The agent must run in its own session so terminate_process_group can reap
+    # the agent and every worker it spawns with a single killpg.
+    request_path = tmp_path / "request.json"
+    request_path.write_text("{}", encoding="utf-8")
+    output_dir = tmp_path / "output"
+    output_dir.mkdir()
+
+    observed = {}
+
+    class FakePopen:
+        def __init__(self, command, **kwargs):
+            observed["kwargs"] = kwargs
+            self.stdout = io.StringIO("")
+            self.stderr = io.StringIO("")
+
+        def poll(self):
+            return 0
+
+        def wait(self, timeout=None):
+            return 0
+
+    monkeypatch.setattr(distributed.time, "sleep", lambda _seconds: None)
+
+    distributed.run_torchrun_worker(
+        module="slide2vec.distributed.direct_embed_worker",
+        num_gpus=2,
+        output_dir=output_dir,
+        request_path=request_path,
+        failure_title="boom",
+        popen_factory=FakePopen,
+    )
+
+    assert observed["kwargs"].get("start_new_session") is True
+
+
+def test_terminate_process_group_reaps_whole_group():
+    # A real grandchild in the same session must die when we tear down the group.
+    parent = subprocess.Popen(
+        [
+            sys.executable,
+            "-c",
+            "import subprocess, sys, time; "
+            "subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(60)']); "
+            "time.sleep(60)",
+        ],
+        start_new_session=True,
+    )
+    import os
+
+    pgid = os.getpgid(parent.pid)
+    distributed.terminate_process_group(parent, grace_seconds=5.0)
+    assert parent.poll() is not None
+    # The whole group is gone: signalling it with 0 must raise.
+    with pytest.raises(ProcessLookupError):
+        os.killpg(pgid, 0)
+
+
 def test_reset_progress_event_logs_is_idempotent(tmp_path: Path):
     import slide2vec.runtime.distributed as distributed