[https://nvbugs/6581065][fix] Reap wedged workers during session drain

tests/test_common/session_reuse.py

@@ -40,6 +40,8 @@
 import os
 import sys
 import threading
+import time
+from typing import Protocol
 
 # The spawn snapshot is shared with the session-prefetch layer (both hand a
 # live pool to a test that did not spawn it — same invariant).
@@ -67,10 +69,18 @@
 }
 
 
-_RETIRE_THREADS: list = []
+_RETIRE_THREADS: list[threading.Thread] = []
 _RETIRE_LOCK = threading.Lock()
 
 
+class _PoolSession(Protocol):
+    _reuse_worker_pids: tuple[tuple[int, bytes | None], ...]
+
+    def shutdown(self) -> None: ...
+
+    def release_exit_joins(self) -> None: ...
+
+
 def _reap_retires(timeout: float = 60.0) -> None:
     """Join in-flight retire threads (bounded); no-op when none are running.
 
@@ -94,6 +104,58 @@ def _reap_retires(timeout: float = 60.0) -> None:
             )
 
 
+def _worker_start_time(pid: int) -> bytes | None:
+    """Read a worker's kernel start time without loading TRT-LLM eagerly."""
+    from tensorrt_llm.llmapi.mpi_session import _process_start_time
+
+    return _process_start_time(pid)
+
+
+def _kill_recorded_workers(real: _PoolSession) -> int:
+    """SIGKILL this pool's recorded workers, guarded against PID reuse.
+
+    Where the kernel supports it, the signal goes through a pidfd. Opening the
+    pidfd binds this loop to one exact process, so the start-time recheck below
+    it can no longer be invalidated by the PID being recycled before the signal
+    lands. Without pidfd the start-time recheck alone still guards the kill: that
+    leaves a microsecond-wide window, but this is the path that reaps wedged
+    workers, so refusing to signal at all would strand them on exactly the
+    platforms the reaper exists for.
+    """
+    import signal
+
+    send_via_pidfd = getattr(signal, "pidfd_send_signal", None)
+    open_pidfd = getattr(os, "pidfd_open", None)
+
+    killed = 0
+    for pid, start_time in getattr(real, "_reuse_worker_pids", ()):
+        if start_time is None:
+            continue
+        handle = None
+        if send_via_pidfd is not None and open_pidfd is not None:
+            try:
+                handle = open_pidfd(pid)
+            except (OSError, ValueError):
+                handle = None
+        try:
+            # Recheck identity AFTER pinning the handle: only kill if the
+            # process at this PID is still the worker recorded at spawn.
+            if _worker_start_time(pid) != start_time:
+                continue
+            try:
+                if handle is not None:
+                    send_via_pidfd(handle, signal.SIGKILL)
+                else:
+                    os.kill(pid, signal.SIGKILL)
+                killed += 1
+            except (ProcessLookupError, PermissionError, OSError):
+                pass
+        finally:
+            if handle is not None:
+                os.close(handle)
+    return killed
+
+
 def _prefetcher():
     """The session-prefetch singleton when that layer is wired, else None.
 
@@ -193,7 +255,7 @@ def max_uses(self) -> int:
         return int(os.environ.get("TRTLLM_TEST_REUSE_MAX_USES", "16"))
 
     @staticmethod
-    def _retire(real, broken: bool = False):
+    def _retire(real: _PoolSession, broken: bool = False) -> None:
         """Dispose of a pool in the background without blocking the test.
 
         Healthy retires (lifetime cap, stale env snapshot, duplicate cache
@@ -209,23 +271,10 @@ def _retire(real, broken: bool = False):
         needs no graceful stop; the driver reclaims GPU memory on process
         death) and then reap the client side.
         """
-        pids = getattr(real, "_reuse_worker_pids", ()) if broken else ()
 
-        def _dispose():
-            import signal
-
-            # Lazy: only runs when a pool exists, so tensorrt_llm is loaded.
-            from tensorrt_llm.llmapi.mpi_session import _process_start_time
-
-            for pid, start_time in pids:
-                # Guard against PID recycling: only kill if the process at
-                # this PID is still the worker we recorded at spawn.
-                if start_time is None or _process_start_time(pid) != start_time:
-                    continue
-                try:
-                    os.kill(pid, signal.SIGKILL)
-                except (ProcessLookupError, PermissionError):
-                    pass
+        def _dispose() -> None:
+            if broken:
+                _kill_recorded_workers(real)
             try:
                 real.shutdown()
             except Exception:
@@ -311,7 +360,10 @@ def acquire(self, real_cls, n_workers):
                 print(
                     "[session-reuse] retiring cached pool: "
                     + _describe_mismatch(
-                        real._reuse_spawn_snapshot, snap, real._reuse_uses, self.max_uses
+                        real._reuse_spawn_snapshot,
+                        snap,
+                        real._reuse_uses,
+                        self.max_uses,
                     ),
                     flush=True,
                 )
@@ -410,7 +462,7 @@ def suspend(self, suspended: bool) -> None:
         """Bypass the cache for the current test (private_mpi_session)."""
         self._suspended = suspended
 
-    def drain(self) -> None:
+    def drain(self, timeout: float = 60.0) -> None:
         """Shut down all cached pools in parallel (frees GPU/CPU footprint).
 
         Also reaps in-flight retire threads: drain runs at natural rendezvous
