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..aa77aba7ed1a --- /dev/null +++ b/tests/test_common/grouped_test_utils.py @@ -0,0 +1,75 @@ +# 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 _barrier_run(fn): + """Runs inside a worker; module-level so it is picklable. + + 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, 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. Each task + therefore opens with a barrier across the worker world (see + ``_barrier_run``), which pins the tasks one-per-worker. + + 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 + + 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: + """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. Run on every worker via ``submit_sync_per_worker``. + """ + 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..145ee5cf2248 --- /dev/null +++ b/tests/test_common/session_reuse.py @@ -0,0 +1,493 @@ +# 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 (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. + +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 + + +_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. + + 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.""" + pid = os.getpid() + return (pid, _proc_start_time(pid)) + + +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. ``submit_sync_per_worker`` runs the + collection exactly once per worker. Best effort — if it fails, the pool + just falls back to graceful shutdown. + """ + try: + return tuple(sorted(submit_sync_per_worker(real, _get_worker_pid))) + except Exception: + return () + + +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, broken: bool = False): + """Dispose of a pool in the background without blocking the test. + + 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", ()) if broken else () + + def _dispose(): + import signal + + 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): + pass + try: + real.shutdown() + except Exception: + pass + + 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 ---- + + 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, broken=True) + 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 + real._reuse_worker_pids = _collect_worker_pids(real) + 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). + + 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: + return + 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: + # 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..7bd4899250a7 --- /dev/null +++ b/tests/test_common/session_reuse_hooks.py @@ -0,0 +1,78 @@ +# 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 + +# 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( + "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 or _is_private_nodeid( + item.nodeid + ) + if opt_out: + REUSE.drain() + REUSE.suspend(opt_out) + 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/_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): diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py new file mode 100644 index 000000000000..3a2f6459a990 --- /dev/null +++ b/tests/unittest/llmapi/test_session_reuse.py @@ -0,0 +1,297 @@ +# 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 + + +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]" + ) + + +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 + + +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 + + 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 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