[https://nvbugs/6523767][fix] Size MPI worker-identity barrier timeout to cover worker bootstrap

tensorrt_llm/llmapi/mpi_session.py

@@ -1,5 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import abc
 import itertools
+import math
 import os
 import socket
 import sys
@@ -220,6 +224,36 @@ def _process_start_time(pid: int) -> Optional[bytes]:
         return None
 
 
+_DEFAULT_IDENTITY_TIMEOUT = 300.0
+
+
+def _identity_barrier_timeout() -> float:
+    """Deadline for the ``wait_shutdown`` worker-identity barrier, in seconds.
+
+    The barrier itself completes in milliseconds, but it is the first work ever
+    submitted to a freshly built ``MPIPoolExecutor``, and mpi4py spawns lazily
+    from its manager thread — so this deadline really bounds the whole worker
+    bootstrap: process spawn plus ``import tensorrt_llm``, measured at ~50-65s
+    on an idle node and up to ~117s on a contended one. Hence a ceiling sized
+    against bootstrap cost rather than barrier latency. The test-session
+    prefetcher derives its own wait budget from this value so it cannot abandon
+    a bootstrap that this layer still considers healthy.
+    ``TRTLLM_MPI_IDENTITY_TIMEOUT`` overrides it.
+    """
+    raw = os.environ.get("TRTLLM_MPI_IDENTITY_TIMEOUT")
+    if not raw:
+        return _DEFAULT_IDENTITY_TIMEOUT
+    try:
+        value = float(raw)
+        if math.isfinite(value) and value > 0:
+            return value
+    except ValueError:
+        pass
+    logger.warning(f"Ignoring invalid TRTLLM_MPI_IDENTITY_TIMEOUT={raw!r}; "
+                   f"using {_DEFAULT_IDENTITY_TIMEOUT}s")
+    return _DEFAULT_IDENTITY_TIMEOUT
+
+
 def _worker_identity_barrier():
     """Runs inside a pool worker; module-level so it is picklable.
 
@@ -315,12 +349,13 @@ def _collect_worker_identities(self) -> Tuple:
         cancel the pending tasks). Instead of handing out such a pool, tear
         it down and raise; callers fall back to a fresh spawn.
         """
+        timeout = _identity_barrier_timeout()
         try:
             futures = [
                 self.mpi_pool.submit(_worker_identity_barrier)
                 for _ in range(self.n_workers)
             ]
-            done, not_done = futures_wait(futures, timeout=60.0)
+            done, not_done = futures_wait(futures, timeout=timeout)
             identities = tuple(f.result() for f in done)
         except Exception as e:
             self._teardown_unidentified_pool(())
@@ -336,7 +371,9 @@ def _collect_worker_identities(self) -> Tuple:
                 "MpiPoolSession(wait_shutdown=True): worker identity "
                 f"collection incomplete ({len(identities)}/{self.n_workers} "
                 "valid identities); pool torn down instead of handing out a "
-                "session that cannot honor the wait_shutdown contract")
+                "session that cannot honor the wait_shutdown contract. Raise "
+                "TRTLLM_MPI_IDENTITY_TIMEOUT if worker bootstrap is merely "
+                f"slow (deadline was {timeout}s)")
         return identities
 
     def _teardown_unidentified_pool(self, partial_identities: Tuple) -> None:

tests/test_common/session_prefetcher.py

