From 5fb2ab4f1b62db879874b5f8780432cc9c13bc5e Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 7 Jul 2026 07:16:24 -0700 Subject: [PATCH 01/12] [None][test] Automatic MPI session reuse via shared pytest infrastructure Eligible tests automatically reuse a shared MpiPoolSession with zero test changes - no signature edits, no wrapper functions. A pytest plugin (loaded via -p in both test trees' pytest.ini, plus a tests/-level fallback conftest) lazily patches the pool-construction seams; released pools return to a per-size cache and serve the next bare LLM(...) of the same size. RPC executors drain the cache and keep private pools at their own seam; env/sys.path spawn snapshots, a use-count cap, a rank-verified dynamo reset, an NVML settle barrier and a per-hit health probe guard every handover; every failure path degrades to a fresh synchronous build. TRTLLM_TEST_REUSE_SESSION=0 disables everything; disabled under xdist. Signed-off-by: qgai --- tests/conftest.py | 22 ++ tests/integration/defs/pytest.ini | 7 +- tests/test_common/grouped_test_utils.py | 54 +++ tests/test_common/session_reuse.py | 388 ++++++++++++++++++++ tests/test_common/session_reuse_hooks.py | 36 ++ tests/unittest/llmapi/test_session_reuse.py | 178 +++++++++ tests/unittest/pytest.ini | 9 +- 7 files changed, 689 insertions(+), 5 deletions(-) create mode 100644 tests/conftest.py create mode 100644 tests/test_common/grouped_test_utils.py create mode 100644 tests/test_common/session_reuse.py create mode 100644 tests/test_common/session_reuse_hooks.py create mode 100644 tests/unittest/llmapi/test_session_reuse.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000000..b0d2feebd7fc --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fallback wiring of the session-reuse plugin for test dirs WITHOUT an ini. + +tests/unittest and tests/integration/defs load the plugin through the ``-p`` +option in their own pytest.ini; their rootdir sits below this file, so this +conftest is never collected there (no double registration). Any other +directory under tests/ (current or future) resolves its rootdir at the repo +root and picks the hooks up from here, so automatic MPI session reuse covers +every test under tests/. +""" + +import os +import sys + +sys.path.insert(0, os.path.dirname(__file__)) # make test_common importable + +from test_common.session_reuse_hooks import ( # noqa: E402,F401 + pytest_configure, + pytest_runtest_setup, + pytest_sessionfinish, +) diff --git a/tests/integration/defs/pytest.ini b/tests/integration/defs/pytest.ini index becab6eae09e..ee5e97c00f74 100644 --- a/tests/integration/defs/pytest.ini +++ b/tests/integration/defs/pytest.ini @@ -1,9 +1,12 @@ [pytest] asyncio_default_fixture_loop_scope = module threadleak = True -threadleak_exclude = asyncio_\d+ +# Thread-\d+ \(_manager_spawn\) / session-reuse-* belong to a pool cached for +# reuse by the NEXT test (tests/test_common/session_reuse.py) and legitimately +# outlive the test they start under. +threadleak_exclude = asyncio_\d+|Thread-\d+ \(_manager_spawn\)|session-reuse-\w+ junit_family=legacy -addopts = --ignore-glob="*perf/test_perf.py" --ignore-glob="*perf/disagg/*" --ignore-glob="*test_list_validation.py" --ignore-glob="*llm-test-workspace*" --durations=0 -W ignore::DeprecationWarning --unused-fixtures +addopts = --ignore-glob="*perf/test_perf.py" --ignore-glob="*perf/disagg/*" --ignore-glob="*test_list_validation.py" --ignore-glob="*llm-test-workspace*" --durations=0 -W ignore::DeprecationWarning --unused-fixtures -p test_common.session_reuse_hooks pythonpath = ../../../examples/auto_deploy ../../ norecursedirs = ./triton/perf ./perf/disagg diff --git a/tests/test_common/grouped_test_utils.py b/tests/test_common/grouped_test_utils.py new file mode 100644 index 000000000000..2f74f0de95b2 --- /dev/null +++ b/tests/test_common/grouped_test_utils.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Worker-side helpers for tests that reuse one ``MpiPoolSession``. + +Used by the automatic session-reuse layer (``test_common/session_reuse.py``). +This is test infrastructure only and must not be used on production inference +paths. Heavy imports are done lazily so importing this module never triggers a +circular import during package initialization. +""" + + +def _run_and_get_worker_rank(fn) -> int: + fn() + from tensorrt_llm._utils import mpi_rank + + return mpi_rank() + + +def submit_sync_per_worker(mpi_session, fn, max_rounds: int = 8) -> None: + """Run ``fn`` at least once on EVERY worker of the pool. + + ``MpiPoolSession.submit_sync`` enqueues ``n_workers`` tasks, but + ``MPIPoolExecutor`` gives no one-task-per-worker guarantee: a fast task can + be drained twice by one idle worker, leaving another untouched. Verify + coverage by collecting the worker ranks that actually ran ``fn`` and + resubmitting until all workers are covered. + """ + expected = mpi_session.n_workers + seen: set = set() + for _ in range(max_rounds): + seen.update(mpi_session.submit_sync(_run_and_get_worker_rank, fn)) + if len(seen) >= expected: + return + raise RuntimeError( + f"task only reached worker ranks {sorted(seen)} after {max_rounds} " + f"rounds; expected {expected} distinct workers" + ) + + +def reset_worker_torch_compile_state() -> None: + """Reset per-worker torch.compile / Dynamo state (runs inside each worker). + + Dynamo's recompile counter is process-global and per-code-object. When + worker processes are reused across LLMs (shared ``MpiPoolSession``), each + ``torch_compile`` case recompiles the same ``model.forward`` code object + under new guards; the count accumulates and eventually trips + ``recompile_limit`` (16), which is a HARD failure under ``fullgraph=True`` + (``FailOnRecompileLimitHit``) and aborts the whole MPI job. Resetting + between cases makes each LLM start from a clean compile cache, like a fresh + process. Submit this to each worker via ``mpi_session.submit_sync(...)``. + """ + import torch + + torch._dynamo.reset() diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py new file mode 100644 index 000000000000..6a616715801d --- /dev/null +++ b/tests/test_common/session_reuse.py @@ -0,0 +1,388 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Automatic MPI session reuse for bare ``LLM(...)`` tests — zero test changes. + +Instead of destroying its ``MpiPoolSession`` at ``LLM`` shutdown, the pool is +returned to a per-size cache and handed to the NEXT bare ``LLM(...)`` of the +same size, saving the ~50-65s worker spawn+import per reuse. This delivers the +same reuse the explicit fixtures in this PR provide, but through shared test +infrastructure only: no test signature changes, no wrapper functions. + +Eligibility is automatic: +- size mismatch -> new pool (cache keeps the old one for later) +- env/sys.path mismatch -> cached pool retired (workers froze that state + at spawn; a stale pool would silently miss it) +- RPC executors -> keep a private, never-cached pool; their seam + drains the cache first (their engine build + cannot share GPUs with cached idle pools) +- tests passing their own ``_mpi_session`` (the explicit fixtures) -> never + reach the patched seam +- ``@pytest.mark.private_mpi_session`` -> explicit opt-out: the cache is + drained and the test gets an untracked fresh pool +- use-count cap -> pool retired after N handouts (default 16), + bounding worker state accumulation + +Between handouts every worker runs a torch.compile/Dynamo reset (rank-verified +via ``grouped_test_utils.submit_sync_per_worker``) and the handover waits for +the previous worker's GPU memory to actually be released (NVML settle barrier) +— both failure modes were observed in validation, not hypothetical. + +Enable/disable with ``TRTLLM_TEST_REUSE_SESSION`` (default on; ``0`` disables). +Disabled under pytest-xdist workers (parallel tests would multiply live pools). +""" + +import os +import sys +import threading +import time + +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. +_PATCH_TARGETS = ( + "tensorrt_llm.executor.proxy", + "tensorrt_llm.llmapi.llm", +) +# RPC executors keep a PRIVATE pool (never cached), but their engine build +# cannot share GPUs with cached idle pools (observed init hang), so their +# seam gets a drain-then-build factory instead of the cache. +_RPC_PATCH_TARGET = "tensorrt_llm.executor.rpc_proxy" +_ALL_PATCH_TARGETS = _PATCH_TARGETS + (_RPC_PATCH_TARGET,) + +# Worker-side HF weight cache for cache-managed pools only: set (if absent) +# around the spawn so the workers freeze it, then restored — private/RPC +# pools keep the production default (cache off). +_WEIGHT_CACHE_ENV = { + "TRTLLM_HF_WEIGHT_CACHE": "1", + "TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES": "1", +} + +# Workers freeze the parent environment AND sys.path at spawn time, so a +# cached pool must not be handed to a test that changed either (silently +# stale env / unimportable monkeypatched modules). Process bookkeeping that +# legitimately drifts between tests is ignored; a false mismatch only costs +# one synchronous rebuild. +_ENV_IGNORE = frozenset( + { + "PYTEST_CURRENT_TEST", + "COLUMNS", + "LINES", + "PWD", + "OLDPWD", + "SHLVL", + "_", + } +) + + +def _spawn_snapshot(): + """The worker-visible state a pool freezes at spawn: env + sys.path.""" + return ( + {k: v for k, v in os.environ.items() if k not in _ENV_IGNORE}, + list(sys.path), + ) + + +# GPU-memory settle barrier at handover: a reused live pool skips the ~50s +# synchronous spawn that used to give the previous LLM's worker time to exit; +# its CUDA memory is only released when the process actually exits. Building +# the next model into that race fails with "insufficient GPU memory". +_SETTLE_MIN_FREE_FRAC = 0.85 +_SETTLE_POLL_S = 0.5 +_SETTLE_FLAT_POLLS = 3 +_SETTLE_EPSILON = 256 << 20 +_SETTLE_TIMEOUT_S = 30.0 + + +def _visible_gpu_indices(count: int): + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if not visible: + return list(range(count)) + indices = [] + for token in visible.split(","): + token = token.strip() + if not token.isdigit() or int(token) >= count: + return list(range(count)) # UUID/MIG form: fall back to all GPUs + indices.append(int(token)) + return indices or list(range(count)) + + +def wait_gpu_memory_settle() -> None: + """Wait until visible GPUs are mostly free or free memory stops rising. + + Never raises: on any NVML problem the handover proceeds as before. + """ + try: + import pynvml + + pynvml.nvmlInit() + except Exception: + return + try: + handles = [ + pynvml.nvmlDeviceGetHandleByIndex(i) + for i in _visible_gpu_indices(pynvml.nvmlDeviceGetCount()) + ] + + def _free_total(): + infos = [pynvml.nvmlDeviceGetMemoryInfo(h) for h in handles] + return [i.free for i in infos], [i.total for i in infos] + + t0 = time.monotonic() + flat, prev = 0, None + while True: + free, total = _free_total() + if all(f >= _SETTLE_MIN_FREE_FRAC * t for f, t in zip(free, total)): + break + if prev is not None and all(f - p < _SETTLE_EPSILON for f, p in zip(free, prev)): + flat += 1 + if flat >= _SETTLE_FLAT_POLLS: + break # not increasing: that memory is legitimately in use + else: + flat = 0 + if time.monotonic() - t0 >= _SETTLE_TIMEOUT_S: + break + prev = free + time.sleep(_SETTLE_POLL_S) + waited = time.monotonic() - t0 + if waited >= _SETTLE_POLL_S: + print( + f"[session-reuse] waited {waited:.1f}s before handover for GPU memory release", + flush=True, + ) + except Exception: + pass + finally: + try: + pynvml.nvmlShutdown() + except Exception: + pass + + +def _describe_mismatch(spawn_snap, now_snap, uses, max_uses): + """One line naming WHY a cached pool cannot be handed out (observability).""" + if 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)] + if changed: + return f"env changed since spawn: {sorted(changed)[:6]}" + if spawn_path != now_path: + added = [p for p in now_path if p not in spawn_path] + removed = [p for p in spawn_path if p not in now_path] + return f"sys.path changed since spawn: +{added[:3]} -{removed[:3]}" + return "snapshot mismatch" + + +class _ReusableSession: + """A pool wrapper whose ``shutdown()`` returns it to the cache. + + Everything else delegates to the real ``MpiPoolSession``. ``shutdown_abort`` + marks the pool dead so a crashed pool is never handed out again. + """ + + def __init__(self, real, cache): + self._real = real + self._cache = cache + self._dead = False + self._released = False + + def __getattr__(self, name): + # Only fires for names NOT set on the wrapper (plain attribute reads + # of _real/_cache/... resolve normally, no recursion hazard). Reads + # after release stay delegated (harmless); destructive calls are + # gated in shutdown()/shutdown_abort() below. + return getattr(self.__dict__["_real"], name) + + def shutdown(self): + if self._dead or self._released: + return + self._released = True + self._cache._release(self._real) + + def shutdown_abort(self, *args, **kwargs): + if self._released: + # The pool went back to the cache at shutdown() and may already + # belong to the NEXT test: never kill it from a late error path. + return None + self._dead = True + self._cache._forget(self._real) + return self._real.shutdown_abort(*args, **kwargs) + + +class SessionReuseCache: + def __init__(self): + self._lock = threading.Lock() + self._pools = {} # n_workers -> MpiPoolSession + self._patched = set() + self._suspended = False + + @property + def enabled(self) -> bool: + if os.environ.get("PYTEST_XDIST_WORKER"): + return False # parallel workers would multiply live pools + return os.environ.get("TRTLLM_TEST_REUSE_SESSION", "1").lower() in ( + "1", + "true", + "yes", + "on", + ) + + @property + def max_uses(self) -> int: + return int(os.environ.get("TRTLLM_TEST_REUSE_MAX_USES", "16")) + + @staticmethod + def _retire(real): + threading.Thread(target=real.shutdown, daemon=True, name="session-reuse-retire").start() + + # ---- factory installed at the pool-creation seam ---- + + def install_pool_factory_if_loaded(self) -> None: + """Lazily patch the pool-creation seams (idempotent). + + Only patches target modules ALREADY imported by the test suite, so + suites that never create MPI pools pay nothing — not even the + tensorrt_llm import. Called from ``pytest_runtest_setup``. + """ + if len(self._patched) == len(_ALL_PATCH_TARGETS): + 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] + if not pending: + return + from tensorrt_llm.llmapi.mpi_session import MpiPoolSession as real_cls + + cache = self + + def factory(n_workers, *args, **kwargs): + if args or kwargs: # unknown calling convention: stay out of the way + print( + "[session-reuse] bypassing reuse: MpiPoolSession called with " + "unexpected arguments (library signature changed?)", + flush=True, + ) + return real_cls(n_workers, *args, **kwargs) + return cache.acquire(real_cls, n_workers) + + def rpc_factory(n_workers, *args, **kwargs): + # Fires exactly when an RPC executor is constructed, whatever the + # test is named — no name heuristics. + cache.drain() + return real_cls(n_workers, *args, **kwargs) + + for name in pending: + mod = sys.modules[name] + if getattr(mod, "MpiPoolSession", None) is real_cls: + mod.MpiPoolSession = rpc_factory if name == _RPC_PATCH_TARGET else factory + self._patched.add(name) + + # ---- cache operations ---- + + def acquire(self, real_cls, n_workers): + """Hand out a cached same-size pool (reset + settled) or build one.""" + if self._suspended or not self.enabled: + # Opt-out test (private_mpi_session) or the kill switch flipped + # after the seams were patched: untracked fresh pool that the LLM + # owns and destroys normally. + return real_cls(n_workers=n_workers) + with self._lock: + real = self._pools.pop(n_workers, None) + if real is not None: + # Compare against the state FROZEN INTO the workers at spawn time: + # if the current test expects different env/sys.path, the cached + # workers would silently miss it. + snap = _spawn_snapshot() + if real._reuse_spawn_snapshot != snap or real._reuse_uses >= self.max_uses: + print( + "[session-reuse] retiring cached pool: " + + _describe_mismatch( + real._reuse_spawn_snapshot, snap, real._reuse_uses, self.max_uses + ), + flush=True, + ) + self._retire(real) # stale worker state or lifetime cap + else: + try: + submit_sync_per_worker(real, reset_worker_torch_compile_state) + wait_gpu_memory_settle() + print( + f"[session-reuse] reusing {n_workers}-worker pool " + f"(use #{real._reuse_uses + 1})", + flush=True, + ) + return _ReusableSession(real, self) + except Exception as e: # unhealthy pool: discard, build fresh + print( + f"[session-reuse] cached pool failed reset, rebuilding: {e}", + flush=True, + ) + self._retire(real) + return _ReusableSession(self._spawn_fresh(real_cls, n_workers), self) + + def _spawn_fresh(self, real_cls, n_workers): + """Spawn a cache-managed pool with the worker-side HF weight cache on. + + The cache env vars must be visible at spawn (workers freeze the env) + and are removed right after, so non-managed pools (private/RPC) and + the rest of the suite keep the production default. The spawn snapshot + is taken BEFORE adding them so later acquire-time comparisons (which + see the restored env) still match. An explicit user setting of either + var is respected and left untouched. + """ + snapshot = _spawn_snapshot() + added = [k for k in _WEIGHT_CACHE_ENV if k not in os.environ] + for k in added: + os.environ[k] = _WEIGHT_CACHE_ENV[k] + try: + real = real_cls(n_workers=n_workers) + finally: + for k in added: + os.environ.pop(k, None) + real._reuse_uses = 0 + real._reuse_spawn_snapshot = snapshot + return real + + def _release(self, real): + real._reuse_uses += 1 + with self._lock: + prior = self._pools.get(real.n_workers) + if prior is not None and prior is not real: + self._retire(real) # a pool of this size is already cached + return + self._pools[real.n_workers] = real + + def _forget(self, real): + with self._lock: + if self._pools.get(real.n_workers) is real: + del self._pools[real.n_workers] + + def suspend(self, suspended: bool) -> None: + """Bypass the cache for the current test (private_mpi_session).""" + self._suspended = suspended + + def drain(self) -> None: + """Shut down all cached pools in parallel (frees GPU/CPU footprint).""" + with self._lock: + pools, self._pools = list(self._pools.values()), {} + if not pools: + return + threads = [threading.Thread(target=p.shutdown, name="session-reuse-drain") for p in pools] + for t in threads: + t.start() + 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) + + +REUSE = SessionReuseCache() diff --git a/tests/test_common/session_reuse_hooks.py b/tests/test_common/session_reuse_hooks.py new file mode 100644 index 000000000000..bb10b73a4740 --- /dev/null +++ b/tests/test_common/session_reuse_hooks.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pytest plugin wiring for automatic MPI session reuse (repo-wide). + +Loaded via ``-p test_common.session_reuse_hooks`` from each test tree's +pytest.ini. Factory installation is LAZY: nothing is patched until the test +suite itself imports tensorrt_llm's executor modules (which only happens for +tests that create MPI pools), so suites that never create pools pay nothing — +not even the tensorrt_llm import. +""" + +from test_common.session_reuse import REUSE + + +def pytest_configure(config): + config.addinivalue_line( + "markers", + "private_mpi_session: opt this test out of automatic MPI session " + "reuse; it gets a fresh private pool and the reuse cache is drained " + "first (use for Ray/custom executors or tests needing isolation; RPC " + "executors are handled automatically at their construction seam)", + ) + + +def pytest_runtest_setup(item): + if not REUSE.enabled: + return + opt_out = item.get_closest_marker("private_mpi_session") is not None + if opt_out: + REUSE.drain() + REUSE.suspend(opt_out) + REUSE.install_pool_factory_if_loaded() + + +def pytest_sessionfinish(session, exitstatus): + REUSE.drain() diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py new file mode 100644 index 000000000000..03514356ccd6 --- /dev/null +++ b/tests/unittest/llmapi/test_session_reuse.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pure-logic tests for automatic MPI session reuse — no MPI, no GPU.""" + +import pytest +from test_common import session_reuse +from test_common.session_reuse import SessionReuseCache + + +class _FakePool: + def __init__(self, n_workers): + self.n_workers = n_workers + self.shut = False + import os + + self.spawn_env_weight_cache = os.environ.get("TRTLLM_HF_WEIGHT_CACHE") + + def shutdown(self): + self.shut = True + + def shutdown_abort(self, *args, **kwargs): + self.shut = True + + +@pytest.fixture +def reuse_cache(monkeypatch): + monkeypatch.setenv("TRTLLM_TEST_REUSE_SESSION", "1") + monkeypatch.delenv("PYTEST_XDIST_WORKER", raising=False) + 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, "wait_gpu_memory_settle", lambda: None) + cache.resets = resets + return cache + + +def _wait(pred, timeout=5.0): + import time + + t0 = time.monotonic() + while time.monotonic() - t0 < timeout: + if pred(): + return True + time.sleep(0.05) + return False + + +def test_reuse_hands_back_same_pool(reuse_cache): + s1 = reuse_cache.acquire(_FakePool, 2) + real = s1._real + s1.shutdown() # released to cache, not killed + assert not real.shut + s2 = reuse_cache.acquire(_FakePool, 2) + assert s2._real is real # the SAME pool, reused + assert reuse_cache.resets # workers were reset between handouts + + +def test_reuse_size_mismatch_builds_new(reuse_cache): + s1 = reuse_cache.acquire(_FakePool, 2) + s1.shutdown() + s2 = reuse_cache.acquire(_FakePool, 4) + assert s2._real.n_workers == 4 and s2._real is not s1._real + + +def test_reuse_env_change_retires_pool(reuse_cache, monkeypatch): + s1 = reuse_cache.acquire(_FakePool, 2) + real = s1._real + s1.shutdown() + monkeypatch.setenv("SOME_TEST_KNOB", "changed-after-spawn") + s2 = reuse_cache.acquire(_FakePool, 2) + assert s2._real is not real # stale-env pool not handed out + assert _wait(lambda: real.shut) # and it was retired in the background + + +def test_reuse_syspath_change_retires_pool(reuse_cache, monkeypatch): + s1 = reuse_cache.acquire(_FakePool, 2) + real = s1._real + s1.shutdown() + monkeypatch.syspath_prepend("/oot/example/path") + s2 = reuse_cache.acquire(_FakePool, 2) + assert s2._real is not real # frozen workers could not import from it + + +def test_reuse_max_uses_retires_pool(reuse_cache, monkeypatch): + monkeypatch.setenv("TRTLLM_TEST_REUSE_MAX_USES", "2") + s = reuse_cache.acquire(_FakePool, 2) + first = s._real + s.shutdown() + s = reuse_cache.acquire(_FakePool, 2) # use #2 (cap) + assert s._real is first + s.shutdown() + s = reuse_cache.acquire(_FakePool, 2) # over the cap -> fresh pool + assert s._real is not first + assert _wait(lambda: first.shut) + + +def test_reuse_abort_never_recycles(reuse_cache): + s1 = reuse_cache.acquire(_FakePool, 2) + real = s1._real + s1.shutdown_abort() + s2 = reuse_cache.acquire(_FakePool, 2) + assert s2._real is not real + + +def test_reuse_failed_reset_rebuilds(reuse_cache, monkeypatch): + s1 = reuse_cache.acquire(_FakePool, 2) + real = s1._real + s1.shutdown() + + def _boom(session, fn): + raise RuntimeError("worker died") + + monkeypatch.setattr(session_reuse, "submit_sync_per_worker", _boom) + s2 = reuse_cache.acquire(_FakePool, 2) # health probe fails -> fresh pool + assert s2._real is not real + assert _wait(lambda: real.shut) + + +def test_suspended_gives_private_untracked_pool(reuse_cache): + reuse_cache.suspend(True) + s = reuse_cache.acquire(_FakePool, 2) + assert isinstance(s, _FakePool) # raw pool: the LLM owns and destroys it + reuse_cache.suspend(False) + + +def test_disabled_under_xdist(monkeypatch): + monkeypatch.setenv("TRTLLM_TEST_REUSE_SESSION", "1") + monkeypatch.setenv("PYTEST_XDIST_WORKER", "gw0") + assert not SessionReuseCache().enabled + + +def test_weight_cache_env_scoped_to_spawn(reuse_cache, monkeypatch): + import os + + monkeypatch.delenv("TRTLLM_HF_WEIGHT_CACHE", raising=False) + s = reuse_cache.acquire(_FakePool, 2) + # Workers froze the cache env at spawn... + assert s._real.spawn_env_weight_cache == "1" + # ...but the suite's environment is untouched afterwards. + assert "TRTLLM_HF_WEIGHT_CACHE" not in os.environ + + +def test_weight_cache_env_respects_user_setting(reuse_cache, monkeypatch): + monkeypatch.setenv("TRTLLM_HF_WEIGHT_CACHE", "0") + s = reuse_cache.acquire(_FakePool, 2) + assert s._real.spawn_env_weight_cache == "0" # explicit user value wins + + +def test_abort_after_release_never_kills_cached_pool(reuse_cache): + # Late executor error paths can call shutdown_abort AFTER shutdown already + # returned the pool to the cache; the pool may belong to the NEXT test by + # then and must not be killed through the stale wrapper. + s1 = reuse_cache.acquire(_FakePool, 2) + real = s1._real + s1.shutdown() # released to cache + s1.shutdown_abort() # late abort on the stale wrapper: must be a no-op + assert not real.shut + s2 = reuse_cache.acquire(_FakePool, 2) + assert s2._real is real # still reusable + + +def test_disable_after_patch_bypasses_cache(reuse_cache, monkeypatch): + # The kill switch must keep meaning "off" even after the seams were + # patched: acquire() consults it on every call. + s1 = reuse_cache.acquire(_FakePool, 2) + s1.shutdown() + monkeypatch.setenv("TRTLLM_TEST_REUSE_SESSION", "0") + s2 = reuse_cache.acquire(_FakePool, 2) + assert isinstance(s2, _FakePool) # raw private pool, cache untouched + + +def test_drain_shuts_cached_pools(reuse_cache): + s = reuse_cache.acquire(_FakePool, 2) + real = s._real + s.shutdown() + reuse_cache.drain() + assert real.shut diff --git a/tests/unittest/pytest.ini b/tests/unittest/pytest.ini index 62e38297239d..5b540eedc8c2 100644 --- a/tests/unittest/pytest.ini +++ b/tests/unittest/pytest.ini @@ -2,9 +2,12 @@ xdist_start_method = spawn asyncio_default_fixture_loop_scope = module threadleak = True -# ThreadPoolExecutor-\d+_\d+ excludes worker threads leaked by torch._dynamo/torch.compile -threadleak_exclude = asyncio_\d+|rpc_client_loop|rpc_client_worker_\d+|rpc_server_worker_\d+|InductorSubproc|subproc_worker_timer|ThreadPoolExecutor-\d+_\d+ -addopts = --durations=0 -W ignore::DeprecationWarning +# ThreadPoolExecutor-\d+_\d+ excludes worker threads leaked by torch._dynamo/torch.compile. +# Thread-\d+ \(_manager_spawn\) is the MPIPoolExecutor manager thread of a pool cached +# for reuse by the NEXT test, and session-reuse-* are the reuse layer's own threads +# (tests/test_common/session_reuse.py); they legitimately outlive the test they start under. +threadleak_exclude = asyncio_\d+|rpc_client_loop|rpc_client_worker_\d+|rpc_server_worker_\d+|InductorSubproc|subproc_worker_timer|ThreadPoolExecutor-\d+_\d+|Thread-\d+ \(_manager_spawn\)|session-reuse-\w+ +addopts = --durations=0 -W ignore::DeprecationWarning -p test_common.session_reuse_hooks pythonpath = auto_deploy/_utils_test ../../examples/auto_deploy From 3de898adb0e8935bcc0f806c434197b71930365f Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 7 Jul 2026 19:24:10 -0700 Subject: [PATCH 02/12] [None][test] Opt eagle3 rejection-sampling test out of MPI session reuse The XQA JIT cubin registry is process-global (DecoderXQARunner::getResourceGlobal) and its lookup key omits compile context fields q_seq_len and is_spec_dec_tree. Running the dynamic-tree and non-dynamic-tree variants of this test inside one reused MPI worker makes the second variant hit the cubin compiled for the first variant's q_seq_len (both fold into the same m_tilesize bucket), which fails at launch with CUDA_ERROR_INVALID_VALUE on Hopper. Reproduced deterministically on H100 with the exact CI wheel: the cross-config pair fails in both orders at pool use #2, while same-config pairs, reuse-disabled runs, and fresh-pool single runs all pass. On B200 the same handover passes because SM100 does not use this XQA JIT path. Opt the test out via the private_mpi_session marker until the registry key is fixed upstream to include the full JIT compile context. Signed-off-by: qgai --- tests/unittest/_torch/speculative/test_eagle3.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 4564ab550526..6298565205eb 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -939,6 +939,12 @@ def test_eagle3_cdl_sampling(disable_overlap_scheduler: bool): @pytest.mark.parametrize("use_cuda_graph", [False, True]) @pytest.mark.high_cuda_memory @skip_blackwell +# Opt out of MPI session reuse: the XQA JIT cubin registry is process-global +# (DecoderXQARunner::getResourceGlobal) and its lookup key does not include +# q_seq_len / is_spec_dec_tree, so running the dynamic-tree and non-dynamic-tree +# variants in one worker process launches a cubin compiled for the other +# config's q_seq_len -> CUDA_ERROR_INVALID_VALUE on Hopper. +@pytest.mark.private_mpi_session @with_mocked_hf_download_for_single_gpu def test_llama_eagle3_rejection_sampling_modes(use_dynamic_tree: bool, use_cuda_graph: bool): From 67d18aedfca0763699175b3c9364c047a682aa8a Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 8 Jul 2026 06:32:02 -0700 Subject: [PATCH 03/12] [None][test] Opt AutoDeploy registry accuracy test out of MPI session reuse On a reused MPI worker (pool use #2) this test deterministically fails with 'The expanded size of the tensor (128) must match the existing size (256)': AutoDeploy worker-process state sized for the previous test's config leaks into the next engine. The same case passes on every recent build without session reuse. Same failure class as the eagle3 XQA JIT opt-out; keep this test on a private session until the state source is root-caused. Signed-off-by: qgai --- tests/integration/defs/accuracy/test_llm_api_autodeploy.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 5235c9329f47..89ff004acccc 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -1563,6 +1563,12 @@ def get_default_sampling_params(self): use_beam_search=False) @pytest.mark.skip_less_device_memory(32000) + # Opt out of MPI session reuse: on a reused worker (pool use #2) this test + # fails with "expanded size of the tensor (128) must match the existing + # size (256)" -- AutoDeploy worker-process state sized for the previous + # test's config leaks into the next engine. Keep private until the state + # source is root-caused. + @pytest.mark.private_mpi_session @pytest.mark.parametrize("accuracy_check", [False, True]) @pytest.mark.parametrize("model_name,config_overrides,tasks", MODEL_REGISTRY_ACCURACY_PARAMS) From e943b4f39f71044158cdec42232091818ed8a365 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 8 Jul 2026 19:16:01 -0700 Subject: [PATCH 04/12] [None][test] Opt all AutoDeploy tests out of MPI session reuse Three AutoDeploy tests have now failed on reused workers with the same failure class (tensor expand-size mismatch such as 'expanded size of the tensor (128) must match the existing size (256)', at pool use #2): registry accuracy, guided decoding, and eagle3 one-model. AutoDeploy's executor adapter keeps worker-process state sized for the previous engine's config, so per-test markers are whack-a-mole. Replace the individual marker with a plugin-level node-id pattern opt-out (autodeploy / auto_deploy / test_ad_) in session_reuse_hooks and revert the now-redundant marker on the registry accuracy test. Remove the patterns once AutoDeploy re-initializes that state per engine. Signed-off-by: qgai --- .../defs/accuracy/test_llm_api_autodeploy.py | 6 ----- tests/test_common/session_reuse_hooks.py | 22 ++++++++++++++++++- tests/unittest/llmapi/test_session_reuse.py | 19 ++++++++++++++++ 3 files changed, 40 insertions(+), 7 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 89ff004acccc..5235c9329f47 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -1563,12 +1563,6 @@ def get_default_sampling_params(self): use_beam_search=False) @pytest.mark.skip_less_device_memory(32000) - # Opt out of MPI session reuse: on a reused worker (pool use #2) this test - # fails with "expanded size of the tensor (128) must match the existing - # size (256)" -- AutoDeploy worker-process state sized for the previous - # test's config leaks into the next engine. Keep private until the state - # source is root-caused. - @pytest.mark.private_mpi_session @pytest.mark.parametrize("accuracy_check", [False, True]) @pytest.mark.parametrize("model_name,config_overrides,tasks", MODEL_REGISTRY_ACCURACY_PARAMS) diff --git a/tests/test_common/session_reuse_hooks.py b/tests/test_common/session_reuse_hooks.py index bb10b73a4740..5982973478e2 100644 --- a/tests/test_common/session_reuse_hooks.py +++ b/tests/test_common/session_reuse_hooks.py @@ -11,6 +11,24 @@ from test_common.session_reuse import REUSE +# Node-id substrings that are opted out of session reuse as a class, in +# addition to the per-test ``private_mpi_session`` marker. AutoDeploy's +# executor adapter keeps worker-process state sized for the previous engine's +# config (graph/sequence-length shaped buffers); on a reused pool the next +# test hits e.g. "The expanded size of the tensor (128) must match the +# existing size (256)". Remove entries once the underlying state is +# re-initialized per engine. +_PRIVATE_NODEID_PATTERNS = ( + "autodeploy", + "auto_deploy", + "/test_ad_", +) + + +def _is_private_nodeid(nodeid: str) -> bool: + lowered = nodeid.lower() + return any(pat in lowered for pat in _PRIVATE_NODEID_PATTERNS) + def pytest_configure(config): config.addinivalue_line( @@ -25,7 +43,9 @@ def pytest_configure(config): def pytest_runtest_setup(item): if not REUSE.enabled: return - opt_out = item.get_closest_marker("private_mpi_session") is not None + opt_out = item.get_closest_marker("private_mpi_session") is not None or _is_private_nodeid( + item.nodeid + ) if opt_out: REUSE.drain() REUSE.suspend(opt_out) diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py index 03514356ccd6..2360f7b0ebcd 100644 --- a/tests/unittest/llmapi/test_session_reuse.py +++ b/tests/unittest/llmapi/test_session_reuse.py @@ -176,3 +176,22 @@ def test_drain_shuts_cached_pools(reuse_cache): s.shutdown() reuse_cache.drain() assert real.shut + + +def test_autodeploy_nodeids_are_private(): + from test_common.session_reuse_hooks import _is_private_nodeid + + assert _is_private_nodeid( + "accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::" + "test_autodeploy_from_registry[m-True]" + ) + 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 not _is_private_nodeid( + "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[a]" + ) + assert not _is_private_nodeid( + "unittest/_torch/speculative/test_eagle3.py::test_llama_eagle3[x]" + ) From a68b695611dd5bdfd460c3d77c855bbe62e761e4 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 8 Jul 2026 19:28:46 -0700 Subject: [PATCH 05/12] [None][test] Harden session reuse against wedged worker pools Two robustness gaps let a single broken pool hang the whole suite, observed when a test's executor dies mid-shutdown ('Failed to send object') and poisons the shared pool: - The health probe called submit_sync, whose future.result() has no timeout. Workers that are alive but wedged in a collective never complete the probe task, so the next acquire blocked forever. Bound each probe round with concurrent.futures.wait(round_timeout) and treat a timeout as probe failure (retire + fresh spawn). - _retire disposed pools with a graceful shutdown(), which blocks indefinitely on wedged workers and leaks their GPU memory into subsequent tests. Prefer shutdown_abort (bounded grace, then kill) with shutdown(wait=False) as fallback. Signed-off-by: qgai --- tests/test_common/grouped_test_utils.py | 21 +++++++++++++++++++-- tests/test_common/session_reuse.py | 21 ++++++++++++++++++++- 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/tests/test_common/grouped_test_utils.py b/tests/test_common/grouped_test_utils.py index 2f74f0de95b2..a5eb9f02fee3 100644 --- a/tests/test_common/grouped_test_utils.py +++ b/tests/test_common/grouped_test_utils.py @@ -16,7 +16,9 @@ def _run_and_get_worker_rank(fn) -> int: return mpi_rank() -def submit_sync_per_worker(mpi_session, fn, max_rounds: int = 8) -> None: +def submit_sync_per_worker( + mpi_session, fn, max_rounds: int = 8, round_timeout: float = 60.0 +) -> None: """Run ``fn`` at least once on EVERY worker of the pool. ``MpiPoolSession.submit_sync`` enqueues ``n_workers`` tasks, but @@ -24,11 +26,26 @@ def submit_sync_per_worker(mpi_session, fn, max_rounds: int = 8) -> None: be drained twice by one idle worker, leaving another untouched. Verify coverage by collecting the worker ranks that actually ran ``fn`` and resubmitting until all workers are covered. + + Each round is bounded by ``round_timeout``: a worker that is alive but + wedged (e.g. stuck in a collective after the previous test's executor died + mid-shutdown) never completes its future, so an unbounded ``result()`` + would hang the whole session. A timed-out round is reported as probe + failure, making the caller retire the pool and spawn a fresh one. """ + import concurrent.futures + expected = mpi_session.n_workers seen: set = set() for _ in range(max_rounds): - seen.update(mpi_session.submit_sync(_run_and_get_worker_rank, fn)) + futures = mpi_session.submit(_run_and_get_worker_rank, fn) + done, not_done = concurrent.futures.wait(futures, timeout=round_timeout) + if not_done: + raise RuntimeError( + f"health probe timed out after {round_timeout}s: " + f"{len(not_done)}/{len(futures)} worker tasks never returned" + ) + seen.update(f.result() for f in done) if len(seen) >= expected: return raise RuntimeError( diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py index 6a616715801d..df382f97ce9c 100644 --- a/tests/test_common/session_reuse.py +++ b/tests/test_common/session_reuse.py @@ -236,7 +236,26 @@ def max_uses(self) -> int: @staticmethod def _retire(real): - threading.Thread(target=real.shutdown, daemon=True, name="session-reuse-retire").start() + """Dispose of a pool in the background without blocking the test. + + Pools are usually retired because they are broken (failed health + probe, stale env, lifetime cap). A graceful ``shutdown()`` on a pool + whose workers are wedged in a collective blocks forever and leaks the + workers' GPU memory into subsequent tests, so prefer + ``shutdown_abort`` (bounded grace, then kill) and fall back to + ``shutdown`` only if abort is unavailable or raises. + """ + + def _dispose(): + try: + real.shutdown_abort(grace=30, reason=TimeoutError("session-reuse retire")) + except Exception: + try: + real.shutdown(wait=False) + except Exception: + pass + + threading.Thread(target=_dispose, daemon=True, name="session-reuse-retire").start() # ---- factory installed at the pool-creation seam ---- From f09acfe47024b542357a501cac6e8e8c921b84ad Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 8 Jul 2026 19:40:12 -0700 Subject: [PATCH 06/12] [None][test] Retire wedged pools by SIGKILLing worker PIDs The previous hardening disposed of broken pools via shutdown_abort, which calls MPI_COMM_WORLD.Abort and kills the parent test process along with the workers (observed as 'MPI_ABORT was invoked on rank 0' taking down the whole pytest session right after a pool rebuild). Record the worker PIDs at spawn (best-effort submit_sync rounds on the fresh, healthy pool) and have _retire SIGKILL them before reaping the client side with a plain shutdown. Retired pools are discarded, so nothing needs a graceful stop, and the driver reclaims GPU memory on process death. Missing PIDs simply fall back to graceful shutdown. Signed-off-by: qgai --- tests/test_common/session_reuse.py | 49 +++++++++++++++++++++++++----- 1 file changed, 41 insertions(+), 8 deletions(-) diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py index df382f97ce9c..76f80b427d2a 100644 --- a/tests/test_common/session_reuse.py +++ b/tests/test_common/session_reuse.py @@ -160,6 +160,30 @@ def _free_total(): pass +def _get_worker_pid() -> int: + """Runs inside a worker; module-level so it is picklable.""" + return os.getpid() + + +def _collect_worker_pids(real, n_workers: int) -> tuple: + """Record the worker PIDs of a freshly spawned pool. + + ``_retire`` uses them to SIGKILL wedged workers: a graceful shutdown + blocks forever on a broken pool and ``shutdown_abort`` would MPI_Abort + the parent test process too. Best effort — a missing PID just means that + worker is left for the graceful-shutdown fallback. + """ + pids: set = set() + try: + for _ in range(4): + pids.update(real.submit_sync(_get_worker_pid)) + if len(pids) >= n_workers: + break + except Exception: + pass + return tuple(sorted(pids)) + + def _describe_mismatch(spawn_snap, now_snap, uses, max_uses): """One line naming WHY a cached pool cannot be handed out (observability).""" if uses >= max_uses: @@ -241,19 +265,27 @@ def _retire(real): Pools are usually retired because they are broken (failed health probe, stale env, lifetime cap). A graceful ``shutdown()`` on a pool whose workers are wedged in a collective blocks forever and leaks the - workers' GPU memory into subsequent tests, so prefer - ``shutdown_abort`` (bounded grace, then kill) and fall back to - ``shutdown`` only if abort is unavailable or raises. + workers' GPU memory into subsequent tests. ``shutdown_abort`` is NOT + an option either: it calls ``MPI_COMM_WORLD.Abort``, which kills the + parent test process along with the workers. Instead SIGKILL the + worker PIDs recorded at spawn (retired pools are discarded, so + nothing needs a graceful stop; the driver reclaims GPU memory on + process death) and then reap the client side. """ + pids = getattr(real, "_reuse_worker_pids", ()) def _dispose(): - try: - real.shutdown_abort(grace=30, reason=TimeoutError("session-reuse retire")) - except Exception: + import signal + + for pid in pids: try: - real.shutdown(wait=False) - except Exception: + os.kill(pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError): pass + try: + real.shutdown() + except Exception: + pass threading.Thread(target=_dispose, daemon=True, name="session-reuse-retire").start() @@ -363,6 +395,7 @@ def _spawn_fresh(self, real_cls, n_workers): os.environ.pop(k, None) real._reuse_uses = 0 real._reuse_spawn_snapshot = snapshot + real._reuse_worker_pids = _collect_worker_pids(real, n_workers) return real def _release(self, real): From 027be09dadc6ceb9c7d25e17ba39d23a39addaeb Mon Sep 17 00:00:00 2001 From: qgai Date: Fri, 10 Jul 2026 08:15:27 -0700 Subject: [PATCH 07/12] [None][test] Restrict SIGKILL disposal to broken pools only Healthy retires (lifetime cap, stale env snapshot, duplicate cache slot) go back to a graceful background shutdown: those workers are idle and exit cleanly, and abnormal termination of MPI-spawned children can upset the MPI runtime in the parent test process. The SIGKILL path is kept only for pools that failed the health probe, where workers may be wedged in a collective and a graceful shutdown would block forever. Signed-off-by: qgai --- tests/test_common/session_reuse.py | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py index 76f80b427d2a..b83d3b10d863 100644 --- a/tests/test_common/session_reuse.py +++ b/tests/test_common/session_reuse.py @@ -259,20 +259,23 @@ def max_uses(self) -> int: return int(os.environ.get("TRTLLM_TEST_REUSE_MAX_USES", "16")) @staticmethod - def _retire(real): + def _retire(real, broken: bool = False): """Dispose of a pool in the background without blocking the test. - Pools are usually retired because they are broken (failed health - probe, stale env, lifetime cap). A graceful ``shutdown()`` on a pool - whose workers are wedged in a collective blocks forever and leaks the - workers' GPU memory into subsequent tests. ``shutdown_abort`` is NOT - an option either: it calls ``MPI_COMM_WORLD.Abort``, which kills the - parent test process along with the workers. Instead SIGKILL the - worker PIDs recorded at spawn (retired pools are discarded, so - nothing needs a graceful stop; the driver reclaims GPU memory on - process death) and then reap the client side. + Healthy retires (lifetime cap, stale env snapshot, duplicate cache + slot) use a graceful ``shutdown()``: the workers are idle and exit + cleanly, and killing MPI-spawned children abnormally can upset the + MPI runtime in the parent process. + + ``broken=True`` (failed health probe) means the workers may be wedged + in a collective: a graceful shutdown would block forever and leak + their GPU memory into subsequent tests, and ``shutdown_abort`` calls + ``MPI_COMM_WORLD.Abort``, which kills the parent test process too. + Instead SIGKILL the worker PIDs recorded at spawn (a discarded pool + needs no graceful stop; the driver reclaims GPU memory on process + death) and then reap the client side. """ - pids = getattr(real, "_reuse_worker_pids", ()) + pids = getattr(real, "_reuse_worker_pids", ()) if broken else () def _dispose(): import signal @@ -371,7 +374,7 @@ def acquire(self, real_cls, n_workers): f"[session-reuse] cached pool failed reset, rebuilding: {e}", flush=True, ) - self._retire(real) + self._retire(real, broken=True) return _ReusableSession(self._spawn_fresh(real_cls, n_workers), self) def _spawn_fresh(self, real_cls, n_workers): From b3e7a1fbd0c240c02ce67768b56c21394084d34d Mon Sep 17 00:00:00 2001 From: qgai Date: Sun, 12 Jul 2026 23:54:37 -0700 Subject: [PATCH 08/12] [None][test] Drain cached pools after any failed test item A failed test's pool can pass the health probe (workers alive and responsive) while still carrying worker-side state the probe cannot see, and would otherwise serve up to max_uses-1 subsequent same-size tests. Add a pytest-level failure fence per review feedback: record failed items in pytest_runtest_logreport and drain ALL cached pools in pytest_runtest_logfinish once the failed item's teardown has returned its pool to the cache. This is a conservative fallback for contamination the probe cannot detect, not a replacement for the broken-pool retirement path. The passing-test path is unchanged and pays no extra cost; a failure adds one pool-shutdown plus a fresh spawn for the next test. Regression tests: after a failed item the next acquire gets a fresh pool; after a passing item reuse stays intact. Signed-off-by: qgai --- tests/test_common/session_reuse_hooks.py | 22 +++++++++++ tests/unittest/llmapi/test_session_reuse.py | 43 +++++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/tests/test_common/session_reuse_hooks.py b/tests/test_common/session_reuse_hooks.py index 5982973478e2..7bd4899250a7 100644 --- a/tests/test_common/session_reuse_hooks.py +++ b/tests/test_common/session_reuse_hooks.py @@ -52,5 +52,27 @@ def pytest_runtest_setup(item): REUSE.install_pool_factory_if_loaded() +# Failure fence: a failed test's pool can pass the health probe (workers alive +# and responsive) while still carrying worker-side state the probe cannot see, +# and would otherwise serve up to max_uses-1 more tests. Treat any failed item +# as untrusted and drain ALL cached pools after its teardown (by which time +# the pool has been returned to the cache). This is a conservative fallback +# for contamination the probe cannot detect — not a replacement for the +# broken-pool retirement path — and costs nothing on the passing-test path. +_FAILED_ITEMS: set = set() + + +def pytest_runtest_logreport(report): + if report.failed: + _FAILED_ITEMS.add(report.nodeid) + + +def pytest_runtest_logfinish(nodeid, location): + if nodeid in _FAILED_ITEMS: + _FAILED_ITEMS.discard(nodeid) + if REUSE.enabled: + REUSE.drain() + + def pytest_sessionfinish(session, exitstatus): REUSE.drain() diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py index 2360f7b0ebcd..717027fc148e 100644 --- a/tests/unittest/llmapi/test_session_reuse.py +++ b/tests/unittest/llmapi/test_session_reuse.py @@ -195,3 +195,46 @@ def test_autodeploy_nodeids_are_private(): assert not _is_private_nodeid( "unittest/_torch/speculative/test_eagle3.py::test_llama_eagle3[x]" ) + + +def test_failed_item_fences_cached_pools(reuse_cache, monkeypatch): + """After a failed item the next acquire must NOT reuse the cached pool.""" + from test_common import session_reuse_hooks as hooks + + monkeypatch.setattr(hooks, "REUSE", reuse_cache) + + s = reuse_cache.acquire(_FakePool, 2) + real = s._real + s.shutdown() # pool returned to the cache + + class _FailedReport: + failed = True + nodeid = "tests/foo.py::test_bar" + + hooks.pytest_runtest_logreport(_FailedReport) + hooks.pytest_runtest_logfinish(_FailedReport.nodeid, None) + + assert real.shut # fence drained the cached pool + s2 = reuse_cache.acquire(_FakePool, 2) + assert s2._real is not real # fresh pool, no reuse across the failure + + +def test_passing_item_keeps_reuse(reuse_cache, monkeypatch): + """The fence must not fire for passing items; reuse stays intact.""" + from test_common import session_reuse_hooks as hooks + + monkeypatch.setattr(hooks, "REUSE", reuse_cache) + + s = reuse_cache.acquire(_FakePool, 2) + real = s._real + s.shutdown() + + class _PassedReport: + failed = False + nodeid = "tests/foo.py::test_ok" + + hooks.pytest_runtest_logreport(_PassedReport) + hooks.pytest_runtest_logfinish(_PassedReport.nodeid, None) + + s2 = reuse_cache.acquire(_FakePool, 2) + assert s2._real is real # pool survived and was reused From c9ae719e40b8b63c6638aba0d32f793fdd20cd46 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 13 Jul 2026 01:28:19 -0700 Subject: [PATCH 09/12] [None][test] Make drain threads daemonic A wedged pool shutdown inside a drain thread must not keep the interpreter alive at process exit: the bounded join only protects the drain caller, while a non-daemon thread would make pytest hang until the CI stage timeout. The retire path already uses a daemon thread; apply the same to drain, per review feedback. Signed-off-by: qgai --- tests/test_common/session_reuse.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py index b83d3b10d863..48b137ce2b26 100644 --- a/tests/test_common/session_reuse.py +++ b/tests/test_common/session_reuse.py @@ -425,7 +425,12 @@ def drain(self) -> None: pools, self._pools = list(self._pools.values()), {} if not pools: return - threads = [threading.Thread(target=p.shutdown, name="session-reuse-drain") for p in pools] + threads = [ + # daemon: a wedged pool shutdown must not keep the interpreter + # alive at exit (a non-daemon thread would hang the CI stage). + threading.Thread(target=p.shutdown, name="session-reuse-drain", daemon=True) + for p in pools + ] for t in threads: t.start() for t in threads: From 9e602688c13247f55ad7db600d39b98c7c997b0f Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 13 Jul 2026 01:31:02 -0700 Subject: [PATCH 10/12] [None][test] Guard the retire SIGKILL against PID recycling If a broken pool's worker has already exited, the OS may recycle its PID before the background retire thread runs, and os.kill would then hit an unrelated process owned by the same user - most plausibly a worker of the replacement pool being spawned concurrently. Record (pid, start_time from /proc/pid/stat) pairs at spawn and re-verify the start time right before SIGKILL; a mismatch or missing process means the PID no longer belongs to the recorded worker and is skipped, falling back to the graceful client-side reap. Per review feedback. Signed-off-by: qgai --- tests/test_common/session_reuse.py | 32 +++++++++++++++++++++++++----- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py index 48b137ce2b26..6f6ddc201353 100644 --- a/tests/test_common/session_reuse.py +++ b/tests/test_common/session_reuse.py @@ -160,9 +160,26 @@ def _free_total(): pass -def _get_worker_pid() -> int: +def _proc_start_time(pid: int): + """Kernel start time (jiffies since boot) of ``pid``, or None if gone. + + PIDs are recycled by the OS, but the (pid, start_time) pair is unique: + verifying it right before SIGKILL prevents killing an unrelated process + (e.g. a replacement pool's worker) that inherited a dead worker's PID. + """ + try: + with open(f"/proc/{pid}/stat", "rb") as f: + stat = f.read() + # Field 2 (comm) may contain spaces/parens; parse after the last ')'. + return stat.rsplit(b")", 1)[1].split()[19] # field 22 overall + except OSError: + return None + + +def _get_worker_pid() -> tuple: """Runs inside a worker; module-level so it is picklable.""" - return os.getpid() + pid = os.getpid() + return (pid, _proc_start_time(pid)) def _collect_worker_pids(real, n_workers: int) -> tuple: @@ -170,8 +187,9 @@ def _collect_worker_pids(real, n_workers: int) -> tuple: ``_retire`` uses them to SIGKILL wedged workers: a graceful shutdown blocks forever on a broken pool and ``shutdown_abort`` would MPI_Abort - the parent test process too. Best effort — a missing PID just means that - worker is left for the graceful-shutdown fallback. + the parent test process too. Records (pid, start_time) pairs so the kill + can verify the PID was not recycled. Best effort — a missing PID just + means that worker is left for the graceful-shutdown fallback. """ pids: set = set() try: @@ -280,7 +298,11 @@ def _retire(real, broken: bool = False): def _dispose(): import signal - for pid in pids: + 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 _proc_start_time(pid) != start_time: + continue try: os.kill(pid, signal.SIGKILL) except (ProcessLookupError, PermissionError): From f75b2dee88a1e581d042f5de9219f96a81999004 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 13 Jul 2026 01:33:59 -0700 Subject: [PATCH 11/12] [None][test] Reap in-flight retire threads at drain points Retire threads were fire-and-forget: a disposal still in flight when the session ends would be cut off by interpreter exit (daemon threads), leaving a graceful shutdown half-done. Track them in a registry and have drain() - which runs at the natural rendezvous points (failure fence, opt-out, session finish) - join them with the same bounded wait used for pool shutdowns, per review feedback. The per-test hot path stays non-blocking, and daemon=True still guarantees a wedged disposal cannot hang interpreter exit. Signed-off-by: qgai --- tests/test_common/session_reuse.py | 30 ++++++++++++++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py index 6f6ddc201353..75475ac358f3 100644 --- a/tests/test_common/session_reuse.py +++ b/tests/test_common/session_reuse.py @@ -160,6 +160,10 @@ def _free_total(): pass +_RETIRE_THREADS: list = [] +_RETIRE_LOCK = threading.Lock() + + def _proc_start_time(pid: int): """Kernel start time (jiffies since boot) of ``pid``, or None if gone. @@ -312,7 +316,14 @@ def _dispose(): except Exception: pass - threading.Thread(target=_dispose, daemon=True, name="session-reuse-retire").start() + t = threading.Thread(target=_dispose, daemon=True, name="session-reuse-retire") + t.start() + # Track in-flight disposals so drain() can reap them at natural + # rendezvous points (failure fence / session finish) with a bounded + # join; the hot path stays non-blocking and daemon=True still + # guarantees a wedged disposal cannot hang interpreter exit. + with _RETIRE_LOCK: + _RETIRE_THREADS.append(t) # ---- factory installed at the pool-creation seam ---- @@ -442,7 +453,22 @@ def suspend(self, suspended: bool) -> None: self._suspended = suspended def drain(self) -> None: - """Shut down all cached pools in parallel (frees GPU/CPU footprint).""" + """Shut down all cached pools in parallel (frees GPU/CPU footprint). + + Also reaps in-flight retire threads: drain runs at natural rendezvous + points (failure fence, opt-out, session finish), so waiting here keeps + disposals from leaking past the session without ever blocking the + per-test hot path. The join is bounded for the same reason as below. + """ + with _RETIRE_LOCK: + in_flight, _RETIRE_THREADS[:] = list(_RETIRE_THREADS), [] + for t in in_flight: + t.join(timeout=60) + if t.is_alive(): + print( + "[session-reuse] WARNING: pool retirement did not finish within 60s", + flush=True, + ) with self._lock: pools, self._pools = list(self._pools.values()), {} if not pools: From b086d20903df1e8ae0d0cb8a414ac1eaddb75387 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 13 Jul 2026 02:18:23 -0700 Subject: [PATCH 12/12] [None][test] Pin per-worker probe tasks with an MPI barrier MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MPIPoolExecutor has no one-task-per-worker guarantee, so both the pool health probe (8 resubmit rounds) and worker PID collection (4 tries) compensated with retry-and-dedup loops. Replace both with one exactly-once primitive: each submitted task opens with an MPI_Barrier across the worker world, so a worker holding a task blocks until every worker holds one — by pigeonhole the n_workers tasks land exactly one per worker, in a single round. This also strengthens the probe: completing the barrier proves the worker world's collectives still function (the failure mode wedged pools actually exhibit), not just the task loop. Timeout semantics and the broken-pool retirement path are unchanged. Validated on a real 4-GPU MPIPoolExecutor spawn: 20/20 probes covered exactly 4 workers with worker COMM_WORLD size 4, and a wedged worker turned the probe into a bounded RuntimeError at the timeout. Signed-off-by: qgai --- tests/test_common/grouped_test_utils.py | 70 +++++++++++---------- tests/test_common/session_reuse.py | 23 +++---- tests/unittest/llmapi/test_session_reuse.py | 57 +++++++++++++++++ 3 files changed, 104 insertions(+), 46 deletions(-) diff --git a/tests/test_common/grouped_test_utils.py b/tests/test_common/grouped_test_utils.py index a5eb9f02fee3..aa77aba7ed1a 100644 --- a/tests/test_common/grouped_test_utils.py +++ b/tests/test_common/grouped_test_utils.py @@ -9,49 +9,53 @@ """ -def _run_and_get_worker_rank(fn) -> int: - fn() - from tensorrt_llm._utils import mpi_rank +def _barrier_run(fn): + """Runs inside a worker; module-level so it is picklable. - return mpi_rank() + The leading barrier is what makes ``submit_sync_per_worker`` exactly-once: + a worker that picked up one of the ``n_workers`` submitted tasks blocks + here until every other worker holds a task of its own, so no worker can + drain a second one (pigeonhole). The workers' ``MPI_COMM_WORLD`` is the + spawned worker world (the parent process is not a member), so the barrier + spans exactly the pool's ``n_workers`` ranks. + """ + from mpi4py import MPI + + MPI.COMM_WORLD.barrier() + return (MPI.COMM_WORLD.Get_rank(), fn()) -def submit_sync_per_worker( - mpi_session, fn, max_rounds: int = 8, round_timeout: float = 60.0 -) -> None: - """Run ``fn`` at least once on EVERY worker of the pool. +def submit_sync_per_worker(mpi_session, fn, timeout: float = 60.0) -> list: + """Run ``fn`` exactly ONCE on every worker; return results ordered by rank. ``MpiPoolSession.submit_sync`` enqueues ``n_workers`` tasks, but ``MPIPoolExecutor`` gives no one-task-per-worker guarantee: a fast task can - be drained twice by one idle worker, leaving another untouched. Verify - coverage by collecting the worker ranks that actually ran ``fn`` and - resubmitting until all workers are covered. + be drained twice by one idle worker, leaving another untouched. Each task + therefore opens with a barrier across the worker world (see + ``_barrier_run``), which pins the tasks one-per-worker. - Each round is bounded by ``round_timeout``: a worker that is alive but - wedged (e.g. stuck in a collective after the previous test's executor died - mid-shutdown) never completes its future, so an unbounded ``result()`` - would hang the whole session. A timed-out round is reported as probe + This doubles as the pool health probe: completing the barrier proves every + worker is alive AND that the worker world's collectives still function — + the thing tests actually depend on. The wait is bounded by ``timeout``: a + wedged worker (e.g. stuck in a collective after the previous test's + executor died mid-shutdown) never completes the barrier, so an unbounded + ``result()`` would hang the whole session. A timeout is reported as probe failure, making the caller retire the pool and spawn a fresh one. """ import concurrent.futures - expected = mpi_session.n_workers - seen: set = set() - for _ in range(max_rounds): - futures = mpi_session.submit(_run_and_get_worker_rank, fn) - done, not_done = concurrent.futures.wait(futures, timeout=round_timeout) - if not_done: - raise RuntimeError( - f"health probe timed out after {round_timeout}s: " - f"{len(not_done)}/{len(futures)} worker tasks never returned" - ) - seen.update(f.result() for f in done) - if len(seen) >= expected: - return - raise RuntimeError( - f"task only reached worker ranks {sorted(seen)} after {max_rounds} " - f"rounds; expected {expected} distinct workers" - ) + futures = mpi_session.submit(_barrier_run, fn) + done, not_done = concurrent.futures.wait(futures, timeout=timeout) + if not_done: + raise RuntimeError( + f"health probe timed out after {timeout}s: " + f"{len(not_done)}/{len(futures)} worker tasks never returned" + ) + by_rank = dict(f.result() for f in done) # .result() re-raises worker errors + missing = set(range(mpi_session.n_workers)) - set(by_rank) + if missing: # unreachable given the barrier; guards scheduler surprises + raise RuntimeError(f"no task ran on worker ranks {sorted(missing)}") + return [result for _, result in sorted(by_rank.items())] def reset_worker_torch_compile_state() -> None: @@ -64,7 +68,7 @@ def reset_worker_torch_compile_state() -> None: ``recompile_limit`` (16), which is a HARD failure under ``fullgraph=True`` (``FailOnRecompileLimitHit``) and aborts the whole MPI job. Resetting between cases makes each LLM start from a clean compile cache, like a fresh - process. Submit this to each worker via ``mpi_session.submit_sync(...)``. + process. Run on every worker via ``submit_sync_per_worker``. """ import torch diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py index 75475ac358f3..145ee5cf2248 100644 --- a/tests/test_common/session_reuse.py +++ b/tests/test_common/session_reuse.py @@ -22,8 +22,9 @@ - use-count cap -> pool retired after N handouts (default 16), bounding worker state accumulation -Between handouts every worker runs a torch.compile/Dynamo reset (rank-verified -via ``grouped_test_utils.submit_sync_per_worker``) and the handover waits for +Between handouts every worker runs a torch.compile/Dynamo reset (exactly once +per worker, barrier-pinned: ``grouped_test_utils.submit_sync_per_worker``) and +the handover waits for the previous worker's GPU memory to actually be released (NVML settle barrier) — both failure modes were observed in validation, not hypothetical. @@ -186,24 +187,20 @@ def _get_worker_pid() -> tuple: return (pid, _proc_start_time(pid)) -def _collect_worker_pids(real, n_workers: int) -> tuple: +def _collect_worker_pids(real) -> tuple: """Record the worker PIDs of a freshly spawned pool. ``_retire`` uses them to SIGKILL wedged workers: a graceful shutdown blocks forever on a broken pool and ``shutdown_abort`` would MPI_Abort the parent test process too. Records (pid, start_time) pairs so the kill - can verify the PID was not recycled. Best effort — a missing PID just - means that worker is left for the graceful-shutdown fallback. + can verify the PID was not recycled. ``submit_sync_per_worker`` runs the + collection exactly once per worker. Best effort — if it fails, the pool + just falls back to graceful shutdown. """ - pids: set = set() try: - for _ in range(4): - pids.update(real.submit_sync(_get_worker_pid)) - if len(pids) >= n_workers: - break + return tuple(sorted(submit_sync_per_worker(real, _get_worker_pid))) except Exception: - pass - return tuple(sorted(pids)) + return () def _describe_mismatch(spawn_snap, now_snap, uses, max_uses): @@ -431,7 +428,7 @@ def _spawn_fresh(self, real_cls, n_workers): os.environ.pop(k, None) real._reuse_uses = 0 real._reuse_spawn_snapshot = snapshot - real._reuse_worker_pids = _collect_worker_pids(real, n_workers) + real._reuse_worker_pids = _collect_worker_pids(real) return real def _release(self, real): diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py index 717027fc148e..3a2f6459a990 100644 --- a/tests/unittest/llmapi/test_session_reuse.py +++ b/tests/unittest/llmapi/test_session_reuse.py @@ -219,6 +219,63 @@ class _FailedReport: assert s2._real is not real # fresh pool, no reuse across the failure +class _FakeBarrierSession: + """Fakes MpiPoolSession.submit for master-side submit_sync_per_worker tests. + + Worker-side behaviour (the barrier itself) needs real MPI and is covered + by the multi-GPU validation runs; here each entry stands for one worker's + future: a (rank, result) tuple, an Exception, or None (never completes, + i.e. a wedged worker). + """ + + def __init__(self, outcomes): + import concurrent.futures + + self.n_workers = len(outcomes) + self._futures = [] + for outcome in outcomes: + f = concurrent.futures.Future() + if isinstance(outcome, Exception): + f.set_exception(outcome) + elif outcome is not None: + f.set_result(outcome) + self._futures.append(f) + + def submit(self, task, *args): + return self._futures + + +def test_probe_returns_results_ordered_by_rank(): + from test_common.grouped_test_utils import submit_sync_per_worker + + session = _FakeBarrierSession([(2, "c"), (0, "a"), (1, "b")]) + assert submit_sync_per_worker(session, lambda: None) == ["a", "b", "c"] + + +def test_probe_times_out_on_wedged_worker(): + from test_common.grouped_test_utils import submit_sync_per_worker + + session = _FakeBarrierSession([(0, "a"), None]) # rank 1 never returns + with pytest.raises(RuntimeError, match="timed out"): + submit_sync_per_worker(session, lambda: None, timeout=0.2) + + +def test_probe_propagates_worker_exception(): + from test_common.grouped_test_utils import submit_sync_per_worker + + session = _FakeBarrierSession([(0, "a"), ValueError("worker died")]) + with pytest.raises(ValueError, match="worker died"): + submit_sync_per_worker(session, lambda: None) + + +def test_probe_flags_missing_rank(): + from test_common.grouped_test_utils import submit_sync_per_worker + + session = _FakeBarrierSession([(0, "a"), (0, "a")]) # rank 1 never covered + with pytest.raises(RuntimeError, match=r"ranks \[1\]"): + submit_sync_per_worker(session, lambda: None) + + def test_passing_item_keeps_reuse(reuse_cache, monkeypatch): """The fence must not fire for passing items; reuse stays intact.""" from test_common import session_reuse_hooks as hooks