[None][test] Add opt-in background prefetch of test MPI sessions and model page cache

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)

tests/conftest.py

@@ -1,22 +1,49 @@
 # 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
 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,
-)
+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)

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)

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 =

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"<pool seam shim for {real_cls!r}>"
+
+    return _SeamMeta("MpiPoolSession", (), {})