@@ -47,6 +47,7 @@
 """
 
 import glob
+import math
 import os
 import sys
 import threading
@@ -70,6 +71,39 @@
     "tensorrt_llm.llmapi.llm",
 )
 
+# Identity collection is followed by one lightweight diagnostic submitted to
+# the workers. Keep enough room for that hand-off after the lower-level
+# bootstrap deadline expires.
+_SHADOW_BUILD_FINISH_GRACE = 30.0
+_FALLBACK_IDENTITY_TIMEOUT = 300.0
+
+
+def _fallback_identity_timeout() -> float:
+    """Mirror the lower-level env contract while its module is still loading."""
+    raw = os.environ.get("TRTLLM_MPI_IDENTITY_TIMEOUT")
+    if raw:
+        try:
+            value = float(raw)
+            if math.isfinite(value) and value > 0:
+                return value
+        except ValueError:
+            pass
+    return _FALLBACK_IDENTITY_TIMEOUT
+
+
+def _shadow_build_wait_timeout() -> float:
+    """Upper-level wait budget derived from the MPI bootstrap deadline.
+
+    Do not import TensorRT-LLM here: this plugin must stay usable by pure-logic
+    tests and suites without built bindings. A real shadow build imports
+    ``mpi_session`` before it can construct a pool, so the live lower-level
+    setting is present by the time ``take()`` waits on that build.
+    """
+    mpi_session = sys.modules.get("tensorrt_llm.llmapi.mpi_session")
+    timeout_fn = getattr(mpi_session, "_identity_barrier_timeout", None)
+    identity_timeout = timeout_fn() if timeout_fn is not None else _fallback_identity_timeout()
+    return identity_timeout + _SHADOW_BUILD_FINISH_GRACE
+
 
 def _reuse_layer_active() -> bool:
     """True when the MPI session-reuse layer owns the pool-creation seams.
@@ -245,8 +279,10 @@ class _Built(NamedTuple):
 class SessionPrefetcher:
     def __init__(self):
         self._lock = threading.Lock()
+        self._drain_lock = threading.Lock()
         self._thread = None
         self._building_spec = None  # spec of the in-flight build, while _thread is set
+        self._build_timed_out = False
         self._build_gen = 0  # bumped when a pending build is abandoned
         self._built = None  # Optional[_Built], set only by _publish()
         self._patched = set()
