From d9cdc0012354236a5b419963f6582aa61b54125c Mon Sep 17 00:00:00 2001
From: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
Date: Fri, 14 Aug 2026 12:27:31 +0000
Subject: [PATCH 1/7] [https://nvbugs/6581065][fix] Reap wedged session workers
 during drain

Signed-off-by: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
---
 tests/test_common/session_reuse.py          | 99 +++++++++++++++------
 tests/unittest/llmapi/test_session_reuse.py | 34 +++++++
 2 files changed, 105 insertions(+), 28 deletions(-)

diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py
index ea1b6272a5c5..9302ee0050e5 100644
--- a/tests/test_common/session_reuse.py
+++ b/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,16 @@
 }
 
 
-_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 _reap_retires(timeout: float = 60.0) -> None:
     """Join in-flight retire threads (bounded); no-op when none are running.
 
@@ -94,6 +102,31 @@ 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."""
+    import signal
+
+    killed = 0
+    for pid, start_time in getattr(real, "_reuse_worker_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 _worker_start_time(pid) != start_time:
+            continue
+        try:
+            os.kill(pid, signal.SIGKILL)
+            killed += 1
+        except (ProcessLookupError, PermissionError):
+            pass
+    return killed
+
+
 def _prefetcher():
     """The session-prefetch singleton when that layer is wired, else None.
 
@@ -193,7 +226,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 +242,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:
@@ -410,7 +430,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 +443,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 +452,38 @@ 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(timeout, 5.0)
+        for _, t in wedged:
+            t.join(timeout=max(0.0, reap_deadline - time.monotonic()))
+
+        still_alive = sum(t.is_alive() for t in threads)
+        if still_alive:
+            print(
+                f"[session-reuse] WARNING: {still_alive} 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()
diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py
index ef693c291c10..5741fb830b8d 100644
--- a/tests/unittest/llmapi/test_session_reuse.py
+++ b/tests/unittest/llmapi/test_session_reuse.py
@@ -312,6 +312,40 @@ 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_autodeploy_nodeids_are_private():
     from test_common.session_reuse_hooks import _is_private_nodeid
 

From f2cd3a36ba9b3f3f9fc971c008fe93320d329e8f Mon Sep 17 00:00:00 2001
From: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
Date: Tue, 18 Aug 2026 00:10:47 +0000
Subject: [PATCH 2/7] [https://nvbugs/6581065][fix] Address drain review
 feedback

Signed-off-by: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
---
 tests/integration/test_lists/waives.txt     |  1 -
 tests/test_common/session_reuse.py          | 37 ++++++++---
 tests/unittest/llmapi/test_session_reuse.py | 71 +++++++++++++++++++--
 3 files changed, 94 insertions(+), 15 deletions(-)

diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt
index a42556295097..ecc1c4d97f8e 100644
--- a/tests/integration/test_lists/waives.txt
+++ b/tests/integration/test_lists/waives.txt
@@ -77,7 +77,6 @@ accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_
 accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torch_compile=False-enable_gemm_allreduce_fusion=False] SKIP (https://nvbugs/6428089)
 accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torch_compile=True-enable_gemm_allreduce_fusion=False] SKIP (https://nvbugs/6427411)
 accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torch_compile=True-enable_gemm_allreduce_fusion=True] SKIP (https://nvbugs/6211441)
-accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp SKIP (https://nvbugs/6581065)
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TP4_PP2] SKIP (https://nvbugs/6427411)
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_4gpus_static_eplb[moe_backend=CUTLASS] SKIP (https://nvbugs/6535767)
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_parallelism[ADP2_PP2] SKIP (https://nvbugs/6427411)
diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py
index 9302ee0050e5..9636234dcf07 100644
--- a/tests/test_common/session_reuse.py
+++ b/tests/test_common/session_reuse.py
@@ -46,7 +46,10 @@
 # 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).
 from test_common._session_utils import _isinstance_transparent_shim, _spawn_snapshot
-from test_common.grouped_test_utils import reset_worker_torch_compile_state, submit_sync_per_worker
+from test_common.grouped_test_utils import (
+    reset_worker_torch_compile_state,
+    submit_sync_per_worker,
+)
 
 # The only places in the library that construct MpiPoolSession for a bare
 # LLM(...); tests passing their own _mpi_session never reach these lines.
@@ -78,6 +81,8 @@ class _PoolSession(Protocol):
 
     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.
@@ -146,7 +151,9 @@ def _describe_mismatch(spawn_snap, now_snap, uses, max_uses):
         return f"lifetime cap reached ({uses}/{max_uses} uses)"
     spawn_env, spawn_path = spawn_snap
     now_env, now_path = now_snap
-    changed = [k for k in set(spawn_env) | set(now_env) if spawn_env.get(k) != now_env.get(k)]
+    changed = [
+        k for k in set(spawn_env) | set(now_env) if spawn_env.get(k) != now_env.get(k)
+    ]
     if changed:
         return f"env changed since spawn: {sorted(changed)[:6]}"
     if spawn_path != now_path:
@@ -273,7 +280,9 @@ def install_pool_factory_if_loaded(self) -> None:
             return  # fully installed: skip the env reads and module scan
         if not self.enabled:
             return
-        pending = [n for n in _ALL_PATCH_TARGETS if n in sys.modules and n not in self._patched]
+        pending = [
+            n for n in _ALL_PATCH_TARGETS if n in sys.modules and n not in self._patched
+        ]
         if not pending:
             return
         from tensorrt_llm.llmapi.mpi_session import MpiPoolSession as real_cls
@@ -331,7 +340,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,
                 )
@@ -395,7 +407,9 @@ def _spawn_fresh(self, real_cls, n_workers):
             # collection still fails, unidentified workers may remain alive;
             # an immediate retry would overlap another MPI bootstrap with
             # them, so propagate the fail-closed error.
-            real = real_cls(n_workers=n_workers, wait_shutdown=True, env_overrides=overrides)
+            real = real_cls(
+                n_workers=n_workers, wait_shutdown=True, env_overrides=overrides
+            )
         if prefetcher is not None:
             try:
                 # Restock only after the current pool is ready. On a shadow
@@ -459,7 +473,9 @@ def drain(self, timeout: float = 60.0) -> None:
         for t in threads:
             t.join(timeout=max(0.0, deadline - time.monotonic()))
 
-        wedged = [(pool, thread) for pool, thread in zip(pools, threads) if thread.is_alive()]
+        wedged = [
+            (pool, thread) for pool, thread in zip(pools, threads) if thread.is_alive()
+        ]
         for pool, _ in wedged:
             killed = _kill_recorded_workers(pool)
             print(
@@ -471,14 +487,17 @@ def drain(self, timeout: float = 60.0) -> None:
         # 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(timeout, 5.0)
+        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 = sum(t.is_alive() for t in threads)
+        still_alive = [(pool, thread) for pool, thread in wedged if thread.is_alive()]
+        for pool, _ in still_alive:
+            pool.release_exit_joins()
+
         if still_alive:
             print(
-                f"[session-reuse] WARNING: {still_alive} pool shutdown thread(s) "
+                f"[session-reuse] WARNING: {len(still_alive)} pool shutdown thread(s) "
                 "remain after worker termination",
                 flush=True,
             )
diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py
index 5741fb830b8d..bc17caddbb0b 100644
--- a/tests/unittest/llmapi/test_session_reuse.py
+++ b/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):
+        self.exit_joins_released = True
+
     def shutdown_abort(self, *args, **kwargs):
         self.shut = True
 
@@ -35,7 +39,9 @@ def reuse_cache(monkeypatch):
     cache = SessionReuseCache()
     # No real MPI / NVML in pure-logic tests: record the calls instead.
     resets = []
-    monkeypatch.setattr(session_reuse, "submit_sync_per_worker", lambda s, fn: resets.append(s))
+    monkeypatch.setattr(
+        session_reuse, "submit_sync_per_worker", lambda s, fn: resets.append(s)
+    )
     cache.resets = resets
 
     # Hermetic: the REAL prefetcher singleton must not start background MPI
@@ -105,7 +111,9 @@ def test_cache_miss_takes_prefetched_shadow(reuse_cache):
     # A shadow armed at the PREVIOUS miss is consumed instantly on this one
     # (no synchronous spawn), and a replacement is restocked for the next
     # miss with the worker-side weight-cache overlay.
-    shadow = _FakePool(2, wait_shutdown=True, env_overrides={"TRTLLM_HF_WEIGHT_CACHE": "1"})
+    shadow = _FakePool(
+        2, wait_shutdown=True, env_overrides={"TRTLLM_HF_WEIGHT_CACHE": "1"}
+    )
     reuse_cache.prefetch.shadow = shadow
     s = reuse_cache.acquire(_FakePool, 2)
     assert s._real is shadow
@@ -323,10 +331,15 @@ def test_drain_kills_recorded_worker_when_shutdown_wedges(
 
     class _WedgedPool(_FakePool):
         def __init__(
-            self, n_workers: int, wait_shutdown: bool = False, env_overrides: dict | None = None
+            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 = subprocess.Popen(
+                [sys.executable, "-c", "import time; time.sleep(300)"]
+            )
             self._worker_identities = ((self.worker.pid, b"owned"),)
 
         def shutdown(self) -> None:
@@ -346,6 +359,52 @@ def shutdown(self) -> None:
         pool.worker.wait()
 
 
+def test_drain_releases_exit_joins_when_shutdown_remains_wedged(
+    reuse_cache: SessionReuseCache,
+) -> None:
+    import threading
+
+    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()
+            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.wait(timeout=1.0)
+
+
+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_autodeploy_nodeids_are_private():
     from test_common.session_reuse_hooks import _is_private_nodeid
 
@@ -356,7 +415,9 @@ def test_autodeploy_nodeids_are_private():
     assert _is_private_nodeid(
         "examples/test_ad_guided_decoding.py::test_autodeploy_guided_decoding_main_json"
     )
-    assert _is_private_nodeid("unittest/_torch/auto_deploy/unit/singlegpu/test_x.py::test_y")
+    assert _is_private_nodeid(
+        "unittest/_torch/auto_deploy/unit/singlegpu/test_x.py::test_y"
+    )
     assert not _is_private_nodeid(
         "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[a]"
     )

From 6fec8c9df536af296b872f56a58e9f35a1f189f4 Mon Sep 17 00:00:00 2001
From: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
Date: Tue, 18 Aug 2026 03:34:50 +0000
Subject: [PATCH 3/7] [https://nvbugs/6581065][chore] Reformat for the ruff
 line-length-100 config from main

The target sync picked up main's ruff line-length=100, so the lines this
PR touched no longer matched the repository formatter. Whitespace only;
no behavior change.

Signed-off-by: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
---
 tests/test_common/session_reuse.py          | 21 +++++----------------
 tests/unittest/llmapi/test_session_reuse.py | 20 +++++---------------
 2 files changed, 10 insertions(+), 31 deletions(-)

diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py
index 9636234dcf07..6aa273390ed1 100644
--- a/tests/test_common/session_reuse.py
+++ b/tests/test_common/session_reuse.py
@@ -46,10 +46,7 @@
 # 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).
 from test_common._session_utils import _isinstance_transparent_shim, _spawn_snapshot
-from test_common.grouped_test_utils import (
-    reset_worker_torch_compile_state,
-    submit_sync_per_worker,
-)
+from test_common.grouped_test_utils import reset_worker_torch_compile_state, submit_sync_per_worker
 
 # The only places in the library that construct MpiPoolSession for a bare
 # LLM(...); tests passing their own _mpi_session never reach these lines.
@@ -151,9 +148,7 @@ def _describe_mismatch(spawn_snap, now_snap, uses, max_uses):
         return f"lifetime cap reached ({uses}/{max_uses} uses)"
     spawn_env, spawn_path = spawn_snap
     now_env, now_path = now_snap
-    changed = [
-        k for k in set(spawn_env) | set(now_env) if spawn_env.get(k) != now_env.get(k)
-    ]
+    changed = [k for k in set(spawn_env) | set(now_env) if spawn_env.get(k) != now_env.get(k)]
     if changed:
         return f"env changed since spawn: {sorted(changed)[:6]}"
     if spawn_path != now_path:
@@ -280,9 +275,7 @@ def install_pool_factory_if_loaded(self) -> None:
             return  # fully installed: skip the env reads and module scan
         if not self.enabled:
             return
-        pending = [
-            n for n in _ALL_PATCH_TARGETS if n in sys.modules and n not in self._patched
-        ]
+        pending = [n for n in _ALL_PATCH_TARGETS if n in sys.modules and n not in self._patched]
         if not pending:
             return
         from tensorrt_llm.llmapi.mpi_session import MpiPoolSession as real_cls
@@ -407,9 +400,7 @@ def _spawn_fresh(self, real_cls, n_workers):
             # collection still fails, unidentified workers may remain alive;
             # an immediate retry would overlap another MPI bootstrap with
             # them, so propagate the fail-closed error.
-            real = real_cls(
-                n_workers=n_workers, wait_shutdown=True, env_overrides=overrides
-            )
+            real = real_cls(n_workers=n_workers, wait_shutdown=True, env_overrides=overrides)
         if prefetcher is not None:
             try:
                 # Restock only after the current pool is ready. On a shadow
@@ -473,9 +464,7 @@ def drain(self, timeout: float = 60.0) -> None:
         for t in threads:
             t.join(timeout=max(0.0, deadline - time.monotonic()))
 
-        wedged = [
-            (pool, thread) for pool, thread in zip(pools, threads) if thread.is_alive()
-        ]
+        wedged = [(pool, thread) for pool, thread in zip(pools, threads) if thread.is_alive()]
         for pool, _ in wedged:
             killed = _kill_recorded_workers(pool)
             print(
diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py
index bc17caddbb0b..4a7761049bc2 100644
--- a/tests/unittest/llmapi/test_session_reuse.py
+++ b/tests/unittest/llmapi/test_session_reuse.py
@@ -39,9 +39,7 @@ def reuse_cache(monkeypatch):
     cache = SessionReuseCache()
     # No real MPI / NVML in pure-logic tests: record the calls instead.
     resets = []
-    monkeypatch.setattr(
-        session_reuse, "submit_sync_per_worker", lambda s, fn: resets.append(s)
-    )
+    monkeypatch.setattr(session_reuse, "submit_sync_per_worker", lambda s, fn: resets.append(s))
     cache.resets = resets
 
     # Hermetic: the REAL prefetcher singleton must not start background MPI
@@ -111,9 +109,7 @@ def test_cache_miss_takes_prefetched_shadow(reuse_cache):
     # A shadow armed at the PREVIOUS miss is consumed instantly on this one
     # (no synchronous spawn), and a replacement is restocked for the next
     # miss with the worker-side weight-cache overlay.
-    shadow = _FakePool(
-        2, wait_shutdown=True, env_overrides={"TRTLLM_HF_WEIGHT_CACHE": "1"}
-    )
+    shadow = _FakePool(2, wait_shutdown=True, env_overrides={"TRTLLM_HF_WEIGHT_CACHE": "1"})
     reuse_cache.prefetch.shadow = shadow
     s = reuse_cache.acquire(_FakePool, 2)
     assert s._real is shadow
@@ -337,9 +333,7 @@ def __init__(
             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 = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(300)"])
             self._worker_identities = ((self.worker.pid, b"owned"),)
 
         def shutdown(self) -> None:
@@ -397,9 +391,7 @@ def test_kill_recorded_workers_skips_recycled_pid(
     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))
-    )
+    monkeypatch.setattr(session_reuse.os, "kill", lambda pid, sig: kills.append((pid, sig)))
 
     assert session_reuse._kill_recorded_workers(pool) == 0
     assert kills == []
@@ -415,9 +407,7 @@ def test_autodeploy_nodeids_are_private():
     assert _is_private_nodeid(
         "examples/test_ad_guided_decoding.py::test_autodeploy_guided_decoding_main_json"
     )
-    assert _is_private_nodeid(
-        "unittest/_torch/auto_deploy/unit/singlegpu/test_x.py::test_y"
-    )
+    assert _is_private_nodeid("unittest/_torch/auto_deploy/unit/singlegpu/test_x.py::test_y")
     assert not _is_private_nodeid(
         "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[a]"
     )

From 190a33dc121f049a2098c84fc472ea87fa45281a Mon Sep 17 00:00:00 2001
From: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
Date: Tue, 18 Aug 2026 05:32:29 +0000
Subject: [PATCH 4/7] [https://nvbugs/6581065][chore] Annotate
 _FakePool.release_exit_joins return type

Matches the subclass override that already declares `-> None`.

Signed-off-by: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
---
 tests/unittest/llmapi/test_session_reuse.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py
index 4a7761049bc2..a560e00ae5c9 100644
--- a/tests/unittest/llmapi/test_session_reuse.py
+++ b/tests/unittest/llmapi/test_session_reuse.py
@@ -25,7 +25,7 @@ def __init__(self, n_workers, wait_shutdown=False, env_overrides=None):
     def shutdown(self):
         self.shut = True
 
-    def release_exit_joins(self):
+    def release_exit_joins(self) -> None:
         self.exit_joins_released = True
 
     def shutdown_abort(self, *args, **kwargs):

From 6f471dc35e4df421987a6fbcdb080c7cc9805331 Mon Sep 17 00:00:00 2001
From: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
Date: Wed, 19 Aug 2026 08:03:02 +0000
Subject: [PATCH 5/7] [https://nvbugs/6581065][fix] Pin worker identity with a
 pidfd before SIGKILL

Opening a pidfd binds the reaper to one exact process, so the start-time
recheck can no longer be invalidated by PID recycling before the signal
lands. Kernels without pidfd keep the start-time-guarded kill: this path
reaps wedged workers, so refusing to signal would strand them.

Signed-off-by: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
---
 tests/test_common/session_reuse.py          | 43 ++++++++++++++----
 tests/unittest/llmapi/test_session_reuse.py | 49 +++++++++++++++++++++
 2 files changed, 84 insertions(+), 8 deletions(-)

diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py
index 6aa273390ed1..c3bef8eb3b50 100644
--- a/tests/test_common/session_reuse.py
+++ b/tests/test_common/session_reuse.py
@@ -112,20 +112,47 @@ def _worker_start_time(pid: int) -> bytes | None:
 
 
 def _kill_recorded_workers(real: _PoolSession) -> int:
-    """SIGKILL this pool's recorded workers, guarded against PID reuse."""
+    """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", ()):
-        # 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 _worker_start_time(pid) != start_time:
+        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:
-            os.kill(pid, signal.SIGKILL)
-            killed += 1
-        except (ProcessLookupError, PermissionError):
-            pass
+            # 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
 
 
diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py
index a560e00ae5c9..c97a5265ac9a 100644
--- a/tests/unittest/llmapi/test_session_reuse.py
+++ b/tests/unittest/llmapi/test_session_reuse.py
@@ -397,6 +397,55 @@ def test_kill_recorded_workers_skips_recycled_pid(
     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."""
+    pool = _FakePool(1)
+    pool._reuse_worker_pids = ((123, b"owned"),)
+    opened, signalled, closed = [], [], []
+    monkeypatch.setattr(session_reuse, "_worker_start_time", lambda _pid: b"owned")
+    monkeypatch.setattr(
+        session_reuse.os, "pidfd_open", lambda pid: opened.append(pid) or 77, raising=False
+    )
+    monkeypatch.setattr(session_reuse.os, "close", lambda fd: closed.append(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: signalled.append((fd, sig)),
+        raising=False,
+    )
+
+    assert session_reuse._kill_recorded_workers(pool) == 1
+    assert opened == [123]
+    assert signalled == [(77, _signal.SIGKILL)]
+    assert closed == [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
 

From 805fd1790ad29969567a7234626ae53702b66b06 Mon Sep 17 00:00:00 2001
From: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
Date: Wed, 19 Aug 2026 11:15:56 +0000
Subject: [PATCH 6/7] [https://nvbugs/6581065][fix] Assert pidfd ordering and
 drop a stale re-waive

The pidfd test recorded opens, signals and closes in separate lists, which
prove each call happened but not that the handle was pinned before the
identity recheck -- the one property the fix exists to guarantee. A single
ordered event log now fails if the recheck moves ahead of pidfd_open.

Also drop the test_nvfp4_4gpus_hopper_w4a16 waiver this branch carried: it
was unwaived on main by bd90276f6, and resolving the merge kept the stale
side, silently reverting that unwaive.

Signed-off-by: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
---
 tests/integration/test_lists/waives.txt     |  1 -
 tests/unittest/llmapi/test_session_reuse.py | 39 +++++++++++++++------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt
index 2fa295bc8cbe..72488310354f 100644
--- a/tests/integration/test_lists/waives.txt
+++ b/tests/integration/test_lists/waives.txt
@@ -62,7 +62,6 @@ accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp
 accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] SKIP (https://nvbugs/6278337)
 accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torch_compile=False-enable_gemm_allreduce_fusion=False] SKIP (https://nvbugs/6428089)
 accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torch_compile=True-enable_gemm_allreduce_fusion=True] SKIP (https://nvbugs/6211441)
-accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_hopper_w4a16 SKIP (https://nvbugs/6478723)
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp SKIP (https://nvbugs/6581065)
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_4gpus_static_eplb[moe_backend=CUTLASS] SKIP (https://nvbugs/6535767)
 accuracy/test_llm_api_pytorch.py::TestQwen3NextInstruct::test_nvfp4[tp1_block_reuse-cutlass] SKIP (https://nvbugs/6535767)
diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py
index c97a5265ac9a..955a4ef68ad7 100644
--- a/tests/unittest/llmapi/test_session_reuse.py
+++ b/tests/unittest/llmapi/test_session_reuse.py
@@ -400,15 +400,29 @@ def test_kill_recorded_workers_skips_recycled_pid(
 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 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"),)
-    opened, signalled, closed = [], [], []
-    monkeypatch.setattr(session_reuse, "_worker_start_time", lambda _pid: b"owned")
-    monkeypatch.setattr(
-        session_reuse.os, "pidfd_open", lambda pid: opened.append(pid) or 77, raising=False
-    )
-    monkeypatch.setattr(session_reuse.os, "close", lambda fd: closed.append(fd))
+    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",
@@ -419,14 +433,17 @@ def test_kill_recorded_workers_signals_through_pidfd(
     monkeypatch.setattr(
         _signal,
         "pidfd_send_signal",
-        lambda fd, sig: signalled.append((fd, sig)),
+        lambda fd, sig: events.append(("signal", fd, sig)),
         raising=False,
     )
 
     assert session_reuse._kill_recorded_workers(pool) == 1
-    assert opened == [123]
-    assert signalled == [(77, _signal.SIGKILL)]
-    assert closed == [77]
+    assert events == [
+        ("open", 123),
+        ("start_time", 123),
+        ("signal", 77, _signal.SIGKILL),
+        ("close", 77),
+    ]
 
 
 def test_kill_recorded_workers_falls_back_without_pidfd(

From 1c170bbb6e1a2c782dad9cf46167078bfed818ce Mon Sep 17 00:00:00 2001
From: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
Date: Thu, 20 Aug 2026 00:14:14 +0000
Subject: [PATCH 7/7] [https://nvbugs/6581065][fix] Join the released shutdown
 before drain returns

release_exit_joins() unparks the wedged shutdown thread rather than just
marking it abandoned, but drain() returned immediately afterwards. The
thread then unwound outside drain -- and drain runs inside a test (RPC
construction seam, opt-out setup, failure fence), so its transport
threads crossed the test boundary and pytest-threadleak charged the leak
to whichever test ran next.

Add a bounded join after the release so the shutdown finishes inside
drain, and only warn about pools that are still alive after it.

Signed-off-by: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
---
 tests/test_common/session_reuse.py          | 17 +++++++++++++++--
 tests/unittest/llmapi/test_session_reuse.py | 13 ++++++++++++-
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py
index c3bef8eb3b50..dc52ac40d2f9 100644
--- a/tests/test_common/session_reuse.py
+++ b/tests/test_common/session_reuse.py
@@ -511,9 +511,22 @@ def drain(self, timeout: float = 60.0) -> None:
         for pool, _ in still_alive:
             pool.release_exit_joins()
 
-        if still_alive:
+        # 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(still_alive)} pool shutdown thread(s) "
+                f"[session-reuse] WARNING: {len(leaked)} pool shutdown thread(s) "
                 "remain after worker termination",
                 flush=True,
             )
diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py
index 955a4ef68ad7..a4f7b403f5ea 100644
--- a/tests/unittest/llmapi/test_session_reuse.py
+++ b/tests/unittest/llmapi/test_session_reuse.py
@@ -357,6 +357,7 @@ def test_drain_releases_exit_joins_when_shutdown_remains_wedged(
     reuse_cache: SessionReuseCache,
 ) -> None:
     import threading
+    import time
 
     class _ManagerWedgedPool(_FakePool):
         def __init__(
@@ -370,6 +371,10 @@ def __init__(
 
         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:
@@ -381,7 +386,13 @@ def release_exit_joins(self) -> None:
     session.shutdown()
     reuse_cache.drain(timeout=0.01)
     assert pool.exit_joins_released
-    assert pool.shutdown_released.wait(timeout=1.0)
+    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(