diff --git a/tensorrt_llm/llmapi/mpi_session.py b/tensorrt_llm/llmapi/mpi_session.py index 93c863ff7b29..31bc7363e624 100644 --- a/tensorrt_llm/llmapi/mpi_session.py +++ b/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: diff --git a/tests/test_common/session_prefetcher.py b/tests/test_common/session_prefetcher.py index 2325df25b26b..188a6aa70cc3 100644 --- a/tests/test_common/session_prefetcher.py +++ b/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 diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py index ebe9eac2ff4d..ea1b6272a5c5 100644 --- a/tests/test_common/session_reuse.py +++ b/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 diff --git a/tests/unittest/llmapi/test_mpi_session.py b/tests/unittest/llmapi/test_mpi_session.py index 98a92e2712b9..f0db0d2ba99e 100644 --- a/tests/unittest/llmapi/test_mpi_session.py +++ b/tests/unittest/llmapi/test_mpi_session.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import os import subprocess # nosec B404 import sys @@ -10,8 +13,10 @@ import pytest from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE -from tensorrt_llm.llmapi.mpi_session import (MPINodeState, MpiPoolSession, +from tensorrt_llm.llmapi.mpi_session import (_DEFAULT_IDENTITY_TIMEOUT, + MPINodeState, MpiPoolSession, RemoteMpiCommSessionClient, + _identity_barrier_timeout, split_mpi_env) # isort: off @@ -216,7 +221,11 @@ def test_wait_workers_exit_bounded_by_timeout_on_live_worker(): assert 0.2 <= waited < 2.0 # bounded: a wedged worker cannot hang teardown -def _collect_identities(monkeypatch, results, pending=0, n_workers=2): +def _collect_identities(monkeypatch, + results, + pending=0, + n_workers=2, + observed_timeouts=None): """Drive _collect_worker_identities on an inert stand-in (no MPI spawn).""" import types from concurrent.futures import Future @@ -230,7 +239,12 @@ def _collect_identities(monkeypatch, results, pending=0, n_workers=2): from tensorrt_llm.llmapi import mpi_session as m - monkeypatch.setattr(m, "futures_wait", lambda fs, timeout: (futs, never)) + def _fake_wait(fs, timeout): + if observed_timeouts is not None: + observed_timeouts.append(timeout) + return futs, never + + monkeypatch.setattr(m, "futures_wait", _fake_wait) killed = [] monkeypatch.setattr(os, "kill", lambda pid, sig: killed.append(pid)) it = iter(futs + never) @@ -275,3 +289,45 @@ def test_identity_collection_fails_closed_on_duplicate_pids(monkeypatch): me = (os.getpid(), _process_start_time(os.getpid())) with _pytest.raises(RuntimeError, match="incomplete"): _collect_identities(monkeypatch, [me, me]) # one worker answered twice + + +def test_identity_collection_uses_configured_timeout(monkeypatch): + from tensorrt_llm.llmapi.mpi_session import _process_start_time + + monkeypatch.setenv("TRTLLM_MPI_IDENTITY_TIMEOUT", "123.5") + observed_timeouts = [] + me = (os.getpid(), _process_start_time(os.getpid())) + _collect_identities(monkeypatch, [me], + n_workers=1, + observed_timeouts=observed_timeouts) + assert observed_timeouts == [123.5] + + +def test_identity_timeout_covers_worker_bootstrap(monkeypatch): + # The deadline bounds spawn + `import tensorrt_llm`, not barrier latency: + # it must exceed the slowest bootstrap the repo measures (~117s busy node). + monkeypatch.delenv("TRTLLM_MPI_IDENTITY_TIMEOUT", raising=False) + assert _identity_barrier_timeout() > 117.0 + + +# Invalid values (unparsable, non-positive) fall back to the default rather +# than turning the barrier into a busy-wait or an unbounded block. +@pytest.mark.parametrize("raw, expected", + [("90", 90.0), ("0.5", 0.5), + ("", _DEFAULT_IDENTITY_TIMEOUT), + ("0", _DEFAULT_IDENTITY_TIMEOUT), + ("-1", _DEFAULT_IDENTITY_TIMEOUT), + ("abc", _DEFAULT_IDENTITY_TIMEOUT), + ("nan", _DEFAULT_IDENTITY_TIMEOUT), + ("inf", _DEFAULT_IDENTITY_TIMEOUT), + ("-inf", _DEFAULT_IDENTITY_TIMEOUT), + ("1e309", _DEFAULT_IDENTITY_TIMEOUT)]) +def test_identity_timeout_env_override(monkeypatch, raw, expected): + monkeypatch.setenv("TRTLLM_MPI_IDENTITY_TIMEOUT", raw) + assert _identity_barrier_timeout() == expected + + +def test_prefetch_fallback_identity_timeout_matches_mpi_default(): + from test_common.session_prefetcher import _FALLBACK_IDENTITY_TIMEOUT + + assert _FALLBACK_IDENTITY_TIMEOUT == _DEFAULT_IDENTITY_TIMEOUT diff --git a/tests/unittest/llmapi/test_session_prefetcher.py b/tests/unittest/llmapi/test_session_prefetcher.py index f88735d3a020..b91869f24aab 100644 --- a/tests/unittest/llmapi/test_session_prefetcher.py +++ b/tests/unittest/llmapi/test_session_prefetcher.py @@ -160,47 +160,137 @@ def test_schedule_shadow_passes_env_overlay_to_build(prefetcher): assert prefetcher.overlays == [{"TRTLLM_HF_WEIGHT_CACHE": "1"}] -def test_take_wrong_size_in_flight_does_not_wait(prefetcher, monkeypatch): - # A miss must not stall behind an in-flight build of ANOTHER size only to - # discard the result — that is slower than no prefetch at all (wait ~one - # spawn, then spawn again). It falls back to sync immediately, and the - # build still lands for a later take of its own size. +def test_take_wrong_size_in_flight_waits_before_miss(prefetcher, monkeypatch): + # A wrong-size shadow is a miss, but the caller must not start its sync + # pool until that in-flight bootstrap finishes. Otherwise a group-size + # transition can make two MPI pools compete on the same allocation. release = threading.Event() def _slow_build(self, spec, gen, env_overlay=None): - release.wait(5) + release.wait() self._publish(spec, _FakePool(spec), session_prefetcher._spawn_snapshot(), gen) monkeypatch.setattr(SessionPrefetcher, "_build", _slow_build) prefetcher.schedule_shadow(2) - t0 = time.monotonic() - assert prefetcher.take(4) is None # wrong size in flight: no join - assert time.monotonic() - t0 < 1.0 - assert prefetcher.stats["pools_skipped_size_in_flight"] == 1 + result = [] + taker = threading.Thread(target=lambda: result.append(prefetcher.take(4))) + taker.start() + time.sleep(0.1) + assert taker.is_alive() # blocked behind the only in-flight bootstrap release.set() - prefetcher._thread.join(timeout=10) - taken = prefetcher.take(2) # the undisturbed build landed for its size - assert isinstance(taken, _FakePool) and taken.n_workers == 2 + taker.join(timeout=10) + assert result == [None] # wrong-size pool was drained, then rejected + + +def test_shadow_wait_budget_tracks_identity_timeout(monkeypatch): + fake_mpi_session = types.SimpleNamespace(_identity_barrier_timeout=lambda: 125.0) + monkeypatch.setitem(sys.modules, "tensorrt_llm.llmapi.mpi_session", fake_mpi_session) + assert ( + session_prefetcher._shadow_build_wait_timeout() + == 125 + session_prefetcher._SHADOW_BUILD_FINISH_GRACE + ) + + +def test_shadow_wait_budget_handles_partially_loaded_mpi_module(monkeypatch): + monkeypatch.setitem(sys.modules, "tensorrt_llm.llmapi.mpi_session", types.SimpleNamespace()) + monkeypatch.setenv("TRTLLM_MPI_IDENTITY_TIMEOUT", "625") + assert ( + session_prefetcher._shadow_build_wait_timeout() + == 625 + session_prefetcher._SHADOW_BUILD_FINISH_GRACE + ) + + +def test_take_timeout_is_terminal_until_build_exits(prefetcher, monkeypatch): + join_timeouts = [] + + class _HungBuild: + alive = True + + def join(self, timeout): + join_timeouts.append(timeout) + + def is_alive(self): + return self.alive + thread = _HungBuild() + prefetcher._thread = thread + prefetcher._building_spec = 2 + monkeypatch.setattr(session_prefetcher, "_shadow_build_wait_timeout", lambda: 321.0) + + with pytest.raises(TimeoutError, match="refusing to start a concurrent MPI pool"): + prefetcher.take(2) + with pytest.raises(TimeoutError, match="previously timed out"): + prefetcher.take(2) + + # Only the first call spends the full wait budget and records a timeout. + # Later calls fail fast until the abandoned build actually exits. + assert join_timeouts == [321.0] + assert prefetcher._thread is thread + assert prefetcher._build_gen == 1 + assert prefetcher.stats["pool_build_timeouts"] == 1 + + thread.alive = False + assert prefetcher.take(2) is None + assert prefetcher._thread is None + assert not prefetcher._build_timed_out + + +def test_concurrent_take_timeout_waits_only_once(prefetcher, monkeypatch): + join_started = threading.Event() + release_join = threading.Event() + join_timeouts = [] + errors = [] + + class _HungBuild: + def join(self, timeout): + join_timeouts.append(timeout) + join_started.set() + release_join.wait() + + def is_alive(self): + return True + + prefetcher._thread = _HungBuild() + prefetcher._building_spec = 2 + monkeypatch.setattr(session_prefetcher, "_shadow_build_wait_timeout", lambda: 321.0) + + def _take(): + try: + prefetcher.take(2) + except TimeoutError as e: + errors.append(e) -def test_factory_degrades_loudly_when_wait_shutdown_spawn_fails(prefetcher): - # The library fails closed when identity collection cannot complete; the - # prefetch layer must not turn that into a test failure: retry once, - # then degrade LOUDLY to a plain pool (pre-prefetch semantics). + first = threading.Thread(target=_take) + second = threading.Thread(target=_take) + first.start() + assert join_started.wait(timeout=5) + second.start() + release_join.set() + first.join(timeout=5) + second.join(timeout=5) + + assert len(errors) == 2 + assert join_timeouts == [321.0] + assert prefetcher._build_gen == 1 + assert prefetcher.stats["pool_build_timeouts"] == 1 + + +def test_factory_spawn_failure_propagates_without_retry(prefetcher): + # Once the full identity deadline expires, an unidentified worker may + # still be alive. Retrying or downgrading to wait_shutdown=False could + # overlap it with another pool, so the failure must propagate. calls = [] class _FailingWaitPool(_FakePool): def __init__(self, n_workers, wait_shutdown=False): calls.append(wait_shutdown) - if wait_shutdown: - raise RuntimeError("identity collection incomplete") - super().__init__(n_workers, wait_shutdown) + raise RuntimeError("identity collection incomplete") factory = prefetcher._make_factory(_FailingWaitPool) - session = factory(2) - assert isinstance(session, _FailingWaitPool) and not session.wait_shutdown - assert calls == [True, True, False] # two contract attempts, then plain - assert prefetcher.stats["pools_spawned_degraded"] == 1 + with pytest.raises(RuntimeError, match="identity collection incomplete"): + factory(2) + assert calls == [True] + assert prefetcher._thread is None # failure did not restock a shadow def test_factory_hit_hands_over_shadow(prefetcher): diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py index a2ea568d5186..ef693c291c10 100644 --- a/tests/unittest/llmapi/test_session_reuse.py +++ b/tests/unittest/llmapi/test_session_reuse.py @@ -71,6 +71,7 @@ def _wait(pred, timeout=5.0): def test_reuse_hands_back_same_pool(reuse_cache): s1 = reuse_cache.acquire(_FakePool, 2) real = s1._real + assert len(reuse_cache.prefetch.restocks) == 1 # armed only on the cache miss # Cache-managed pools block their (real) shutdown on worker exit, so a # replacement spawned right after a retire cannot race the GPU release. assert real.wait_shutdown @@ -79,6 +80,7 @@ def test_reuse_hands_back_same_pool(reuse_cache): s2 = reuse_cache.acquire(_FakePool, 2) assert s2._real is real # the SAME pool, reused assert reuse_cache.resets # workers were reset between handouts + assert len(reuse_cache.prefetch.restocks) == 1 # reuse does not create a shadow def test_cached_handover_reaps_in_flight_retires(reuse_cache): @@ -113,27 +115,79 @@ def test_cache_miss_takes_prefetched_shadow(reuse_cache): def test_cache_miss_falls_back_to_sync_spawn_and_restocks(reuse_cache): - s = reuse_cache.acquire(_FakePool, 2) # nothing armed: sync spawn + events = [] + + class _OrderedPool(_FakePool): + def __init__(self, *args, **kwargs): + events.append("sync pool ready") + super().__init__(*args, **kwargs) + + original_schedule = reuse_cache.prefetch.schedule_shadow + + def _record_schedule(*args, **kwargs): + events.append("shadow scheduled") + original_schedule(*args, **kwargs) + + reuse_cache.prefetch.schedule_shadow = _record_schedule + s = reuse_cache.acquire(_OrderedPool, 2) # nothing armed: sync spawn assert isinstance(s._real, _FakePool) and s._real.wait_shutdown assert s._real.env_overrides.get("TRTLLM_HF_WEIGHT_CACHE") == "1" assert reuse_cache.prefetch.restocks # shadow armed for the NEXT miss + assert events == ["sync pool ready", "shadow scheduled"] + + +def test_shadow_timeout_does_not_start_sync_pool(reuse_cache): + calls = [] + + def _timeout(_n_workers): + raise TimeoutError("shadow bootstrap still running") + + reuse_cache.prefetch.take = _timeout + + class _RecordingPool(_FakePool): + def __init__(self, *args, **kwargs): + calls.append((args, kwargs)) + super().__init__(*args, **kwargs) + + with pytest.raises(TimeoutError, match="still running"): + reuse_cache.acquire(_RecordingPool, 2) + assert calls == [] + assert reuse_cache.prefetch.restocks == [] + + +def test_prefetcher_programming_error_propagates(reuse_cache): + calls = [] + + def _fail(_n_workers): + raise ValueError("prefetch lifecycle bug") + + reuse_cache.prefetch.take = _fail + + class _RecordingPool(_FakePool): + def __init__(self, *args, **kwargs): + calls.append((args, kwargs)) + super().__init__(*args, **kwargs) + + with pytest.raises(ValueError, match="prefetch lifecycle bug"): + reuse_cache.acquire(_RecordingPool, 2) + assert calls == [] + assert reuse_cache.prefetch.restocks == [] -def test_spawn_failure_retries_once(reuse_cache): - # wait_shutdown spawns fail closed (identity collection must complete); - # one loud retry absorbs a transient slow node, a second failure means - # the node is genuinely broken and must propagate. +def test_spawn_failure_propagates_without_retry(reuse_cache): + # The attempt already waited for the full bootstrap deadline. Retrying + # could overlap workers that the failed identity barrier could not reap. calls = [] - class _FlakyPool(_FakePool): + class _FailingPool(_FakePool): def __init__(self, n_workers, wait_shutdown=False, env_overrides=None): calls.append(1) - if len(calls) == 1: - raise RuntimeError("identity collection incomplete") - super().__init__(n_workers, wait_shutdown, env_overrides) + raise RuntimeError("identity collection incomplete") - s = reuse_cache.acquire(_FlakyPool, 2) - assert isinstance(s._real, _FlakyPool) and len(calls) == 2 + with pytest.raises(RuntimeError, match="identity collection incomplete"): + reuse_cache.acquire(_FailingPool, 2) + assert calls == [1] + assert reuse_cache.prefetch.restocks == [] def test_reuse_size_mismatch_builds_new(reuse_cache):