@@ -332,8 +368,10 @@ def schedule_shadow(self, spec: int, env_overlay=None) -> None:
         """Start building a spare ``spec``-worker pool in the background.
 
         Heuristic: the next test most likely needs a pool of the same size as
-        the current one. A miss is discarded at ``take()`` and the sync build
-        is no slower than without prefetch.
+        the current one. A mismatched in-flight build is drained before a
+        synchronous miss to preserve allocation-wide single-flight. This can
+        add latency when the size prediction is wrong, but avoids two MPI
+        bootstraps contending on the same allocation.
 
         ``env_overlay``: extra env vars to freeze into the WORKERS at spawn
         (session_reuse restocks shadows with its worker-side weight cache
@@ -347,6 +385,7 @@ def schedule_shadow(self, spec: int, env_overlay=None) -> None:
             if self._thread is not None or (self._built is not None and self._built.spec == spec):
                 return  # already building / built
             self._building_spec = spec
+            self._build_timed_out = False
             self._thread = threading.Thread(
                 target=self._build,
                 args=(spec, self._build_gen, env_overlay),
@@ -398,44 +437,64 @@ def _publish(self, spec, session, snapshot, gen: int) -> None:
         print("[session-prefetch] discarding superseded background build", flush=True)
         session.shutdown()
 
-    def _drain(self, timeout: float):
-        """Join a pending build (abandoning it on timeout) and pop the slot."""
-        # Read _thread under the lock: schedule_shadow() assigns-then-starts
-        # inside its critical section, and an unlocked read here can observe
-        # the assigned-but-not-yet-started thread ("cannot join thread before
-        # it is started" when a test creates LLMs concurrently).
-        with self._lock:
-            thread = self._thread
-        if thread is not None:
-            thread.join(timeout=timeout)
-        with self._lock:
-            if thread is not None and thread.is_alive():
-                # Abandon the overdue build: bump the generation so its late
-                # _publish() shuts the pool down instead of landing.
-                self._build_gen += 1
-            self._thread = None
-            built, self._built = self._built, None
-        return built
-
-    def take(self, spec: int):
+    def _drain(self, timeout: float | None = None) -> _Built | None:
+        """Join a pending build and pop the completed shadow slot.
+
+        A live build is kept registered and marked terminal on timeout.
+        Callers fail closed instead of starting a second pool while the first
+        bootstrap is still running; later calls fail immediately until the
+        thread exits, then clear the terminal state.
+        """
+        # Serialize drains so concurrent LLM construction cannot make multiple
+        # callers wait through the full deadline before one records the
+        # terminal timeout.
+        with self._drain_lock:
+            # Read _thread under the lock: schedule_shadow() assigns-then-starts
+            # inside its critical section, and an unlocked read here can observe
+            # the assigned-but-not-yet-started thread ("cannot join thread before
+            # it is started" when a test creates LLMs concurrently).
+            with self._lock:
+                thread = self._thread
+                build_timed_out = self._build_timed_out
+            if thread is not None:
+                if build_timed_out and thread.is_alive():
+                    raise TimeoutError(
+                        "session-prefetch shadow build previously timed out and "
+                        "is still running; refusing to start a concurrent MPI pool"
+                    )
+                if not build_timed_out:
+                    if timeout is None:
+                        timeout = _shadow_build_wait_timeout()
+                    thread.join(timeout=timeout)
+            with self._lock:
+                if thread is not None and thread.is_alive():
+                    # Invalidate a late publish but retain _thread so
+                    # schedule_shadow() cannot start another build alongside it.
+                    if not self._build_timed_out:
+                        self._build_timed_out = True
+                        self._build_gen += 1
+                        self.stats["pool_build_timeouts"] += 1
+                    raise TimeoutError(
+                        "session-prefetch shadow build did not finish within "
+                        f"{timeout}s; refusing to start a concurrent MPI pool"
+                    )
+                self._thread = None
+                self._building_spec = None
+                self._build_timed_out = False
+                built, self._built = self._built, None
+            return built
+
+    def take(self, spec: int) -> object | None:
         """Return a prefetched session for ``spec``, or None to build sync."""
         if not self.enabled:
             return None
-        with self._lock:
-            wrong_size_in_flight = (
-                self._thread is not None and self._built is None and self._building_spec != spec
-            )
-        if wrong_size_in_flight:
-            # Joining would stall this caller for most of a spawn only to
-            # discard the mismatched result — slower than no prefetch at all.
-            # Fall back to the synchronous spawn now and leave the build to
-            # land for a later take of its own size.
-            self.stats["pools_skipped_size_in_flight"] += 1
-            return None
-        # Slowest legitimate build measured is ~117s (busy node); 180s gives
-        # 1.5x margin. On a genuine hang we give up and fall back to a
-        # synchronous build instead of stalling the suite.
-        built = self._drain(timeout=180)
+        # The upper-level deadline is derived from the identity barrier's
+        # bootstrap deadline plus a small post-bootstrap diagnostic grace.
+        # It must never expire while the lower layer still considers the
+        # in-flight pool healthy. A wrong-size build is also drained before
+        # returning a miss: starting the requested size alongside it would
+        # recreate the same concurrent-bootstrap contention.
+        built = self._drain()
         if built is None:
             return None
         if built.spec == spec and built.snapshot == _spawn_snapshot():
@@ -467,27 +526,12 @@ def factory(n_workers, *args, **kwargs):
             # workers exited (and released GPU memory) — the NEXT pool is
             # handed over instantly, without the ~50s sync spawn that used to
             # hide the release window. Such spawns fail closed when identity
-            # collection cannot complete; a prefetch layer must not turn that
-            # into a test failure, so retry once and then degrade LOUDLY to a
-            # plain pool (pre-prefetch semantics: shutdown returns at
-            # disconnect).
+            # collection cannot complete. Do not immediately retry or degrade
+            # to a plain pool: unidentified workers may still be exiting, and
+            # wait_shutdown=False cannot protect the next handover.
             session = self.take(n_workers)
             if session is None:
-                try:
-                    session = real_cls(n_workers=n_workers, wait_shutdown=True)
-                except Exception as e:
-                    print(f"[session-prefetch] pool spawn failed, retrying once: {e}", flush=True)
-                    try:
-                        session = real_cls(n_workers=n_workers, wait_shutdown=True)
-                    except Exception as e2:
-                        print(
-                            "[session-prefetch] wait_shutdown spawn failed twice: "
-                            f"{e2}; falling back to a plain pool (handover "
-                            "protection degraded for the NEXT pool on this node)",
-                            flush=True,
-                        )
-                        self.stats["pools_spawned_degraded"] += 1
-                        session = real_cls(n_workers=n_workers)
+                session = real_cls(n_workers=n_workers, wait_shutdown=True)
             self.schedule_shadow(n_workers)  # re-arm for the NEXT test
             return session
 
@@ -531,16 +575,24 @@ def install_pool_factory_if_loaded(self) -> None:
     def dispose(self) -> None:
         """Shut down any unconsumed shadow pool (end-of-session cleanup).
 
