diff --git a/tensorrt_llm/llmapi/mpi_session.py b/tensorrt_llm/llmapi/mpi_session.py index 182ca1ed9d76..02a1415fb847 100644 --- a/tensorrt_llm/llmapi/mpi_session.py +++ b/tensorrt_llm/llmapi/mpi_session.py @@ -8,6 +8,7 @@ import traceback from collections.abc import Callable from concurrent.futures import Future, ThreadPoolExecutor, as_completed +from concurrent.futures import wait as futures_wait from typing import Any, Dict, List, NamedTuple, Optional, Tuple, TypeVar import zmq @@ -153,12 +154,67 @@ def shutdown_abort(self, grace: float = 60, reason=None): killer.join() +def _process_start_time(pid: int) -> Optional[bytes]: + """Kernel start time (jiffies since boot) of ``pid``, or None if gone. + + PIDs are recycled by the OS, but the (pid, start_time) pair uniquely + identifies a process incarnation — comparing it prevents waiting on an + unrelated process 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 _worker_identity_barrier(): + """Runs inside a pool worker; module-level so it is picklable. + + The leading barrier pins the ``n_workers`` submitted tasks one-per-worker + (a worker holding one task blocks until every other worker holds its own, + so no worker can drain a second one), collecting every worker's identity + exactly once. The workers' ``MPI_COMM_WORLD`` is the spawned worker world + (the parent process is not a member). + """ + from mpi4py import MPI + MPI.COMM_WORLD.barrier() + pid = os.getpid() + return (pid, _process_start_time(pid)) + + class MpiPoolSession(MpiSession): - def __init__(self, n_workers: int): + def __init__(self, + n_workers: int, + wait_shutdown: bool = False, + env_overrides: Optional[Dict[str, str]] = None): + """Args: + n_workers: number of MPI workers to spawn. + wait_shutdown: when True, ``shutdown()`` blocks until the spawned + worker processes have actually exited. ``MPIPoolExecutor.shutdown`` + returns at disconnect, but a worker's GPU memory is only released + when its process exits; callers that start new GPU work right + after ``shutdown()`` (e.g. CI test infrastructure handing a + pre-spawned pool to the next test) race that release and can OOM. + Off by default: production teardown does not need the barrier and + keeps its current latency. + env_overrides: extra environment variables to set in the WORKERS at + spawn, on top of the TRTLLM*/TLLM* variables forwarded from the + parent. The parent process environment is never touched — this + replaces the racy "set os.environ around the spawn, then restore" + pattern for callers that spawn pools from background threads. + """ self.n_workers = n_workers + self._wait_shutdown = wait_shutdown + self._env_overrides = dict(env_overrides) if env_overrides else {} + self._worker_identities: Tuple = () self.mpi_pool: Optional[MPIPoolExecutor] = None self._start_mpi_pool() + if wait_shutdown: + self._worker_identities = self._collect_worker_identities() if ENABLE_MULTI_DEVICE: self.comm = mpi4py.MPI.COMM_WORLD @@ -187,6 +243,87 @@ def shutdown(self, wait=True): self.mpi_pool.shutdown(wait=wait) logger.info("MpiPoolSession.shutdown: done") self.mpi_pool = None + if self._wait_shutdown: + self._wait_workers_exit() + + def _collect_worker_identities(self) -> Tuple: + """(pid, start_time) of every worker, recorded right after spawn. + + FAIL-CLOSED (review requirement): ``wait_shutdown=True`` is a + contract — shutdown blocks until the workers exited. A pool without + complete identities cannot honor it, and returning it anyway would + silently downgrade to the old non-waiting behavior (the timeout can + trip on a slow-but-healthy bootstrap, and ``futures_wait`` does not + cancel the pending tasks). Instead of handing out such a pool, tear + it down and raise; callers fall back to a fresh spawn. + """ + try: + futures = [ + self.mpi_pool.submit(_worker_identity_barrier) + for _ in range(self.n_workers) + ] + done, not_done = futures_wait(futures, timeout=60.0) + identities = tuple(f.result() for f in done) + except Exception as e: + self._teardown_unidentified_pool(()) + raise RuntimeError( + f"MpiPoolSession(wait_shutdown=True): worker identity " + f"collection failed ({e}); pool torn down") from e + if (not_done or len(identities) != self.n_workers + or len({pid + for pid, _ in identities}) != self.n_workers + or any(start is None for _, start in identities)): + self._teardown_unidentified_pool(identities) + raise RuntimeError( + "MpiPoolSession(wait_shutdown=True): worker identity " + f"collection incomplete ({len(identities)}/{self.n_workers} " + "valid identities); pool torn down instead of handing out a " + "session that cannot honor the wait_shutdown contract") + return identities + + def _teardown_unidentified_pool(self, partial_identities: Tuple) -> None: + """Dispose of a pool whose identity collection failed. + + The workers may be stuck in the collection barrier (one of them + never picked up its task), so a graceful blocking shutdown could + hang; disconnect without waiting and SIGKILL the workers we did + identify (with the pid-recycling guard). Workers we never identified + exit with the MPI runtime teardown; if one is truly wedged it leaks + until job end — the same bounded leak class as any wedged pool. + """ + import signal + + try: + self.mpi_pool.shutdown(wait=False) + except Exception: + pass + self.mpi_pool = None + for pid, start in partial_identities: + if start is None or _process_start_time(pid) != start: + continue # gone already, or the PID was recycled + try: + os.kill(pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError): + pass + + def _wait_workers_exit(self, timeout: float = 30.0) -> None: + """Block until the spawned worker processes have actually exited. + + Bounded: a wedged worker stops blocking the caller after ``timeout`` + (its memory is not coming back anyway; the caller's own recovery — + e.g. an OOM retry or a fresh spawn — takes over from there). + """ + deadline = time.monotonic() + timeout + for pid, start in self._worker_identities: + if start is None: + continue + while _process_start_time(pid) == start: + if time.monotonic() >= deadline: + logger.warning( + f"MpiPoolSession.shutdown: worker pid {pid} still " + f"alive after {timeout}s; not waiting further") + return + time.sleep(0.05) def abort(self): self.get_comm().Abort(1) @@ -199,6 +336,7 @@ def _start_mpi_pool(self): for key, value in os.environ.items() if key.startswith("TRTLLM") or key.startswith("TLLM") } + env.update(self._env_overrides) self.mpi_pool = MPIPoolExecutor(max_workers=self.n_workers, path=sys.path, env=env) diff --git a/tests/conftest.py b/tests/conftest.py index b0d2feebd7fc..b40dd597ce76 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,13 +1,18 @@ # 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/. +"""Fallback wiring of the session-reuse and session-prefetch plugins. + +tests/unittest and tests/integration/defs load the plugins through their own +pytest.ini / conftest; their rootdir sits below this file, so this conftest +is never collected there (no double registration). Any other directory under +tests/ (current or future) picks the hooks up from here, so MPI session +reuse and prefetch cover every test under tests/. + +Both plugins define same-named hooks, and ``pytest_plugins`` is only allowed +in a rootdir conftest (this file is not one when pytest runs from the repo +root) — so dispatch to both modules explicitly instead of importing their +hook functions into this namespace (the second import would silently shadow +the first). """ import os @@ -15,8 +20,30 @@ 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, -) +from test_common import session_prefetcher_hooks as _prefetch # noqa: E402 +from test_common import session_reuse_hooks as _reuse # noqa: E402 + + +def pytest_configure(config): + _reuse.pytest_configure(config) + _prefetch.pytest_configure(config) + + +def pytest_runtest_setup(item): + _reuse.pytest_runtest_setup(item) + _prefetch.pytest_runtest_setup(item) + + +def pytest_runtest_logreport(report): + # Reuse's failure fence (drain pools after a failed test); previously not + # wired in the fallback dirs at all. + _reuse.pytest_runtest_logreport(report) + + +def pytest_runtest_logfinish(nodeid, location): + _reuse.pytest_runtest_logfinish(nodeid, location) + + +def pytest_sessionfinish(session, exitstatus): + _reuse.pytest_sessionfinish(session, exitstatus) + _prefetch.pytest_sessionfinish(session, exitstatus) diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index 2ec7f4d81f11..691ae634d85e 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -39,6 +39,14 @@ import tqdm import yaml from _pytest.mark import ParameterSet +# Dispatched explicitly (not via pytest_plugins, which pytest forbids in a +# non-top-level conftest: a repo-root invocation like `pytest tests` loads +# this file as a NESTED conftest and would fail collection; and not via "-p" +# in pytest.ini addopts, which imports at preparse, before the ini pythonpath +# entries are usable). The wrappers below forward to the plugin; hooks are +# idempotent, so a repo-root run that also dispatches from tests/conftest.py +# is harmless. +from test_common import session_prefetcher_hooks as _prefetch_hooks from tensorrt_llm.bindings import ipc_nvls_supported from tensorrt_llm.llmapi.mpi_session import get_mpi_world_size @@ -2195,6 +2203,7 @@ def pytest_collection_modifyitems(session, config, items): def pytest_configure(config): + _prefetch_hooks.pytest_configure(config) os.environ.setdefault("TRTLLM_NO_USAGE_STATS", "1") # avoid thread leak of tqdm's TMonitor @@ -2666,3 +2675,11 @@ def torch_empty_cache() -> None: gc.collect() torch.cuda.empty_cache() gc.collect() + + +def pytest_runtest_setup(item): + _prefetch_hooks.pytest_runtest_setup(item) + + +def pytest_sessionfinish(session, exitstatus): + _prefetch_hooks.pytest_sessionfinish(session, exitstatus) diff --git a/tests/integration/defs/pytest.ini b/tests/integration/defs/pytest.ini index ee5e97c00f74..7e0cc6a053cf 100644 --- a/tests/integration/defs/pytest.ini +++ b/tests/integration/defs/pytest.ini @@ -1,10 +1,11 @@ [pytest] asyncio_default_fixture_loop_scope = module threadleak = True -# 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 +# Thread-\d+ \(_manager_spawn\) / session-reuse-* / session-prefetch-* belong to +# a pool cached for reuse by the NEXT test (tests/test_common/session_reuse.py) +# or prefetched for it (tests/test_common/session_prefetcher.py) and legitimately # outlive the test they start under. -threadleak_exclude = asyncio_\d+|Thread-\d+ \(_manager_spawn\)|session-reuse-\w+ +threadleak_exclude = asyncio_\d+|Thread-\d+ \(_manager_spawn\)|session-reuse-\w+|session-prefetch-\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 -p test_common.session_reuse_hooks pythonpath = diff --git a/tests/test_common/_session_utils.py b/tests/test_common/_session_utils.py new file mode 100644 index 000000000000..1c5f0cd76d09 --- /dev/null +++ b/tests/test_common/_session_utils.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Helpers shared by the MPI session reuse and session prefetch layers. + +``session_reuse.py`` (keeps pools alive across tests) and +``session_prefetcher.py`` (spawns the next pool in the background) manage the +same object — a live ``MpiPoolSession`` handed to a test that did not spawn +it — so they share the same invariant: the worker-visible state a pool +freezes at spawn. Keeping it in one place means a fix applies to both layers. + +Policy stays in the layers: what to DO with a snapshot mismatch (discard vs +proceed) is each layer's own decision. +""" + +import os +import sys + +# Workers freeze the parent environment AND sys.path at spawn time, so a +# pool spawned earlier 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", # changes every test phase by design + "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), + ) + + +def _isinstance_transparent_shim(real_cls, factory): + """A seam replacement that intercepts construction but stays a real type. + + The pool-creation seams used to hold a plain FUNCTION in place of + ``MpiPoolSession``. Library code that does ``isinstance(x, + MpiPoolSession)`` against the patched module attribute then raises + ``TypeError: isinstance() arg 2 must be a type`` — proxy.py's + killed-worker detection added exactly such a check and every bare + ``LLM()`` creation failed until it was worked around with an + exclusion-based match. This shim removes the hazard for good: a real + class whose metaclass routes construction to ``factory`` and + instance/subclass checks to ``real_cls``, so both usage patterns keep + working — including consumers added after this layer was written. + """ + + class _SeamMeta(type): + def __call__(cls, *args, **kwargs): + return factory(*args, **kwargs) + + def __instancecheck__(cls, obj): + return isinstance(obj, real_cls) + + def __subclasscheck__(cls, sub): + return issubclass(sub, real_cls) + + def __repr__(cls): + return f"" + + return _SeamMeta("MpiPoolSession", (), {}) diff --git a/tests/test_common/session_prefetcher.py b/tests/test_common/session_prefetcher.py new file mode 100644 index 000000000000..10f23d73cad9 --- /dev/null +++ b/tests/test_common/session_prefetcher.py @@ -0,0 +1,575 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Background prefetch of the NEXT test's MPI session — zero test changes. + +Multi-GPU LLM-API tests pay ~50-65s per bare ``LLM(...)`` to spawn an MPI +pool whose workers import ``tensorrt_llm``. That is pure CPU/IO work, so +while the CURRENT test runs on the GPUs a spare pool for the next test can +be spawned in a background thread — hiding the spawn cost behind the +previous test's runtime. Prefetched workers run no kernels and allocate +nothing before handover; depending on the library version, importing +tensorrt_llm may leave an idle CUDA context (~a few hundred MB), which +safely coexists with the running test. + +Mechanism (wired by ``tests/test_common/session_prefetcher_hooks.py``, +loaded from each test tree's top-level conftest): ``pytest_runtest_setup`` +lazily patches the library seams that construct ``MpiPoolSession`` for a +bare ``LLM(...)`` with a factory that (a) hands over the prefetched pool +when its size and spawn-time env/sys.path still match, and (b) re-arms a +spare pool of the same size for the next test. A miss falls back to the +normal synchronous spawn, so a wrong prefetch can only cost time, never +correctness. + +Weight page-cache warming: when the NEXT test's model differs from the +current one, its weight files are read in a background thread so the kernel +page cache is hot by the time that test loads weights. The next model is +discovered automatically from the accuracy-harness ``MODEL_PATH`` class +attribute or a ``model_folder``-style test parameter (modeling unit tests), +or declared explicitly with +``@pytest.mark.prefetch_model_dir("/path/to/model")``. This complements +pool prefetch (pool reuse does not cover model IO). Page cache is +reclaimable memory, so warming cannot OOM the host; a wasted warm (test +skipped or reordered) costs only IO bandwidth. + +Coexistence with MPI session reuse: when ``test_common/session_reuse.py`` is +wired and enabled it owns the same pool-creation seams and eliminates the +respawn outright, so the prefetcher automatically stays off the seams (see +``_reuse_layer_active``). Weight warming stays active, and reuse consumes +this layer's shadow pools on its cache misses (first pool of a size, +post-drain rebuild, post-retire replacement) via ``take``/``schedule_shadow`` +— the two layers compose: reuse covers the steady state, prefetch covers the +misses. + +Rollout: currently in a CANARY phase (reviewer request) — with no explicit +setting the plugin is active only on the stage groups in +``_CANARY_STAGE_PREFIXES`` (CI exports the stage name as ``stageName``). +``TRTLLM_TEST_PREFETCH_SESSION=1``/``0`` overrides in either direction and +remains the permanent kill switch. Suites that never import tensorrt_llm's +executor modules pay nothing — not even the tensorrt_llm import. +""" + +import glob +import os +import sys +import threading +import time +from collections import Counter +from concurrent.futures import ThreadPoolExecutor +from typing import NamedTuple + +# The spawn snapshot is shared with the session-reuse layer (both hand a +# live pool to a test that did not spawn it — same invariant: workers freeze +# the FULL env + sys.path at spawn). +from test_common._session_utils import _isinstance_transparent_shim, _spawn_snapshot + +# The only places in the library that construct MpiPoolSession for a bare +# LLM(...); tests passing their own _mpi_session never reach these lines. +# test_patch_targets_cover_all_library_construction_sites keeps this list +# honest against new construction sites appearing in the library. +_PATCH_TARGETS = ( + "tensorrt_llm.executor.proxy", + "tensorrt_llm.executor.rpc_proxy", + "tensorrt_llm.llmapi.llm", +) + +# Canary rollout stage GROUPS (one multi-GPU LLM-dense, one single-GPU +# LLM-dense). Prefixes, not exact names: the numbered shard suffix (-1/-2/…) +# is assigned by dynamic load balancing and cannot be targeted stably. +# Graduation plan: after a week of clean canary runs (zero failures +# attributed to the plugin, healthy session-summary counters) this gate is +# removed and the plugin returns to enabled-by-default. +_CANARY_STAGE_PREFIXES = ( + "DGX_H100-4_GPUs-PyTorch-DeepSeek", + "A10-PyTorch", +) + + +def _reuse_layer_active() -> bool: + """True when the MPI session-reuse layer owns the pool-creation seams. + + ``test_common.session_reuse`` keeps pools alive across tests at the SAME + seams this module would patch, and saves the whole respawn rather than + just hiding it — strictly better where it applies. When it is wired and + enabled, the prefetcher must stay off the seams so the two factories + don't fight over them (whoever patches first would silently disable the + other). Weight page-cache warming is orthogonal and stays on either way. + """ + mod = sys.modules.get("test_common.session_reuse") + if mod is None: + return False # not wired into this suite: seams are ours to patch + try: + return bool(mod.REUSE.is_active()) + except Exception: + return True # loaded but unreadable: err on staying out of the way + + +_READ_CHUNK = 64 << 20 # 64MB + + +def _weight_files(model_dir: str): + """The weight files the loader will actually read, in loader order. + + Mirrors HfWeightLoader.load_weights' selection: safetensors first — + minus "consolidated" copies, which the loader deliberately skips (they + duplicate the shards and can be enormous) — else *.bin, else *.pth. + Warming anything else is pure wasted IO. + """ + files = [ + f + for f in glob.glob(os.path.join(model_dir, "*.safetensors")) + if "consolidated" not in os.path.basename(f) + ] + for fallback in ("*.bin", "*.pth"): + if files: + break + files = glob.glob(os.path.join(model_dir, fallback)) + return sorted(files) + + +def _available_host_memory(): + """MemAvailable from /proc/meminfo in bytes, or None when unreadable.""" + try: + with open("/proc/meminfo") as fh: + for line in fh: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) * 1024 + except (OSError, ValueError, IndexError): + pass + return None + + +# Parametrized-test convention: the model lives in a parameter with one of +# these names (e.g. test_modeling_* files), holding either an absolute path +# or a directory name under LLM_MODELS_ROOT. +_MODEL_PARAM_NAMES = ("model_folder", "model_dir", "model_path") + + +def _models_root(): + """The models root the tests themselves resolve against, or None.""" + root = os.environ.get("LLM_MODELS_ROOT") + if root: + return root + try: # same fallback the test suites use (CI default scratch path) + from test_common.llm_data import llm_models_root + + root = llm_models_root() + return str(root) if root else None + except Exception: + return None + + +def _model_dir_of(item): + """A test item's model dir: marker, else class or parameter convention. + + Discovery order: the explicit ``prefetch_model_dir`` marker; the accuracy + harness's ``MODEL_PATH`` class attribute (120+ classes across + tests/integration/defs/accuracy); a ``model_folder``-style test parameter + (modeling unit tests), resolved under LLM_MODELS_ROOT unless absolute. + All of it is guess-tolerant: a value that is not a real directory of + weight files (e.g. an HF model id) makes ``warm_page_cache`` a silent + no-op, so a wrong guess costs nothing. + """ + marker = item.get_closest_marker("prefetch_model_dir") + if marker is not None and marker.args: + return marker.args[0] + model_path = getattr(getattr(item, "cls", None), "MODEL_PATH", None) + if isinstance(model_path, str): + return model_path + params = getattr(getattr(item, "callspec", None), "params", None) or {} + for name in _MODEL_PARAM_NAMES: + value = params.get(name) + if isinstance(value, str) and value: + if os.path.isabs(value): + return value + root = _models_root() + if root: + return os.path.join(root, value) + return None + + +def warm_page_cache(model_dir: str) -> float: + """Read ``model_dir``'s weight files to keep them in the OS page cache. + + The next LLM create then loads the weights from RAM, not disk. Pure file + IO — never touches CUDA, safe to run while another test owns the GPUs. + Returns the number of GiB read. + """ + files = _weight_files(model_dir) + if not files: + return 0.0 # not a local weight dir (e.g. an HF model id): nothing to warm + total_bytes = sum(os.stat(f).st_size for f in files) + available = _available_host_memory() + if available is not None and total_bytes > available: + # Larger than RAM: pages would be evicted before the test loads them — + # pure filer traffic with zero benefit (e.g. multi-hundred-GB models). + print( + f"[session-prefetch] skipping warm of {model_dir}: {total_bytes >> 30} GiB " + f"exceeds available host memory ({available >> 30} GiB)", + flush=True, + ) + return 0.0 + t0 = time.monotonic() + + def _read(path): + n = 0 + with open(path, "rb") as fh: + while True: + chunk = fh.read(_READ_CHUNK) + if not chunk: + return n + n += len(chunk) + + # thread_name_prefix keeps the IO workers inside the pytest.ini + # threadleak_exclude pattern (session-prefetch-\w+): a large warm can + # legitimately still be reading during the next test's threadleak check. + with ThreadPoolExecutor(max_workers=4, thread_name_prefix="session-prefetch-io") as ex: + total = sum(ex.map(_read, files)) + gib = total / (1 << 30) + print( + f"[session-prefetch] warmed page cache: {gib:.1f} GiB from " + f"{model_dir} in {time.monotonic() - t0:.1f}s", + flush=True, + ) + return gib + + +def _worker_import_report_cuda() -> bool: + """Import tensorrt_llm (the expensive part) and report CUDA state. + + Some library versions initialize a CUDA context at import time; that + idle context (~a few hundred MB, no kernels/allocations) is acceptable + and coexists with the running test, so it is reported, not asserted. + """ + import torch + + import tensorrt_llm # noqa: F401 + + return torch.cuda.is_initialized() + + +class _Built(NamedTuple): + """A finished background build: everything published (and consumed) together.""" + + spec: int + session: object + snapshot: object + + +class SessionPrefetcher: + def __init__(self): + self._lock = threading.Lock() + self._thread = None + self._building_spec = None # spec of the in-flight build, while _thread is set + self._build_gen = 0 # bumped when a pending build is abandoned + self._built = None # Optional[_Built], set only by _publish() + self._patched = set() + self._next_model = None # item -> next model dir; built lazily + self._warmed_dirs = set() + self._disposed = False + # Activity counters, reported once per session by dispose(). pytest + # captures per-test stdout (swallowing the per-event prints for + # passing tests), but pytest_sessionfinish runs OUTSIDE capture, so + # the summary is the one line guaranteed to reach the CI console. + self.stats = Counter() + self._warmed_gib = 0.0 + + @property + def enabled(self) -> bool: + # Under pytest-xdist every worker sees the FULL collection but runs a + # scheduler-assigned subset, and N workers would each hold a live + # pool plus a spare. Disable in xdist workers. + if os.environ.get("PYTEST_XDIST_WORKER"): + return False + explicit = os.environ.get("TRTLLM_TEST_PREFETCH_SESSION") + if explicit is not None: + return explicit.lower() in ("1", "true", "yes", "on") + # Canary rollout (reviewer request): with no explicit setting, enable + # only on the canary stage GROUPS. CI exports the stage name as + # ``stageName``; prefix matching covers dynamically numbered shards + # (per-shard targeting is impossible — the split is rebalanced as the + # test list changes). Everywhere else (including runs without a + # stageName) the plugin stays off until the canary graduates. + stage = os.environ.get("stageName", "") + return any(stage.startswith(p) for p in _CANARY_STAGE_PREFIXES) + + @staticmethod + def _next_model_map(items): + """Per item, the model dir of the NEXT test declaring one. + + One reverse pass, one lookup per item — O(n) once; ``on_test_setup`` + then costs a single dict lookup per test. (The naive alternative — + scanning the remaining collection at every test setup — is O(n^2) + lookups per session, seconds to minutes on large suites.) + """ + next_model, mapping = None, {} + for item in reversed(items): + mapping[item] = next_model + model = _model_dir_of(item) + if model is not None: + next_model = model + return mapping + + def on_test_setup(self, item) -> None: + """Warm the NEXT test's model weights while this test runs. + + Fires when the next model differs from the current one — including + between tests that share a pool (pool prefetch does not cover model + IO). The next-model map is built lazily on the FIRST test setup, from + ``session.items``: by then every reordering/deselecting plugin + (pytest-split runs trylast, --test-list filtering, -k/-m) has produced + the final run order, which a ``pytest_collection_modifyitems`` hook + could not guarantee. + """ + if self._next_model is None: + items = getattr(item.session, "items", None) or [item] + self._next_model = self._next_model_map(items) if self.enabled else {} + nxt = self._next_model.get(item) + if not nxt or nxt == _model_dir_of(item): + return + # Main-pytest-thread only (pytest_runtest_setup): no lock needed. + if nxt in self._warmed_dirs: + return # already warmed (or being warmed) this session + self._warmed_dirs.add(nxt) + threading.Thread( + target=self._warm, args=(nxt,), daemon=True, name="session-prefetch-warm" + ).start() + + def _warm(self, model_dir: str) -> None: + try: + gib = warm_page_cache(model_dir) + with self._lock: + if gib > 0: + self.stats["warms"] += 1 + self._warmed_gib += gib + else: + self.stats["warm_noops"] += 1 # no local weights / RAM guard + except Exception as e: # warming must never break the tests + print(f"[session-prefetch] page-cache warm failed (harmless): {e}", flush=True) + + def schedule_shadow(self, spec: int, env_overlay=None) -> None: + """Start building a spare ``spec``-worker pool in the background. + + Heuristic: the next test most likely needs a pool of the same size as + the current one. A miss is discarded at ``take()`` and the sync build + is no slower than without prefetch. + + ``env_overlay``: extra env vars to freeze into the WORKERS at spawn + (session_reuse restocks shadows with its worker-side weight cache + on). Passed through the library's worker-env channel, so the parent + process environment — and therefore the take()-time snapshot + comparison — is never touched. + """ + if not self.enabled or spec < 1: + return + with self._lock: + if self._thread is not None or (self._built is not None and self._built.spec == spec): + return # already building / built + self._building_spec = spec + self._thread = threading.Thread( + target=self._build, + args=(spec, self._build_gen, env_overlay), + daemon=True, + name="session-prefetch-build", + ) + self._thread.start() + + def _build(self, spec: int, gen: int, env_overlay=None) -> None: + try: + from tensorrt_llm._utils import mpi_disabled + + if mpi_disabled(): + return + from tensorrt_llm.llmapi.mpi_session import MpiPoolSession + + snapshot = _spawn_snapshot() # workers freeze env+sys.path at spawn + # wait_shutdown: see _make_factory — every pool this layer hands + # out blocks its shutdown on actual worker exit. + session = MpiPoolSession(n_workers=spec, wait_shutdown=True, env_overrides=env_overlay) + if any(session.submit_sync(_worker_import_report_cuda)): + print( + "[session-prefetch] note: tensorrt_llm import initialized an idle " + "CUDA context in the prefetched workers (library version behavior)", + flush=True, + ) + self._publish(spec, session, snapshot, gen) + except Exception as e: # prefetch must never break the tests + print( + f"[session-prefetch] background build failed (falling back to synchronous): {e}", + flush=True, + ) + + def _publish(self, spec, session, snapshot, gen: int) -> None: + """Publish a finished background build, unless it was abandoned. + + ``_drain()`` bumps ``_build_gen`` when a build outlives its join + timeout; such a late build must shut its pool down instead of + publishing (a late publish would overwrite — and leak — a newer pool, + or hand a stale pool to a future test). The empty-slot check likewise + prevents overwriting an unconsumed pool. + """ + with self._lock: + if gen == self._build_gen and self._built is None: + self._built = _Built(spec, session, snapshot) + self.stats["pools_built"] += 1 + return + self.stats["pools_discarded_superseded"] += 1 + print("[session-prefetch] discarding superseded background build", flush=True) + session.shutdown() + + def _drain(self, timeout: float): + """Join a pending build (abandoning it on timeout) and pop the slot.""" + # Read _thread under the lock: schedule_shadow() assigns-then-starts + # inside its critical section, and an unlocked read here can observe + # the assigned-but-not-yet-started thread ("cannot join thread before + # it is started" when a test creates LLMs concurrently). + with self._lock: + thread = self._thread + if thread is not None: + thread.join(timeout=timeout) + with self._lock: + if thread is not None and thread.is_alive(): + # Abandon the overdue build: bump the generation so its late + # _publish() shuts the pool down instead of landing. + self._build_gen += 1 + self._thread = None + built, self._built = self._built, None + return built + + def take(self, spec: int): + """Return a prefetched session for ``spec``, or None to build sync.""" + if not self.enabled: + return None + with self._lock: + wrong_size_in_flight = ( + self._thread is not None and self._built is None and self._building_spec != spec + ) + if wrong_size_in_flight: + # Joining would stall this caller for most of a spawn only to + # discard the mismatched result — slower than no prefetch at all. + # Fall back to the synchronous spawn now and leave the build to + # land for a later take of its own size. + self.stats["pools_skipped_size_in_flight"] += 1 + return None + # Slowest legitimate build measured is ~117s (busy node); 180s gives + # 1.5x margin. On a genuine hang we give up and fall back to a + # synchronous build instead of stalling the suite. + built = self._drain(timeout=180) + if built is None: + return None + if built.spec == spec and built.snapshot == _spawn_snapshot(): + # An instant handover is safe against the previous worker's GPU + # memory: every pool these layers hand out is built with + # wait_shutdown=True, so its shutdown() blocked until the workers + # actually exited (and released their memory). + self.stats["pools_handed_over"] += 1 + print(f"[session-prefetch] handing over prefetched {spec}-worker pool", flush=True) + return built.session + # Spec/env/sys.path mismatch (test skipped, reordered, or changed + # state the frozen workers would not see): discard. + self.stats["pools_discarded_stale"] += 1 + threading.Thread( + target=built.session.shutdown, daemon=True, name="session-prefetch-discard" + ).start() + return None + + def _make_factory(self, real_cls): + """A drop-in for ``MpiPoolSession`` that consumes and re-arms the shadow.""" + + def factory(n_workers, *args, **kwargs): + if args or kwargs: + return real_cls(n_workers, *args, **kwargs) + # n_workers == 1 included: the default single-GPU path also spawns + # a 1-worker pool (executor.py -> proxy.py) costing ~50s of + # spawn+import, the same as multi-GPU pools. + # wait_shutdown: this pool's shutdown must not return until its + # workers exited (and released GPU memory) — the NEXT pool is + # handed over instantly, without the ~50s sync spawn that used to + # hide the release window. Such spawns fail closed when identity + # collection cannot complete; a prefetch layer must not turn that + # into a test failure, so retry once and then degrade LOUDLY to a + # plain pool (pre-prefetch semantics: shutdown returns at + # disconnect). + session = self.take(n_workers) + if session is None: + try: + session = real_cls(n_workers=n_workers, wait_shutdown=True) + except Exception as e: + print(f"[session-prefetch] pool spawn failed, retrying once: {e}", flush=True) + try: + session = real_cls(n_workers=n_workers, wait_shutdown=True) + except Exception as e2: + print( + "[session-prefetch] wait_shutdown spawn failed twice: " + f"{e2}; falling back to a plain pool (handover " + "protection degraded for the NEXT pool on this node)", + flush=True, + ) + self.stats["pools_spawned_degraded"] += 1 + session = real_cls(n_workers=n_workers) + self.schedule_shadow(n_workers) # re-arm for the NEXT test + return session + + return factory + + def install_pool_factory_if_loaded(self) -> None: + """Lazily patch the pool-creation seams for zero-test-change prefetch. + + Only patches target modules ALREADY imported by the test suite, so + suites that never touch tensorrt_llm pay nothing (not even the + import). Idempotent — called from ``pytest_runtest_setup``. + Only the ``mpi_session is None`` branches construct ``MpiPoolSession`` + directly, so tests passing their own session (shared/grouped pools) + are never intercepted. + """ + if not self.enabled: + return + if len(self._patched) == len(_PATCH_TARGETS): + return # everything already patched: per-test fast path + if _reuse_layer_active(): + # session_reuse owns the seams: skip MPI-pool prefetch entirely + # (reuse eliminates the respawn; prefetch could only hide it). + self.stats["mpi_yielded_to_reuse"] = 1 + return + pending = [n for n in _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 + + factory = self._make_factory(real_cls) + for name in pending: + mod = sys.modules[name] + if getattr(mod, "MpiPoolSession", None) is real_cls: + # A real type, not a bare function: library code may run + # isinstance(x, MpiPoolSession) against the patched attribute + # (proxy.py's killed-worker detection did) — a function there + # raises TypeError and kills every LLM creation. + mod.MpiPoolSession = _isinstance_transparent_shim(real_cls, factory) + self._patched.add(name) + + def dispose(self) -> None: + """Shut down any unconsumed shadow pool (end-of-session cleanup). + + 60s (vs take()'s 180s): at session end there is no test left to hand + the pool to, so a still-running build is only worth a short grace + before it is abandoned to its generation-bump cleanup. Idempotent: a + repository-root run dispatches sessionfinish from both the repo-root + and the subtree conftest. + """ + if self._disposed: + return + self._disposed = True + built = self._drain(timeout=60) + if built is not None: + built.session.shutdown() + # One line per session, emitted OUTSIDE pytest's per-test capture + # (pytest_sessionfinish) so it reaches the CI console: the per-event + # prints above are swallowed for passing tests. Silent when the + # prefetcher never did anything (non-LLM suites). + if self.stats: + parts = ", ".join(f"{k}={v}" for k, v in sorted(self.stats.items())) + if self._warmed_gib: + parts += f", warmed_gib={self._warmed_gib:.1f}" + print(f"[session-prefetch] session summary: {parts}", flush=True) + + +PREFETCHER = SessionPrefetcher() diff --git a/tests/test_common/session_prefetcher_hooks.py b/tests/test_common/session_prefetcher_hooks.py new file mode 100644 index 000000000000..26c5d7201472 --- /dev/null +++ b/tests/test_common/session_prefetcher_hooks.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pytest plugin wiring for the session prefetcher (repo-wide, demand-driven). + +Loaded via ``pytest_plugins`` in each test tree's top-level conftest. 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. +""" + +import os + +from test_common.session_prefetcher import PREFETCHER + + +def pytest_configure(config): + config.addinivalue_line( + "markers", + "prefetch_model_dir(path): model dir the test loads, for page-cache warming", + ) + + +def pytest_runtest_setup(item): + # Last-line fail-open: prefetch is an optimization wired into EVERY test's + # setup, so an unexpected error here must degrade to baseline speed via + # the kill switch — never error the suite. + try: + PREFETCHER.install_pool_factory_if_loaded() + PREFETCHER.on_test_setup(item) + except Exception as e: + os.environ["TRTLLM_TEST_PREFETCH_SESSION"] = "0" + print(f"[session-prefetch] disabled by unexpected error: {e}", flush=True) + + +def pytest_sessionfinish(session, exitstatus): + try: + PREFETCHER.dispose() + except Exception as e: # never fail the session over cleanup + print(f"[session-prefetch] dispose failed: {e}", flush=True) diff --git a/tests/test_common/session_reuse.py b/tests/test_common/session_reuse.py index 145ee5cf2248..ebe9eac2ff4d 100644 --- a/tests/test_common/session_reuse.py +++ b/tests/test_common/session_reuse.py @@ -23,10 +23,15 @@ 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. +per worker, barrier-pinned: ``grouped_test_utils.submit_sync_per_worker``). +Handover cannot race the previous worker's GPU-memory release: every pool +these layers build is constructed with ``wait_shutdown=True``, so its +shutdown blocks until the workers actually exited. + +Cache misses (first pool of a size, post-drain rebuild, post-retire +replacement) take a shadow pool pre-spawned by the session-prefetch layer +when it is wired (``_prefetcher``), hiding the ~50s spawn; each miss restocks +one shadow for the next. Enable/disable with ``TRTLLM_TEST_REUSE_SESSION`` (default on; ``0`` disables). Disabled under pytest-xdist workers (parallel tests would multiply live pools). @@ -35,8 +40,10 @@ import os import sys import threading -import time +# The spawn snapshot is shared with the session-prefetch layer (both hand a +# live pool to a test that did not spawn it — same invariant). +from test_common._session_utils import _isinstance_transparent_shim, _spawn_snapshot 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 @@ -59,148 +66,45 @@ "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)) +_RETIRE_THREADS: list = [] +_RETIRE_LOCK = threading.Lock() -def wait_gpu_memory_settle() -> None: - """Wait until visible GPUs are mostly free or free memory stops rising. +def _reap_retires(timeout: float = 60.0) -> None: + """Join in-flight retire threads (bounded); no-op when none are running. - Never raises: on any NVML problem the handover proceeds as before. + A retired pool's workers hold their (full-model) GPU memory until they + exit; the retire thread blocks on that exit (``wait_shutdown=True``), but + it is a BACKGROUND thread — the test hot path never waits on it. Before + an instant cached-pool handover, joining in-flight retires is what makes + the handover safe against a corpse still releasing (e.g. the duplicate + retired by ``_release`` moments earlier); every other path spawns fresh + (~50s), which outlasts the release naturally. Also called at drain + rendezvous points so disposals cannot leak past the session. """ - 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: + with _RETIRE_LOCK: + in_flight, _RETIRE_THREADS[:] = list(_RETIRE_THREADS), [] + for t in in_flight: + t.join(timeout=timeout) + if t.is_alive(): print( - f"[session-reuse] waited {waited:.1f}s before handover for GPU memory release", + "[session-reuse] WARNING: pool retirement did not finish within 60s", flush=True, ) - except Exception: - pass - finally: - try: - pynvml.nvmlShutdown() - except Exception: - pass - -_RETIRE_THREADS: list = [] -_RETIRE_LOCK = threading.Lock() +def _prefetcher(): + """The session-prefetch singleton when that layer is wired, else None. -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. + Mirror of the prefetcher's own reuse probe: coordination goes through + sys.modules so neither layer imports the other at module load (a suite + wired with only one layer pays nothing for the other). The prefetcher + yields the pool SEAMS to reuse; reuse in turn consumes prefetched + shadows on its cache misses — the two layers compose, not compete. """ - try: - return tuple(sorted(submit_sync_per_worker(real, _get_worker_pid))) - except Exception: - return () + mod = sys.modules.get("test_common.session_prefetcher") + return getattr(mod, "PREFETCHER", None) def _describe_mismatch(spawn_snap, now_snap, uses, max_uses): @@ -273,6 +177,17 @@ def enabled(self) -> bool: "on", ) + def is_active(self) -> bool: + """Public probe for sibling layers: does reuse own the pool seams? + + The session prefetcher yields the ``MpiPoolSession`` seams when this + returns True (reuse eliminates the respawn outright; prefetch could + only hide it). Deliberately ignores ``_suspended``: a per-test + cache bypass (``private_mpi_session``) does not change seam + ownership. + """ + return self.enabled + @property def max_uses(self) -> int: return int(os.environ.get("TRTLLM_TEST_REUSE_MAX_USES", "16")) @@ -299,10 +214,13 @@ def _retire(real, broken: bool = False): def _dispose(): import signal + # Lazy: only runs when a pool exists, so tensorrt_llm is loaded. + from tensorrt_llm.llmapi.mpi_session import _process_start_time + 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: + if start_time is None or _process_start_time(pid) != start_time: continue try: os.kill(pid, signal.SIGKILL) @@ -356,23 +274,32 @@ 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) + if args or kwargs: + return real_cls(n_workers, *args, **kwargs) + # wait_shutdown: the private pool dies at LLM shutdown; block + # there until its workers exited so the next pool (often handed + # over instantly from the cache) cannot race the GPU release. + return real_cls(n_workers=n_workers, wait_shutdown=True) 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 + mod.MpiPoolSession = _isinstance_transparent_shim( + real_cls, 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.""" + """Hand out a cached same-size pool (workers reset) 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) + # owns and destroys normally (wait_shutdown: its shutdown blocks + # until the workers exited, so the next handover cannot race the + # GPU-memory release). + return real_cls(n_workers=n_workers, wait_shutdown=True) with self._lock: real = self._pools.pop(n_workers, None) if real is not None: @@ -391,8 +318,15 @@ def acquire(self, real_cls, n_workers): self._retire(real) # stale worker state or lifetime cap else: try: + # An instant handover must not race a corpse still + # releasing its GPU memory. Retire threads block on the + # workers' exit (wait_shutdown=True) but run in the + # BACKGROUND, so join any in flight (a duplicate retired + # by _release moments ago held full model memory). No-op + # on the common path; every non-cached path spawns fresh + # (~50s), which outlasts the release naturally. + _reap_retires() 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})", @@ -408,27 +342,55 @@ def acquire(self, real_cls, n_workers): 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. + """Obtain a cache-managed pool: prefetched if one is armed, else spawn. + + Every cache miss lands here (first pool of a size, post-drain + rebuild, post-retire replacement). When the session-prefetch layer is + wired, a shadow pool armed at the PREVIOUS miss is taken instantly — + hiding the ~50s spawn the miss would otherwise pay — and a + replacement shadow is armed for the next miss of this size. Without + the prefetch layer (or on a shadow miss) the synchronous spawn is + unchanged. + + The worker-side HF weight cache env is frozen into the workers via + the library's ``env_overrides`` channel (parent env untouched, so + acquire-time snapshot comparisons still match); an explicit user + setting of either var is respected and left untouched. Prefetched + shadows were armed with the same overlay. wait_shutdown: shutdown of + this pool blocks until its workers exited, so a successor cannot + race the GPU-memory release. """ 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) + overrides = {k: v for k, v in _WEIGHT_CACHE_ENV.items() if k not in os.environ} + real = None + prefetcher = _prefetcher() + if prefetcher is not None: + try: + real = prefetcher.take(n_workers) + except Exception as e: # prefetch is an optimization: fall back + print(f"[session-reuse] prefetched-pool take failed: {e}", flush=True) + real = None + try: + # Restock ONE shadow for the next miss of this size (no-op if + # one is already armed/building, or prefetch is disabled). + prefetcher.schedule_shadow(n_workers, env_overlay=overrides) + except Exception: + pass + if real is None: + try: + real = real_cls(n_workers=n_workers, wait_shutdown=True, env_overrides=overrides) + except Exception as e: + # wait_shutdown spawns fail closed (identity collection must + # complete); a transient slow node deserves one loud retry — + # a second failure means the node is genuinely broken. + print(f"[session-reuse] pool spawn failed, retrying once: {e}", flush=True) + real = real_cls(n_workers=n_workers, wait_shutdown=True, env_overrides=overrides) real._reuse_uses = 0 real._reuse_spawn_snapshot = snapshot - real._reuse_worker_pids = _collect_worker_pids(real) + # (pid, start_time) per worker, recorded by the library at spawn + # (wait_shutdown=True above). _retire uses them to SIGKILL wedged + # workers; best effort — empty means graceful shutdown only. + real._reuse_worker_pids = getattr(real, "_worker_identities", ()) return real def _release(self, real): @@ -457,15 +419,7 @@ def drain(self) -> None: 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, - ) + _reap_retires() with self._lock: pools, self._pools = list(self._pools.values()), {} if not pools: diff --git a/tests/unittest/auto_deploy/multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py b/tests/unittest/auto_deploy/multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py index 139d76cd5a97..a08d633a5381 100644 --- a/tests/unittest/auto_deploy/multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py +++ b/tests/unittest/auto_deploy/multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py @@ -476,7 +476,9 @@ def _run_with_retries(worker_fn, world_size, **kwargs): max_retries = 5 last_exc = None for _ in range(max_retries): - pool = MpiPoolSession(n_workers=world_size) + # wait_shutdown: block shutdown until the workers exited, so a test + # handed a live pool right after this one cannot race the GPU release. + pool = MpiPoolSession(n_workers=world_size, wait_shutdown=True) try: return pool.submit_sync(worker_fn, port=None, world_size=world_size, **kwargs) except DistNetworkError as e: diff --git a/tests/unittest/auto_deploy/multigpu/transformations/library/test_allreduce_residual_rmsnorm_fusion.py b/tests/unittest/auto_deploy/multigpu/transformations/library/test_allreduce_residual_rmsnorm_fusion.py index f9602d84b8e9..94a2f47c649f 100644 --- a/tests/unittest/auto_deploy/multigpu/transformations/library/test_allreduce_residual_rmsnorm_fusion.py +++ b/tests/unittest/auto_deploy/multigpu/transformations/library/test_allreduce_residual_rmsnorm_fusion.py @@ -203,7 +203,9 @@ def test_allreduce_fusion(device_count, ModuleCls, strategy, rmsnorm_op): max_retries = 5 last_exc: Exception | None = None for _ in range(max_retries): - mpi_pool = MpiPoolSession(n_workers=n_workers) + # wait_shutdown: block shutdown until the workers exited, so a test + # handed a live pool right after this one cannot race the GPU release. + mpi_pool = MpiPoolSession(n_workers=n_workers, wait_shutdown=True) try: mpi_pool.submit_sync( _test_allreduce_fusion, diff --git a/tests/unittest/conftest.py b/tests/unittest/conftest.py index e1d5943a0a57..9ceca2315958 100644 --- a/tests/unittest/conftest.py +++ b/tests/unittest/conftest.py @@ -39,7 +39,15 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from integration.defs import test_list_parser +# Dispatched explicitly (not via pytest_plugins, which pytest forbids in a +# non-top-level conftest: a repo-root invocation like `pytest tests` loads +# this file as a NESTED conftest and would fail collection; and not via "-p" +# in pytest.ini addopts, which imports at preparse, before the ini pythonpath +# entries are usable). The wrappers below forward to the plugin; hooks are +# idempotent, so a repo-root run that also dispatches from tests/conftest.py +# is harmless. from test_common import s3_output +from test_common import session_prefetcher_hooks as _prefetch_hooks def dump_threads(signum, frame): @@ -47,6 +55,7 @@ def dump_threads(signum, frame): def pytest_configure(config): + _prefetch_hooks.pytest_configure(config) os.environ.setdefault("TRTLLM_NO_USAGE_STATS", "1") # avoid thread leak of tqdm's TMonitor @@ -510,3 +519,11 @@ def setup_ray_cluster() -> Generator[int, None, None]: finally: if ray.is_initialized(): ray.shutdown() + + +def pytest_runtest_setup(item): + _prefetch_hooks.pytest_runtest_setup(item) + + +def pytest_sessionfinish(session, exitstatus): + _prefetch_hooks.pytest_sessionfinish(session, exitstatus) diff --git a/tests/unittest/executor/test_base_worker.py b/tests/unittest/executor/test_base_worker.py index 3c5b54b96569..a60922edb72f 100644 --- a/tests/unittest/executor/test_base_worker.py +++ b/tests/unittest/executor/test_base_worker.py @@ -176,7 +176,9 @@ def setup_method(self): self.session = self.create_worker_session() def create_worker_session(self): - session = MpiPoolSession(n_workers=2) + # wait_shutdown: block shutdown until the workers exited, so a test + # handed a live pool right after this one cannot race the GPU release. + session = MpiPoolSession(n_workers=2, wait_shutdown=True) return session @pytest.mark.gpu2 diff --git a/tests/unittest/llmapi/test_mpi_session.py b/tests/unittest/llmapi/test_mpi_session.py index 781947fe65bf..98a92e2712b9 100644 --- a/tests/unittest/llmapi/test_mpi_session.py +++ b/tests/unittest/llmapi/test_mpi_session.py @@ -166,3 +166,112 @@ def read_stream(stream, output_stream): if return_code != 0: raise subprocess.CalledProcessError(return_code, command) + + +# ---- wait_shutdown: shutdown blocks until worker processes actually exit ---- + + +def _wait_workers_exit(identities, timeout: float) -> None: + """Call the unbound method on an inert stand-in (no MPI spawn). + + ``_wait_workers_exit`` only reads ``self._worker_identities``; a real + ``MpiPoolSession`` shell would trigger the base class's abort machinery + at garbage collection. + """ + import types + + stand_in = types.SimpleNamespace(_worker_identities=identities) + MpiPoolSession._wait_workers_exit(stand_in, timeout=timeout) + + +def test_process_start_time_live_and_gone(): + from tensorrt_llm.llmapi.mpi_session import _process_start_time + + assert _process_start_time(os.getpid()) is not None + child = Popen(["true"]) # nosec B603, B607 + child.wait() + assert _process_start_time(child.pid) is None # reaped: /proc entry gone + + +def test_wait_workers_exit_returns_once_workers_are_gone(): + from tensorrt_llm.llmapi.mpi_session import _process_start_time + + child = Popen(["true"]) # nosec B603, B607 + identity = (child.pid, _process_start_time(child.pid)) + child.wait() + # Dead worker -> returns immediately; a None start_time is skipped + # (identity collection failed for that worker: nothing to wait on). + _wait_workers_exit((identity, (os.getpid(), None)), timeout=5.0) + + +def test_wait_workers_exit_bounded_by_timeout_on_live_worker(): + import time as _time + + from tensorrt_llm.llmapi.mpi_session import _process_start_time + + me = (os.getpid(), _process_start_time(os.getpid())) + t0 = _time.monotonic() + _wait_workers_exit((me, ), timeout=0.2) # this process will not exit + waited = _time.monotonic() - t0 + assert 0.2 <= waited < 2.0 # bounded: a wedged worker cannot hang teardown + + +def _collect_identities(monkeypatch, results, pending=0, n_workers=2): + """Drive _collect_worker_identities on an inert stand-in (no MPI spawn).""" + import types + from concurrent.futures import Future + + futs = [] + for r in results: + f = Future() + f.set_result(r) + futs.append(f) + never = [Future() for _ in range(pending)] # never resolve + + from tensorrt_llm.llmapi import mpi_session as m + + monkeypatch.setattr(m, "futures_wait", lambda fs, timeout: (futs, never)) + killed = [] + monkeypatch.setattr(os, "kill", lambda pid, sig: killed.append(pid)) + it = iter(futs + never) + stand_in = types.SimpleNamespace( + n_workers=n_workers, + mpi_pool=types.SimpleNamespace(submit=lambda fn: next(it), + shutdown=lambda wait=True: None), + _teardown_unidentified_pool=lambda ids: MpiPoolSession. + _teardown_unidentified_pool(stand_in, ids), + ) + result = MpiPoolSession._collect_worker_identities(stand_in) + return result, killed + + +def test_identity_collection_complete_returns_identities(monkeypatch): + from tensorrt_llm.llmapi.mpi_session import _process_start_time + + me = (os.getpid(), _process_start_time(os.getpid())) + other = (1, b"1") # pid 1: exists but start_time won't match -> unique pid + ids, killed = _collect_identities(monkeypatch, [me, other]) + assert set(ids) == {me, other} and not killed + + +def test_identity_collection_fails_closed_on_timeout(monkeypatch): + # A pending barrier task means the pool cannot honor wait_shutdown: + # the session must be torn down and rejected, NOT handed out with the + # contract silently downgraded (review requirement). + import pytest as _pytest + + from tensorrt_llm.llmapi.mpi_session import _process_start_time + + me = (os.getpid(), _process_start_time(os.getpid())) + with _pytest.raises(RuntimeError, match="incomplete"): + _collect_identities(monkeypatch, [me], pending=1) + + +def test_identity_collection_fails_closed_on_duplicate_pids(monkeypatch): + import pytest as _pytest + + from tensorrt_llm.llmapi.mpi_session import _process_start_time + + me = (os.getpid(), _process_start_time(os.getpid())) + with _pytest.raises(RuntimeError, match="incomplete"): + _collect_identities(monkeypatch, [me, me]) # one worker answered twice diff --git a/tests/unittest/llmapi/test_session_prefetcher.py b/tests/unittest/llmapi/test_session_prefetcher.py new file mode 100644 index 000000000000..2df7dfe5d50c --- /dev/null +++ b/tests/unittest/llmapi/test_session_prefetcher.py @@ -0,0 +1,533 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pure-logic tests for the session prefetcher — no MPI, no GPU.""" + +import re +import sys +import threading +import time +import types +from pathlib import Path + +import pytest +from test_common import session_prefetcher +from test_common.session_prefetcher import SessionPrefetcher, warm_page_cache + + +class _FakeMarker: + def __init__(self, *args): + self.args = args + + +class _FakeItem: + def __init__(self, model_dir=None, cls=None, params=None): + self._marker = _FakeMarker(model_dir) if model_dir else None + self.cls = cls + if params is not None: + self.callspec = types.SimpleNamespace(params=params) + + def get_closest_marker(self, name): + return self._marker if name == "prefetch_model_dir" else None + + +def _as_session(*items): + """Link fake items into a fake pytest session (final run order).""" + session = types.SimpleNamespace(items=list(items)) + for item in items: + item.session = session + return list(items) + + +def _wait_for(cond, timeout=5.0): + """Poll ``cond`` (fire-and-forget background work) until true or timeout.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if cond(): + return True + time.sleep(0.05) + return cond() + + +@pytest.fixture +def prefetcher(monkeypatch): + monkeypatch.setenv("TRTLLM_TEST_PREFETCH_SESSION", "1") + # These pure-logic tests may themselves run under xdist; pin the worker + # marker off so the prefetcher's xdist guard does not disable it here. + monkeypatch.delenv("PYTEST_XDIST_WORKER", raising=False) + p = SessionPrefetcher() + # Record build/warm triggers instead of spawning MPI pools/reading weights. + built, warmed, overlays = [], [], [] + + def _fake_build(self, spec, gen, env_overlay=None): + built.append(spec) + overlays.append(env_overlay) + + monkeypatch.setattr(SessionPrefetcher, "_build", _fake_build) + monkeypatch.setattr(SessionPrefetcher, "_warm", lambda self, d: warmed.append(d)) + p.built, p.warmed, p.overlays = built, warmed, overlays + return p + + +class _FakePool: + def __init__(self, n_workers, wait_shutdown=False): + self.n_workers = n_workers + self.wait_shutdown = wait_shutdown + self.shut = False + + def shutdown(self): + self.shut = True + + +def _arm(prefetcher, pool, spec=4): + """Publish ``pool`` into the shadow slot through the real API.""" + prefetcher._publish(spec, pool, session_prefetcher._spawn_snapshot(), prefetcher._build_gen) + + +def test_canary_enabled_on_canary_stage_groups(monkeypatch): + # Canary phase: with no explicit setting, active only on the canary + # stage groups; prefix match covers dynamically numbered shards. + monkeypatch.delenv("TRTLLM_TEST_PREFETCH_SESSION", raising=False) + monkeypatch.delenv("PYTEST_XDIST_WORKER", raising=False) + monkeypatch.setenv("stageName", "A10-PyTorch-2") + assert SessionPrefetcher().enabled + monkeypatch.setenv("stageName", "DGX_H100-4_GPUs-PyTorch-DeepSeek-1") + assert SessionPrefetcher().enabled + + +def test_canary_disabled_elsewhere(monkeypatch): + monkeypatch.delenv("TRTLLM_TEST_PREFETCH_SESSION", raising=False) + monkeypatch.delenv("PYTEST_XDIST_WORKER", raising=False) + monkeypatch.setenv("stageName", "DGX_B200-PyTorch-4") + assert not SessionPrefetcher().enabled + monkeypatch.delenv("stageName", raising=False) # local run, no stage + assert not SessionPrefetcher().enabled + + +def test_explicit_env_overrides_canary_gate(monkeypatch): + monkeypatch.delenv("PYTEST_XDIST_WORKER", raising=False) + monkeypatch.setenv("stageName", "DGX_B200-PyTorch-4") # not a canary stage + monkeypatch.setenv("TRTLLM_TEST_PREFETCH_SESSION", "1") + assert SessionPrefetcher().enabled # manual opt-in anywhere + monkeypatch.setenv("stageName", "A10-PyTorch-1") # canary stage + monkeypatch.setenv("TRTLLM_TEST_PREFETCH_SESSION", "0") + assert not SessionPrefetcher().enabled # kill switch beats the canary + + +def test_disabled_is_noop(monkeypatch): + monkeypatch.setenv("TRTLLM_TEST_PREFETCH_SESSION", "0") + p = SessionPrefetcher() + p.schedule_shadow(2) + assert p._thread is None + assert p.take(2) is None + + +def test_disabled_in_xdist_worker(monkeypatch): + # Under xdist each worker runs a scheduler-assigned subset; N workers + # would each hold a live pool plus a spare. + monkeypatch.setenv("TRTLLM_TEST_PREFETCH_SESSION", "1") + monkeypatch.setenv("PYTEST_XDIST_WORKER", "gw0") + assert not SessionPrefetcher().enabled + + +def test_factory_miss_builds_sync_and_arms_shadow(prefetcher): + # _build is stubbed by the fixture to record specs instead of spawning. + factory = prefetcher._make_factory(_FakePool) + session = factory(4) # nothing prefetched yet -> sync build + assert isinstance(session, _FakePool) and session.n_workers == 4 + # Every pool this layer hands out must block its shutdown on worker + # exit: the NEXT pool is handed over instantly, without the ~50s sync + # spawn that used to hide the GPU-memory release window. + assert session.wait_shutdown + prefetcher._thread.join(timeout=10) + assert prefetcher.built == [4] # shadow armed for the next test + + +def test_factory_single_worker_also_prefetches(prefetcher): + # The default single-GPU path spawns a 1-worker pool too (executor.py -> + # proxy.py), paying the same ~50s spawn+import: it must benefit as well. + factory = prefetcher._make_factory(_FakePool) + session = factory(1) + assert isinstance(session, _FakePool) + prefetcher._thread.join(timeout=10) + assert prefetcher.built == [1] # shadow armed for the next 1-GPU test + + +def test_schedule_shadow_passes_env_overlay_to_build(prefetcher): + # session_reuse restocks shadows with its worker-side weight-cache env; + # the overlay must reach the build (frozen into workers via the library's + # env_overrides channel — the parent env is never touched). + prefetcher.schedule_shadow(2, env_overlay={"TRTLLM_HF_WEIGHT_CACHE": "1"}) + prefetcher._thread.join(timeout=10) + assert prefetcher.built == [2] + assert prefetcher.overlays == [{"TRTLLM_HF_WEIGHT_CACHE": "1"}] + + +def test_take_wrong_size_in_flight_does_not_wait(prefetcher, monkeypatch): + # A miss must not stall behind an in-flight build of ANOTHER size only to + # discard the result — that is slower than no prefetch at all (wait ~one + # spawn, then spawn again). It falls back to sync immediately, and the + # build still lands for a later take of its own size. + release = threading.Event() + + def _slow_build(self, spec, gen, env_overlay=None): + release.wait(5) + self._publish(spec, _FakePool(spec), session_prefetcher._spawn_snapshot(), gen) + + monkeypatch.setattr(SessionPrefetcher, "_build", _slow_build) + prefetcher.schedule_shadow(2) + t0 = time.monotonic() + assert prefetcher.take(4) is None # wrong size in flight: no join + assert time.monotonic() - t0 < 1.0 + assert prefetcher.stats["pools_skipped_size_in_flight"] == 1 + release.set() + prefetcher._thread.join(timeout=10) + taken = prefetcher.take(2) # the undisturbed build landed for its size + assert isinstance(taken, _FakePool) and taken.n_workers == 2 + + +def test_factory_degrades_loudly_when_wait_shutdown_spawn_fails(prefetcher): + # The library fails closed when identity collection cannot complete; the + # prefetch layer must not turn that into a test failure: retry once, + # then degrade LOUDLY to a plain pool (pre-prefetch semantics). + calls = [] + + class _FailingWaitPool(_FakePool): + def __init__(self, n_workers, wait_shutdown=False): + calls.append(wait_shutdown) + if wait_shutdown: + raise RuntimeError("identity collection incomplete") + super().__init__(n_workers, wait_shutdown) + + factory = prefetcher._make_factory(_FailingWaitPool) + session = factory(2) + assert isinstance(session, _FailingWaitPool) and not session.wait_shutdown + assert calls == [True, True, False] # two contract attempts, then plain + assert prefetcher.stats["pools_spawned_degraded"] == 1 + + +def test_factory_hit_hands_over_shadow(prefetcher): + pool = _FakePool(4) + _arm(prefetcher, pool, spec=4) + factory = prefetcher._make_factory(_FakePool) + assert factory(4) is pool # prefetched pool handed over + + +def test_take_spec_mismatch_returns_none(prefetcher): + pool = _FakePool(4) + _arm(prefetcher, pool, spec=4) + assert prefetcher.take(2) is None # wrong size: sync fallback + + +def test_take_discards_on_env_mismatch(prefetcher, monkeypatch): + pool = _FakePool(4) + _arm(prefetcher, pool, spec=4) + monkeypatch.setenv("TLLM_TEST_ONLY_FLAG", "changed-after-spawn") + assert prefetcher.take(4) is None # frozen workers would miss the new env + assert _wait_for(lambda: pool.shut) # stale shadow torn down in background + + +def test_take_discards_on_nonprefixed_env_mismatch(prefetcher, monkeypatch): + # Workers inherit the WHOLE parent env at spawn; test knobs outside any + # TRTLLM*/TLLM* prefix (e.g. OVERRIDE_QUANT_ALGO, read inside workers by + # model_config.py) must also invalidate a prefetched pool, else workers + # silently run with stale env (review finding on the prefix allowlist). + pool = _FakePool(4) + _arm(prefetcher, pool, spec=4) + monkeypatch.setenv("OVERRIDE_QUANT_ALGO", "W4A16_MXFP4") + assert prefetcher.take(4) is None + + +def test_pytest_current_test_drift_does_not_discard(prefetcher, monkeypatch): + # PYTEST_CURRENT_TEST changes every test phase by design; it must not + # invalidate the snapshot or no prefetched pool would ever be handed over. + monkeypatch.setenv("PYTEST_CURRENT_TEST", "test_a (call)") + pool = _FakePool(4) + _arm(prefetcher, pool, spec=4) + monkeypatch.setenv("PYTEST_CURRENT_TEST", "test_b (setup)") + assert prefetcher.take(4) is pool + + +def test_take_discards_on_syspath_mismatch(prefetcher, monkeypatch): + # test_modeling_out_of_tree monkeypatches sys.path before + # LLM(); pool workers freeze sys.path at spawn (MPIPoolExecutor(path=...)), + # so a pool spawned earlier can't import the out-of-tree module and dies + # during initialization. sys.path must be part of the handover guard. + pool = _FakePool(4) + _arm(prefetcher, pool, spec=4) + monkeypatch.syspath_prepend("/oot/example/path") + assert prefetcher.take(4) is None # frozen workers would miss the new path + + +def test_take_does_not_join_unstarted_shadow_thread(prefetcher, monkeypatch): + # A test creating LLMs concurrently (ThreadPoolExecutor) + # raced take()'s unlocked read of _thread against schedule_shadow()'s + # assign-then-start critical section, joining a thread that had not been + # started yet ("cannot join thread before it is started"). A slow start() + # widens the assign->start window deterministically. + class _SlowStartThread(threading.Thread): + def start(self): + time.sleep(0.3) # hold the assigned-but-unstarted state visible + super().start() + + monkeypatch.setattr(session_prefetcher.threading, "Thread", _SlowStartThread) + errors = [] + + def _taker(): + time.sleep(0.1) # let schedule_shadow enter its critical section first + try: + prefetcher.take(1) + except RuntimeError as e: # pre-fix: "cannot join thread before it is started" + errors.append(e) + + taker = threading.Thread(target=_taker) + taker.start() + prefetcher.schedule_shadow(1) + taker.join(timeout=10) + assert errors == [] + + +def test_abandoned_build_publish_discards_pool(prefetcher): + # A build that outlives _drain()'s join timeout is abandoned (generation + # bump); its late _publish() must shut the pool down, not land it — + # landing would overwrite (and leak) a newer pool or hand stale state + # to a future test. + pool = _FakePool(4) + gen = prefetcher._build_gen + prefetcher._build_gen += 1 # what _drain() does on abandonment + prefetcher._publish(4, pool, session_prefetcher._spawn_snapshot(), gen) + assert pool.shut + assert prefetcher._built is None + + +def test_publish_never_overwrites_unconsumed_pool(prefetcher): + first, second = _FakePool(4), _FakePool(4) + _arm(prefetcher, first, spec=4) + _arm(prefetcher, second, spec=4) + assert prefetcher._built.session is first # slot kept + assert second.shut and not first.shut # newcomer discarded, not the slot + + +def test_session_summary_counters_and_emission(prefetcher, capfd): + # Handover / stale-discard / superseded each count once; + # dispose() emits ONE summary line (outside pytest capture in real runs, + # the only guaranteed console-visible record of prefetch activity). + hit = _FakePool(4) + _arm(prefetcher, hit, spec=4) + assert prefetcher.take(4) is hit + stale = _FakePool(4) + _arm(prefetcher, stale, spec=4) + assert prefetcher.take(2) is None # spec mismatch -> stale discard + late = _FakePool(4) + prefetcher._publish(4, late, session_prefetcher._spawn_snapshot(), prefetcher._build_gen - 1) + assert prefetcher.stats["pools_handed_over"] == 1 + assert prefetcher.stats["pools_discarded_stale"] == 1 + assert prefetcher.stats["pools_discarded_superseded"] == 1 + prefetcher.dispose() + out = capfd.readouterr().out + assert "[session-prefetch] session summary:" in out + assert "pools_handed_over=1" in out + + +def test_no_summary_when_prefetch_never_fired(prefetcher, capfd): + prefetcher.dispose() + assert "session summary" not in capfd.readouterr().out + + +def test_warm_counters_track_gib(tmp_path, monkeypatch): + monkeypatch.setenv("TRTLLM_TEST_PREFETCH_SESSION", "1") + monkeypatch.delenv("PYTEST_XDIST_WORKER", raising=False) + p = SessionPrefetcher() # real _warm (fixture would stub it) + (tmp_path / "model-00001.safetensors").write_bytes(b"x" * (1 << 20)) + p._warm(str(tmp_path)) + assert p.stats["warms"] == 1 and p._warmed_gib > 0 + p._warm("not/a/real/dir") + assert p.stats["warm_noops"] == 1 + + +def test_model_switch_triggers_warm_and_dedups(prefetcher): + items = _as_session(_FakeItem("/models/a"), _FakeItem("/models/b")) + prefetcher.on_test_setup(items[0]) + # Warm threads are fire-and-forget; poll briefly for the recorded call. + assert _wait_for(lambda: prefetcher.warmed) + assert prefetcher.warmed == ["/models/b"] + # Same next-model again: deduplicated. + prefetcher.on_test_setup(items[0]) + time.sleep(0.2) + assert prefetcher.warmed == ["/models/b"] + + +def test_same_next_model_does_not_warm(prefetcher): + # Consecutive tests on the same model: its weights are already hot. + items = _as_session(_FakeItem("/models/a"), _FakeItem("/models/a")) + prefetcher.on_test_setup(items[0]) + time.sleep(0.2) + assert prefetcher.warmed == [] + + +def test_no_marker_suite_never_warms(prefetcher): + items = _as_session(_FakeItem(), _FakeItem(), _FakeItem()) + prefetcher.on_test_setup(items[0]) + time.sleep(0.2) + assert prefetcher.warmed == [] + + +def test_auto_model_dir_from_accuracy_class_attr(prefetcher): + # Accuracy-harness classes declare MODEL_PATH; warming must pick it up + # automatically, with no marker on the test. + class _TestLlama: + MODEL_PATH = "/models/llama" + + class _TestQwen: + MODEL_PATH = "/models/qwen" + + items = _as_session(_FakeItem(cls=_TestLlama), _FakeItem(cls=_TestQwen)) + prefetcher.on_test_setup(items[0]) + assert _wait_for(lambda: prefetcher.warmed) + assert prefetcher.warmed == ["/models/qwen"] + + +def test_model_param_discovery(monkeypatch): + # Modeling unit tests carry the model as a `model_folder`-style parameter + # (a name under LLM_MODELS_ROOT, or an absolute path) — discovered + # without any test-file changes. + monkeypatch.setenv("LLM_MODELS_ROOT", "/models-root") + item = _FakeItem(params={"model_folder": "Nemotron-H-8B-Base-8K"}) + assert session_prefetcher._model_dir_of(item) == "/models-root/Nemotron-H-8B-Base-8K" + item = _FakeItem(params={"model_dir": "/abs/path/model"}) + assert session_prefetcher._model_dir_of(item) == "/abs/path/model" + # Non-model params never produce a guess. + assert session_prefetcher._model_dir_of(_FakeItem(params={"dtype": "fp8"})) is None + + +def test_marker_overrides_class_model_path(): + class _TestCls: + MODEL_PATH = "/models/from-class" + + item = _FakeItem(model_dir="/models/from-marker", cls=_TestCls) + assert session_prefetcher._model_dir_of(item) == "/models/from-marker" + + +def test_accuracy_harness_still_declares_model_path(): + # _model_dir_of auto-discovers models via the MODEL_PATH class attribute of + # the accuracy harnesses (accuracy_core.py); renaming that attribute would + # silently kill warming repo-wide. Textual check — importing accuracy_core + # would drag integration-only dependencies into this unit test. + core = Path(__file__).parents[2] / "integration" / "defs" / "accuracy" / "accuracy_core.py" + assert re.search(r"^\s+MODEL_PATH\s*=", core.read_text(), re.MULTILINE), ( + "accuracy_core.py no longer declares MODEL_PATH — update " + "session_prefetcher._model_dir_of to the harness's new convention" + ) + + +def test_warm_selects_files_like_the_weight_loader(tmp_path): + # Selection must mirror HfWeightLoader.load_weights: safetensors first + # (minus huge "consolidated" copies the loader skips), so the .bin copy + # and the consolidated file must NOT be read here. + payload = b"x" * (1 << 20) + (tmp_path / "model-00001.safetensors").write_bytes(payload) + (tmp_path / "consolidated.safetensors").write_bytes(payload * 4) + (tmp_path / "pytorch_model.bin").write_bytes(payload) + (tmp_path / "config.json").write_bytes(b"{}") # not a weight file + assert warm_page_cache(str(tmp_path)) == pytest.approx(1 / 1024, rel=1e-3) + + +def test_warm_falls_back_to_bin_then_pth(tmp_path): + payload = b"x" * (1 << 20) + bin_dir, pth_dir = tmp_path / "bin", tmp_path / "pth" + bin_dir.mkdir(), pth_dir.mkdir() + (bin_dir / "pytorch_model.bin").write_bytes(payload) + (pth_dir / "model.pth").write_bytes(payload) + assert warm_page_cache(str(bin_dir)) == pytest.approx(1 / 1024, rel=1e-3) + assert warm_page_cache(str(pth_dir)) == pytest.approx(1 / 1024, rel=1e-3) + + +def test_warm_page_cache_ignores_non_weight_dirs(tmp_path): + # MODEL_PATH may be an HF model id or a dir without local weights: no-op. + assert warm_page_cache(str(tmp_path)) == 0.0 + assert warm_page_cache("not/a/real/dir") == 0.0 + + +def test_warm_skips_models_larger_than_host_memory(tmp_path, monkeypatch): + # Warming a model bigger than free RAM is pure filer traffic: the pages + # would be evicted before the test loads them (DeepSeek-R1-class dirs). + (tmp_path / "model-00001.safetensors").write_bytes(b"x" * (1 << 20)) + monkeypatch.setattr(session_prefetcher, "_available_host_memory", lambda: 1 << 10) + assert warm_page_cache(str(tmp_path)) == 0.0 + + +def test_warm_io_thread_names_covered_by_threadleak_exclude(): + # Both pytest.ini threadleak_exclude lists contain r"session-prefetch-\w+"; + # the warm executor's thread_name_prefix must keep its IO workers inside + # that pattern (a large warm can outlive the test that started it). + assert re.fullmatch(r"session-prefetch-\w+", "session-prefetch-io_0") + + +def _fake_reuse_module(enabled): + return types.SimpleNamespace(REUSE=types.SimpleNamespace(is_active=lambda: enabled)) + + +def test_yields_mpi_seams_to_active_session_reuse(monkeypatch): + # session_reuse owns the same seams and saves the whole respawn; when it + # is enabled the prefetcher must not install its factory (whoever patched + # first would silently disable the other layer). + monkeypatch.setenv("TRTLLM_TEST_PREFETCH_SESSION", "1") + monkeypatch.delenv("PYTEST_XDIST_WORKER", raising=False) + monkeypatch.setitem(sys.modules, "test_common.session_reuse", _fake_reuse_module(enabled=True)) + p = SessionPrefetcher() + p.install_pool_factory_if_loaded() + assert not p._patched + assert p.stats["mpi_yielded_to_reuse"] == 1 + + +def test_reuse_layer_inactive_or_absent_does_not_block_prefetch(monkeypatch): + monkeypatch.setitem(sys.modules, "test_common.session_reuse", _fake_reuse_module(enabled=False)) + assert not session_prefetcher._reuse_layer_active() + monkeypatch.delitem(sys.modules, "test_common.session_reuse") + assert not session_prefetcher._reuse_layer_active() + + +def test_unreadable_reuse_module_errs_on_yielding(monkeypatch): + # Module present but attribute layout changed: stay out of the way. + monkeypatch.setitem(sys.modules, "test_common.session_reuse", types.SimpleNamespace()) + assert session_prefetcher._reuse_layer_active() + + +def test_install_wraps_seam_in_isinstance_transparent_shim(prefetcher, monkeypatch): + # The patched seam must stay a real TYPE: proxy.py's killed-worker + # detection runs isinstance(x, MpiPoolSession) against this attribute, + # and a bare function there raises TypeError (the #16338 breakage class). + mpi_mod = pytest.importorskip("tensorrt_llm.llmapi.mpi_session") + fake = types.ModuleType("fake_seam_mod") + fake.MpiPoolSession = mpi_mod.MpiPoolSession + monkeypatch.setitem(sys.modules, "fake_seam_mod", fake) + monkeypatch.setattr(session_prefetcher, "_PATCH_TARGETS", ("fake_seam_mod",)) + monkeypatch.setattr(session_prefetcher, "_reuse_layer_active", lambda: False) + prefetcher.install_pool_factory_if_loaded() + assert fake.MpiPoolSession is not mpi_mod.MpiPoolSession # patched + # isinstance must not raise, and must answer for the real class. + assert isinstance(object(), fake.MpiPoolSession) is False + assert issubclass(mpi_mod.MpiPoolSession, fake.MpiPoolSession) + + +def test_patch_targets_cover_all_library_construction_sites(): + # The factory only intercepts the modules listed in _PATCH_TARGETS. If the + # library grows another MpiPoolSession(...) construction site, prefetch + # would silently stop covering it (armed spare pools would idle next to + # directly-constructed ones) — turn that drift into a red test. + import tensorrt_llm + + root = Path(tensorrt_llm.__file__).parent + # mpi_session.py defines the class (and the MGMN server path, which + # legitimately builds its own pool outside the bare-LLM() seams). + exempt = {"tensorrt_llm.llmapi.mpi_session"} + offenders = [] + for py in root.rglob("*.py"): + if re.search(r"(?