[TRTLLM-15302][fix] Reject dead prefetched MPI pools

tests/integration/test_lists/test-db/l0_cpu.yml

@@ -96,6 +96,7 @@ l0_cpu:
   - unittest/llmapi/test_sampling_params.py
   - unittest/llmapi/test_serialization.py
   - unittest/llmapi/test_serve_report_addr.py
+  - unittest/llmapi/test_session_prefetcher.py
   - unittest/llmapi/test_utils.py
   - unittest/metrics/test_collector.py
   - unittest/models/test_quant_config_utils.py

tests/test_common/session_prefetcher.py

@@ -276,6 +276,63 @@ class _Built(NamedTuple):
     snapshot: object
 
 
+def _prefetched_workers_alive(session: object) -> bool:
+    """Whether every recorded worker is still the process spawned for this pool.
+
+    ``MpiPoolSession(wait_shutdown=True)`` records a complete set of
+    ``(pid, start_time)`` identities before a pool can be published. Checking
+    both values rejects workers that exited after publication without
+    mistaking a recycled PID for the original process.
+    """
+    identities = getattr(session, "_worker_identities", ())
+    if len(identities) != getattr(session, "n_workers", None):
+        return False
+    mpi_session = sys.modules.get("tensorrt_llm.llmapi.mpi_session")
+    process_start_time = getattr(mpi_session, "_process_start_time", None)
+    if process_start_time is None:
+        return False
+    return all(
+        start_time is not None and process_start_time(pid) == start_time
+        for pid, start_time in identities
+    )
+
+
+def _reap_dead_pool(session: object) -> None:
+    """Abandon a part-dead pool and SIGKILL whichever workers are still up.
+
+    ``abandon()`` only disconnects the parent side. With a rank already gone the
+    manager thread stays wedged in MPI, so the survivors are never told to exit
+    and keep their GPU memory while the replacement pool spawns on top of them.
+    Mirrors ``MpiPoolSession._teardown_unidentified_pool``, recycling guard
+    included.
+    """
+    import signal
+
+    try:
+        session.abandon()
+    except Exception as exc:  # noqa: BLE001
+        # Broad on purpose: this pool is already known part-dead, so abandon()
+        # reaches into mpi4py with a rank missing and can fail in ways MPI does
+        # not enumerate. Swallowing is required -- the SIGKILL sweep below is
+        # what actually reclaims the survivors' GPU memory, and it must run
+        # whatever abandon() did. Surface it instead of passing silently.
+        print(
+            f"[session-prefetch] WARNING: abandon() failed on a part-dead pool: {exc!r}",
+            flush=True,
+        )
+    mpi_session = sys.modules.get("tensorrt_llm.llmapi.mpi_session")
+    process_start_time = getattr(mpi_session, "_process_start_time", None)
+    if process_start_time is None:
+        return
+    for pid, start_time in getattr(session, "_worker_identities", ()):
+        if start_time is None or process_start_time(pid) != start_time:
+            continue  # already gone, or the PID was recycled
+        try:
+            os.kill(pid, signal.SIGKILL)
+        except (ProcessLookupError, PermissionError):
+            pass
+
+
 class SessionPrefetcher:
     def __init__(self):
         self._lock = threading.Lock()
@@ -497,6 +554,24 @@ def take(self, spec: int) -> object | None:
         built = self._drain()
         if built is None:
             return None
+        if not _prefetched_workers_alive(built.session):
+            # A worker may die after the background build publishes its pool.
+            # Never hand that unusable executor to the next test: its queued
+            # initialization task would have no worker and can wait until the
+            # outer pytest timeout. The dead world cannot be joined, so
+            # abandon it before the synchronous fallback.
+            self.stats["pools_discarded_dead"] += 1
+            print(
+                "[session-prefetch] discarding prefetched pool with dead worker",
+                flush=True,
+            )
+            threading.Thread(
+                target=_reap_dead_pool,
+                args=(built.session,),
+                daemon=True,
+                name="session-prefetch-discard-dead",
+            ).start()
+            return None
         if built.spec == spec and built.snapshot == _spawn_snapshot():
             # An instant handover is safe against the previous worker's GPU
             # memory: every pool these layers hand out is built with