@@ -423,6 +475,7 @@ def drain(self) -> None:
             pools, self._pools = list(self._pools.values()), {}
         if not pools:
             return
+
         threads = [
             # daemon: a wedged pool shutdown must not keep the interpreter
             # alive at exit (a non-daemon thread would hang the CI stage).
@@ -431,16 +484,54 @@ def drain(self) -> None:
         ]
         for t in threads:
             t.start()
+
+        # Bound the whole parallel drain, rather than waiting ``timeout`` for
+        # every pool in sequence. A healthy shutdown remains graceful.
+        deadline = time.monotonic() + timeout
         for t in threads:
-            # Bounded wait: one wedged pool shutdown must not turn a drain at
-            # a shared seam (sessionfinish / RPC construction) into a
-            # suite-wide hang; a leaked wedged pool is the lesser evil.
-            t.join(timeout=60)
-            if t.is_alive():
-                print(
-                    "[session-reuse] WARNING: pool shutdown did not finish within 60s", flush=True
-                )
-        print(f"[session-reuse] drained {len(pools)} cached pool(s)", flush=True)
+            t.join(timeout=max(0.0, deadline - time.monotonic()))
+
+        wedged = [(pool, thread) for pool, thread in zip(pools, threads) if thread.is_alive()]
+        for pool, _ in wedged:
+            killed = _kill_recorded_workers(pool)
+            print(
+                "[session-reuse] WARNING: pool shutdown did not finish within "
+                f"{timeout:g}s; sent SIGKILL to {killed} recorded worker(s)",
+                flush=True,
+            )
+
+        # Killing a wedged worker should release its GPU allocation and let
+        # the already-running shutdown return. Give that original thread a
+        # short bounded reap window; never start a concurrent second shutdown.
+        reap_deadline = time.monotonic() + min(max(timeout, 1.0), 5.0)
+        for _, t in wedged:
+            t.join(timeout=max(0.0, reap_deadline - time.monotonic()))
+
+        still_alive = [(pool, thread) for pool, thread in wedged if thread.is_alive()]
+        for pool, _ in still_alive:
+            pool.release_exit_joins()
+
+        # release_exit_joins() *unblocks* the wedged shutdown (it drops the
+        # exit joins the thread is parked on) rather than merely marking it
+        # abandoned, so the thread normally finishes just after. Join it here
+        # instead of returning immediately: drain runs inside a test (RPC
+        # construction seam, opt-out setup, failure fence), so a thread that
+        # terminates a moment later takes its transport threads with it across
+        # the test boundary, and pytest-threadleak charges the leak to whatever
+        # test happens to be running next.
+        release_deadline = time.monotonic() + min(max(timeout, 1.0), 30.0)
+        for _, t in still_alive:
+            t.join(timeout=max(0.0, release_deadline - time.monotonic()))
+
+        leaked = [pool for pool, thread in still_alive if thread.is_alive()]
+        if leaked:
+            print(
+                f"[session-reuse] WARNING: {len(leaked)} pool shutdown thread(s) "
+                "remain after worker termination",
+                flush=True,
+            )
+        else:
+            print(f"[session-reuse] drained {len(pools)} cached pool(s)", flush=True)
 
 
 REUSE = SessionReuseCache()

tests/unittest/llmapi/test_session_reuse.py

@@ -13,6 +13,7 @@ def __init__(self, n_workers, wait_shutdown=False, env_overrides=None):
         self.wait_shutdown = wait_shutdown
         self.env_overrides = dict(env_overrides or {})
         self.shut = False
+        self.exit_joins_released = False
         import os
 
         # What the workers freeze at spawn: TRTLLM* forwarded from the parent
@@ -24,6 +25,9 @@ def __init__(self, n_workers, wait_shutdown=False, env_overrides=None):
     def shutdown(self):
         self.shut = True
 
+    def release_exit_joins(self) -> None:
+        self.exit_joins_released = True
+
     def shutdown_abort(self, *args, **kwargs):
         self.shut = True
 
@@ -312,6 +316,164 @@ def test_drain_shuts_cached_pools(reuse_cache):
     assert real.shut
 
 