-        60s (vs take()'s 180s): at session end there is no test left to hand
-        the pool to, so a still-running build is only worth a short grace
-        before it is abandoned to its generation-bump cleanup. Idempotent: a
-        repository-root run dispatches sessionfinish from both the repo-root
-        and the subtree conftest.
+        Uses the same coordinated deadline as ``take()``. Ending the pytest
+        session is not permission to abandon a bootstrap while the lower
+        layer still considers it healthy. Idempotent: a repository-root run
+        dispatches sessionfinish from both the repo-root and the subtree
+        conftest.
         """
         if self._disposed:
             return
         self._disposed = True
-        built = self._drain(timeout=60)
+        try:
+            built = self._drain()
+        except TimeoutError as e:
+            # Do not turn pytest_sessionfinish into an internal error. The
+            # build remains registered, preventing another shadow from being
+            # launched; the daemon thread is bounded by the lower-level
+            # identity timeout unless the MPI runtime itself is wedged.
+            print(f"[session-prefetch] cleanup timed out: {e}", flush=True)
+            built = None
         if built is not None:
             built.session.shutdown()
         # One line per session, emitted OUTSIDE pytest's per-test capture

tests/test_common/session_reuse.py

@@ -364,27 +364,26 @@ def _spawn_fresh(self, real_cls, n_workers):
         overrides = {k: v for k, v in _WEIGHT_CACHE_ENV.items() if k not in os.environ}
         real = None
         prefetcher = _prefetcher()
+        if prefetcher is not None:
+            # A timeout must fail closed: starting a synchronous replacement
+            # would create two MPI pools concurrently on the same allocation.
+            # Unexpected prefetcher errors also propagate instead of silently
+            # hiding lifecycle bugs behind a synchronous fallback.
+            real = prefetcher.take(n_workers)
+        if real is None:
+            # One attempt gets the full worker-bootstrap deadline. If identity
+            # 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)
         if prefetcher is not None:
             try:
-                real = prefetcher.take(n_workers)
-            except Exception as e:  # prefetch is an optimization: fall back
-                print(f"[session-reuse] prefetched-pool take failed: {e}", flush=True)
-                real = None
-            try:
-                # Restock ONE shadow for the next miss of this size (no-op if
-                # one is already armed/building, or prefetch is disabled).
+                # Restock only after the current pool is ready. On a shadow
+                # miss, scheduling before the synchronous spawn would make
+                # two MPI pools bootstrap concurrently on the same GPUs.
                 prefetcher.schedule_shadow(n_workers, env_overlay=overrides)
             except Exception:
                 pass
-        if real is None:
-            try:
-                real = real_cls(n_workers=n_workers, wait_shutdown=True, env_overrides=overrides)
-            except Exception as e:
-                # wait_shutdown spawns fail closed (identity collection must
-                # complete); a transient slow node deserves one loud retry —
-                # a second failure means the node is genuinely broken.
-                print(f"[session-reuse] pool spawn failed, retrying once: {e}", flush=True)
-                real = real_cls(n_workers=n_workers, wait_shutdown=True, env_overrides=overrides)
         real._reuse_uses = 0
         real._reuse_spawn_snapshot = snapshot
         # (pid, start_time) per worker, recorded by the library at spawn