From cfb3df019287521e4604de753d3372ff17ff96ff Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 7 Jul 2026 02:05:07 -0700 Subject: [PATCH 1/8] [None][test] Automatic MPI session reuse for tests via shared pytest infra 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 for ini-less dirs) lazily patches the pool-construction seams; 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 each time. Eligibility is automatic; tests with special needs keep their private lifecycle: RPC executors get a drain-then-build factory at their own seam (private, never-cached pools); tests passing their own _mpi_session never reach the patched seam; @pytest.mark.private_mpi_session opts out explicitly; pools are retired when env or sys.path changed since spawn (workers freeze both), after a use-count cap (default 16), or when the per-handout health probe fails - every failure path degrades to a fresh synchronous build. Disabled under pytest-xdist. TRTLLM_TEST_REUSE_SESSION=0 turns everything off. Between handouts every worker runs a rank-verified torch._dynamo reset and the handover waits for the previous worker's GPU memory release (NVML settle barrier). Cache-managed pool spawns also enable the worker-side HF weight cache (companion weight-loader PR) scoped to those pools only. Validated on B200: 12 pure-logic unit tests; unmodified single-GPU tests reuse a 1-worker pool; multi-GPU suite 15 passed with 10 handovers; DeepSeek-V3-Lite accuracy tp4/ep4/tp2pp2/pp4 4 passed on one shared pool, GSM8K at reference; plus a ~2.5h 550-case soak (42 handovers, 2 full retire/rebuild cycles, 0 failures). Split out of #15777 per review (test-side only; works with or without the companion MPI-session-ownership and weight-cache PRs, degrading gracefully). 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 | 366 ++++++++++++++++++++ tests/test_common/session_reuse_hooks.py | 36 ++ tests/unittest/llmapi/test_session_reuse.py | 155 +++++++++ tests/unittest/pytest.ini | 9 +- 7 files changed, 644 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..06e2ec52b642 --- /dev/null +++ b/tests/test_common/session_reuse.py @@ -0,0 +1,366 @@ +# 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(timeout: float = _SETTLE_TIMEOUT_S) -> 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 >= timeout: + 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): + # Set via __dict__ to avoid __getattr__ recursion. + self.__dict__["_real"] = real + self.__dict__["_cache"] = cache + self.__dict__["_dead"] = False + + def __getattr__(self, name): + return getattr(self.__dict__["_real"], name) + + def shutdown(self): + if self.__dict__["_dead"]: + return + self.__dict__["_cache"]._release(self.__dict__["_real"]) + + def shutdown_abort(self, *args, **kwargs): + self.__dict__["_dead"] = True + self.__dict__["_cache"]._forget(self.__dict__["_real"]) + return self.__dict__["_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 + 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: + # Opt-out test (private_mpi_session): 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: + t.join() + 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..dbba1035cb59 --- /dev/null +++ b/tests/unittest/llmapi/test_session_reuse.py @@ -0,0 +1,155 @@ +# 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_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 26a983932a41a1e552f3ce7506740621ab5a4bf5 Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 7 Jul 2026 02:44:47 -0700 Subject: [PATCH 2/8] [None][test] Harden the reuse wrapper per second review pass - A wrapper whose pool was already released to the cache ignores late shutdown_abort calls (executor error paths can fire after shutdown; the pool may belong to the NEXT test by then); shutdown is idempotent; plain attributes replace the __dict__ indirection. - acquire() consults the TRTLLM_TEST_REUSE_SESSION kill switch on every call, so flipping it off after the seams were patched really restores private-pool behavior. - drain() joins each pool shutdown with a 60s bound and warns instead of hanging the suite on one wedged pool; the factory bypass branch logs instead of silently disabling reuse; wait_gpu_memory_settle drops its never-passed timeout parameter. Two new unit tests cover the released-wrapper abort guard and the mid-suite kill switch (14 pass). Signed-off-by: qgai --- tests/test_common/session_reuse.py | 52 +++++++++++++++------ tests/unittest/llmapi/test_session_reuse.py | 23 +++++++++ 2 files changed, 60 insertions(+), 15 deletions(-) diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py index 06e2ec52b642..6a616715801d 100644 --- a/tests/test_common/session_reuse.py +++ b/tests/test_common/session_reuse.py @@ -108,7 +108,7 @@ def _visible_gpu_indices(count: int): return indices or list(range(count)) -def wait_gpu_memory_settle(timeout: float = _SETTLE_TIMEOUT_S) -> None: +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. @@ -141,7 +141,7 @@ def _free_total(): break # not increasing: that memory is legitimately in use else: flat = 0 - if time.monotonic() - t0 >= timeout: + if time.monotonic() - t0 >= _SETTLE_TIMEOUT_S: break prev = free time.sleep(_SETTLE_POLL_S) @@ -184,23 +184,32 @@ class _ReusableSession: """ def __init__(self, real, cache): - # Set via __dict__ to avoid __getattr__ recursion. - self.__dict__["_real"] = real - self.__dict__["_cache"] = cache - self.__dict__["_dead"] = False + 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.__dict__["_dead"]: + if self._dead or self._released: return - self.__dict__["_cache"]._release(self.__dict__["_real"]) + self._released = True + self._cache._release(self._real) def shutdown_abort(self, *args, **kwargs): - self.__dict__["_dead"] = True - self.__dict__["_cache"]._forget(self.__dict__["_real"]) - return self.__dict__["_real"].shutdown_abort(*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: @@ -251,6 +260,11 @@ def install_pool_factory_if_loaded(self) -> None: 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) @@ -270,9 +284,10 @@ def rpc_factory(n_workers, *args, **kwargs): def acquire(self, real_cls, n_workers): """Hand out a cached same-size pool (reset + settled) or build one.""" - if self._suspended: - # Opt-out test (private_mpi_session): untracked fresh pool that - # the LLM owns and destroys normally. + 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) @@ -359,7 +374,14 @@ def drain(self) -> None: for t in threads: t.start() for t in threads: - t.join() + # 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) diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py index dbba1035cb59..03514356ccd6 100644 --- a/tests/unittest/llmapi/test_session_reuse.py +++ b/tests/unittest/llmapi/test_session_reuse.py @@ -147,6 +147,29 @@ def test_weight_cache_env_respects_user_setting(reuse_cache, monkeypatch): 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 From b7e421aa7c9f1773008e0b6bc2c38fa4d466915d Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 7 Jul 2026 19:24:47 -0700 Subject: [PATCH 3/8] [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 d8a9f42000b40dbd3314febbe12e76105d6cf56b Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 8 Jul 2026 06:32:34 -0700 Subject: [PATCH 4/8] [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 d1af17d3e637836795d095f5e8808790d07758b9 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 8 Jul 2026 19:16:55 -0700 Subject: [PATCH 5/8] [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 61d61ec59388ba84dd6a7003edff9844b202996e Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 8 Jul 2026 19:34:28 -0700 Subject: [PATCH 6/8] [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 b696e73961eb5764be39df22d53d0f976b2b7c5d Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 8 Jul 2026 19:40:41 -0700 Subject: [PATCH 7/8] [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 ae405f8dc0e39b89e058dc980f9600011fd7bac8 Mon Sep 17 00:00:00 2001 From: qgai Date: Fri, 10 Jul 2026 08:16:07 -0700 Subject: [PATCH 8/8] [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):