+def test_drain_kills_recorded_worker_when_shutdown_wedges(
+    reuse_cache: SessionReuseCache, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    import signal
+    import subprocess
+    import sys
+
+    monkeypatch.setattr(session_reuse, "_worker_start_time", lambda _pid: b"owned")
+
+    class _WedgedPool(_FakePool):
+        def __init__(
+            self,
+            n_workers: int,
+            wait_shutdown: bool = False,
+            env_overrides: dict | None = None,
+        ) -> None:
+            super().__init__(n_workers, wait_shutdown, env_overrides)
+            self.worker = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(300)"])
+            self._worker_identities = ((self.worker.pid, b"owned"),)
+
+        def shutdown(self) -> None:
+            self.worker.wait()
+            self.shut = True
+
+    session = reuse_cache.acquire(_WedgedPool, 1)
+    pool = session._real
+    session.shutdown()
+    try:
+        reuse_cache.drain(timeout=0.2)
+        assert pool.worker.returncode == -signal.SIGKILL
+        assert pool.shut
+    finally:
+        if pool.worker.poll() is None:
+            pool.worker.kill()
+        pool.worker.wait()
+
+
+def test_drain_releases_exit_joins_when_shutdown_remains_wedged(
+    reuse_cache: SessionReuseCache,
+) -> None:
+    import threading
+    import time
+
+    class _ManagerWedgedPool(_FakePool):
+        def __init__(
+            self,
+            n_workers: int,
+            wait_shutdown: bool = False,
+            env_overrides: dict | None = None,
+        ) -> None:
+            super().__init__(n_workers, wait_shutdown, env_overrides)
+            self.shutdown_released = threading.Event()
+
+        def shutdown(self) -> None:
+            self.shutdown_released.wait()
+            # Dropping the exit joins only unparks the thread; the real
+            # shutdown still has to unwind (close the transport, reap its
+            # connection threads) before it returns.
+            time.sleep(0.05)
+            self.shut = True
+
+        def release_exit_joins(self) -> None:
+            super().release_exit_joins()
+            self.shutdown_released.set()
+
+    session = reuse_cache.acquire(_ManagerWedgedPool, 1)
+    pool = session._real
+    session.shutdown()
+    reuse_cache.drain(timeout=0.01)
+    assert pool.exit_joins_released
+    assert pool.shutdown_released.is_set()
+    # drain must not return while the released shutdown is still unwinding:
+    # it runs inside a test, so a thread that finishes a moment later drags
+    # its transport threads across the test boundary and pytest-threadleak
+    # blames whichever test runs next.
+    assert pool.shut
+    assert not [t for t in threading.enumerate() if t.name == "session-reuse-drain"]
+
+
+def test_kill_recorded_workers_skips_recycled_pid(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    pool = _FakePool(1)
+    pool._reuse_worker_pids = ((123, b"owned"),)
+    kills = []
+    monkeypatch.setattr(session_reuse, "_worker_start_time", lambda _pid: b"recycled")
+    monkeypatch.setattr(session_reuse.os, "kill", lambda pid, sig: kills.append((pid, sig)))
+
+    assert session_reuse._kill_recorded_workers(pool) == 0
+    assert kills == []
+
+
+def test_kill_recorded_workers_signals_through_pidfd(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The pidfd handle is pinned BEFORE the identity recheck, and closed after.
+
+    The ordering is the whole invariant: opening the pidfd first pins the process
+    so the start-time recheck cannot be invalidated by the PID being recycled
+    before the signal lands. One ordered event log is what proves that. Separate
+    per-call lists record only that each call happened, so a regression that
+    rechecks first and opens the handle afterwards still satisfies them.
+    """
+    pool = _FakePool(1)
+    pool._reuse_worker_pids = ((123, b"owned"),)
+    events: list[tuple] = []
+
+    def _start_time(pid):
+        events.append(("start_time", pid))
+        return b"owned"
+
+    def _pidfd_open(pid):
+        events.append(("open", pid))
+        return 77
+
+    monkeypatch.setattr(session_reuse, "_worker_start_time", _start_time)
+    monkeypatch.setattr(session_reuse.os, "pidfd_open", _pidfd_open, raising=False)
+    monkeypatch.setattr(session_reuse.os, "close", lambda fd: events.append(("close", fd)))
+    monkeypatch.setattr(
+        session_reuse.os,
+        "kill",
+        lambda pid, sig: pytest.fail("os.kill used while pidfd was available"),
+    )
+    import signal as _signal
+
+    monkeypatch.setattr(
+        _signal,
+        "pidfd_send_signal",
+        lambda fd, sig: events.append(("signal", fd, sig)),
+        raising=False,
+    )
+
+    assert session_reuse._kill_recorded_workers(pool) == 1
+    assert events == [
+        ("open", 123),
+        ("start_time", 123),
+        ("signal", 77, _signal.SIGKILL),
+        ("close", 77),
+    ]
+
+
+def test_kill_recorded_workers_falls_back_without_pidfd(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """No pidfd support must still reap: the wedged worker is the whole point."""
+    pool = _FakePool(1)
+    pool._reuse_worker_pids = ((123, b"owned"),)
+    kills = []
+    monkeypatch.setattr(session_reuse, "_worker_start_time", lambda _pid: b"owned")
+    monkeypatch.delattr(session_reuse.os, "pidfd_open", raising=False)
+    monkeypatch.setattr(session_reuse.os, "kill", lambda pid, sig: kills.append((pid, sig)))
+
+    import signal as _signal
+
+    assert session_reuse._kill_recorded_workers(pool) == 1
+    assert kills == [(123, _signal.SIGKILL)]
+
+
 def test_autodeploy_nodeids_are_private():
     from test_common.session_reuse_hooks import _is_private_nodeid