tests/unittest/llmapi/test_session_prefetcher.py

@@ -13,6 +13,8 @@
 from test_common import session_prefetcher
 from test_common.session_prefetcher import SessionPrefetcher, warm_page_cache
 
+pytestmark = pytest.mark.cpu_only
+
 
 class _FakeMarker:
     def __init__(self, *args):
@@ -64,6 +66,7 @@ def _fake_build(self, spec, gen, env_overlay=None):
 
     monkeypatch.setattr(SessionPrefetcher, "_build", _fake_build)
     monkeypatch.setattr(SessionPrefetcher, "_warm", lambda self, d: warmed.append(d))
+    monkeypatch.setattr(session_prefetcher, "_prefetched_workers_alive", lambda session: True)
     p.built, p.warmed, p.overlays = built, warmed, overlays
     return p
 
@@ -73,10 +76,15 @@ def __init__(self, n_workers, wait_shutdown=False):
         self.n_workers = n_workers
         self.wait_shutdown = wait_shutdown
         self.shut = False
+        self.abandoned = False
 
     def shutdown(self):
         self.shut = True
 
+    def abandon(self) -> None:
+        self.abandoned = True
+        self.shutdown()
+
 
 def _arm(prefetcher, pool, spec=4):
     """Publish ``pool`` into the shadow slot through the real API."""
@@ -300,6 +308,46 @@ def test_factory_hit_hands_over_shadow(prefetcher):
     assert factory(4) is pool  # prefetched pool handed over
 
 
+def test_factory_dead_shadow_builds_fresh_pool(prefetcher, monkeypatch) -> None:
+    dead = _FakePool(1)
+    _arm(prefetcher, dead, spec=1)
+    monkeypatch.setattr(session_prefetcher, "_prefetched_workers_alive", lambda session: False)
+
+    factory = prefetcher._make_factory(_FakePool)
+    replacement = factory(1)
+
+    assert replacement is not dead
+    assert replacement.wait_shutdown
+    # The discard is backgrounded so the replacement spawn is not stuck behind a
+    # 30s wait on workers that will never answer; join it before asserting.
+    for thread in threading.enumerate():
+        if thread.name == "session-prefetch-discard-dead":
+            thread.join(timeout=10)
+    assert dead.abandoned and dead.shut
+    assert prefetcher.stats["pools_discarded_dead"] == 1
+    assert prefetcher.stats["pools_handed_over"] == 0
+
+
+def test_prefetched_worker_identity_check_rejects_dead_or_recycled_pid(monkeypatch) -> None:
+    starts = {101: b"start-a", 102: b"start-b"}
+    fake_mpi_session = types.SimpleNamespace(_process_start_time=starts.get)
+    monkeypatch.setitem(sys.modules, "tensorrt_llm.llmapi.mpi_session", fake_mpi_session)
+    pool = types.SimpleNamespace(
+        n_workers=2,
+        _worker_identities=((101, b"start-a"), (102, b"start-b")),
+    )
+
+    assert session_prefetcher._prefetched_workers_alive(pool)
+    pool._worker_identities = ((101, b"start-a"),)
+    assert not session_prefetcher._prefetched_workers_alive(pool)
+    pool._worker_identities = ((101, b"start-a"), (102, b"start-b"))
+
+    starts.pop(102)
+    assert not session_prefetcher._prefetched_workers_alive(pool)
+    starts[102] = b"recycled-process"
+    assert not session_prefetcher._prefetched_workers_alive(pool)
+
+
 def test_take_spec_mismatch_returns_none(prefetcher):
     pool = _FakePool(4)
     _arm(prefetcher, pool, spec=4)