[None][test] Automatic MPI session reuse for tests via shared pytest infra - #16055
Closed
sunnyqgg wants to merge 8 commits into
Closed
[None][test] Automatic MPI session reuse for tests via shared pytest infra#16055sunnyqgg wants to merge 8 commits into
sunnyqgg wants to merge 8 commits into
Conversation
This was referenced Jul 7, 2026
Contributor
📝 WalkthroughWalkthroughAdds a test infrastructure feature that automatically reuses MpiPoolSession instances across consecutive bare LLM(...) tests. Implements SessionReuseCache with spawn-time snapshotting, GPU memory settle checks, and lifecycle management; wires it into pytest via new hooks, conftest.py, and pytest.ini updates; adds worker-side helper utilities and unit tests. ChangesMPI Session Reuse Infrastructure
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Test
participant SessionReuseCache
participant MpiPoolSession
participant NVML
Test->>SessionReuseCache: acquire(worker_count)
SessionReuseCache->>SessionReuseCache: check cached pool snapshot (env/sys.path/use cap)
alt cached pool valid
SessionReuseCache->>MpiPoolSession: reset_worker_torch_compile_state
SessionReuseCache->>NVML: wait_gpu_memory_settle
SessionReuseCache-->>Test: reused pool
else no valid pool
SessionReuseCache->>MpiPoolSession: _spawn_fresh (with env overrides)
SessionReuseCache-->>Test: new pool
end
Test->>SessionReuseCache: shutdown() / shutdown_abort()
alt normal shutdown
SessionReuseCache->>SessionReuseCache: _release (store pool for reuse)
else abort
SessionReuseCache->>SessionReuseCache: _forget (discard pool)
end
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (5)
tests/unittest/llmapi/test_session_reuse.py (2)
14-16: 📐 Maintainability & Code Quality | 🔵 TrivialMove local
import os/import timeto module level.These imports inside
_FakePool.__init__,_wait, andtest_weight_cache_env_scoped_to_spawncould be hoisted to the top of the file for consistency and to avoid repeated imports on every call.♻️ Proposed fix
+import os +import time + import pytest from test_common import session_reuse from test_common.session_reuse import SessionReuseCache class _FakePool: def __init__(self, n_workers): self.n_workers = n_workers self.shut = False - import os - self.spawn_env_weight_cache = os.environ.get("TRTLLM_HF_WEIGHT_CACHE")def _wait(pred, timeout=5.0): - import time - t0 = time.monotonic()def test_weight_cache_env_scoped_to_spawn(reuse_cache, monkeypatch): - import os - monkeypatch.delenv("TRTLLM_HF_WEIGHT_CACHE", raising=False)Also applies to: 39-40, 133-135
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/llmapi/test_session_reuse.py` around lines 14 - 16, The file has repeated local imports that should be hoisted to module scope for consistency and to avoid re-importing on each call. Move the `os` and `time` imports used by `_FakePool.__init__`, `_wait`, and `test_weight_cache_env_scoped_to_spawn` to the top-level imports, then remove the inner import statements while keeping the existing references to those symbols unchanged.
25-35: 📐 Maintainability & Code Quality | 🔵 TrivialMissing return type annotations on fixtures/test functions.
None of the fixtures or test functions in this file declare
-> None. As per coding guidelines, "Always annotate functions with return types (useNoneif no return); annotate class members and variables when necessary, especially for dataclasses and NamedTuple."Also applies to: 38-46, 49-155
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/llmapi/test_session_reuse.py` around lines 25 - 35, Add explicit return type annotations to every fixture and test function in this file, using -> None for functions that do not return a value. Update the fixture and test definitions such as reuse_cache and the other test helpers/functions in the module so they comply with the coding guideline. Keep the behavior unchanged; this is a typing-only cleanup across all affected functions in the file.Source: Coding guidelines
tests/test_common/session_reuse.py (3)
79-346: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFunction signatures are missing type annotations throughout.
Most module-level and class functions in this file lack parameter and/or return type annotations (e.g.
_spawn_snapshot()at Line 79,_visible_gpu_indices(count: int)at Line 98,_describe_mismatch(...)at Line 163,SessionReuseCache.__init__at Line 207,acquire(self, real_cls, n_workers)at Line 271,_spawn_fresh,_release,_forget), while others in the same class (suspend,drain,install_pool_factory_if_loaded) are fully annotated — inconsistent within the file.As per coding guidelines: "Always annotate functions with return types (use
Noneif no return); annotate class members and variables when necessary" and "Prefer built-in typeslist,dict,tuple... Prefer specifying argument types inCallabletype hints".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_common/session_reuse.py` around lines 79 - 346, Several functions in SessionReuseCache and its helpers are missing type annotations, creating inconsistency with the rest of the file. Add explicit parameter and return type hints to the unannotated helpers and methods such as _spawn_snapshot, _describe_mismatch, SessionReuseCache.__init__, acquire, _spawn_fresh, _release, and _forget, and use built-in collection types where applicable. Keep the existing annotated style used by _visible_gpu_indices and the other fully typed methods in SessionReuseCache.Source: Coding guidelines
111-161: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGPU settle barrier fails completely silently.
Both
exceptblocks that guard the main polling logic (Lines 154-155, 159-160) swallow every exception with no logging at all, unlike the rest of the file which prints on every notable branch (retiring, reuse, rebuild). If this barrier silently breaks (e.g. due to an NVML API change or a bug in_free_total/_visible_gpu_indices), the very "insufficient GPU memory" failures this feature was built to prevent would resurface with no diagnostic trail pointing back here.🔧 Proposed fix: log on unexpected failure
except Exception: - pass + print("[session-reuse] GPU memory settle check failed, proceeding without it", flush=True) finally: try: pynvml.nvmlShutdown() except Exception: pass🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_common/session_reuse.py` around lines 111 - 161, The GPU settle barrier in wait_gpu_memory_settle currently swallows unexpected exceptions in the main polling path and shutdown path, which makes failures completely silent. Update the broad except blocks around the NVML polling logic and nvmlShutdown cleanup to emit a concise diagnostic message (with the exception details) before returning or continuing, so issues in _free_total, _visible_gpu_indices, or NVML API changes are visible in test output.
234-267: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFailed patch guard silently and permanently disables reuse for that seam.
If
getattr(mod, "MpiPoolSession", None) is real_clsisFalse(e.g. an upstream refactor changes howtensorrt_llm.executor.proxy/tensorrt_llm.llmapi.llmimportMpiPoolSession), the module is still added toself._patchedat Line 267, so the seam is never retried and reuse for that path silently becomes a permanent no-op for the rest of the session with zero log output. Given the whole point of this PR is CI-time savings, a silent regression here would go unnoticed for a long time.🔧 Proposed fix: log when the guard fails
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 + else: + print( + f"[session-reuse] {name}.MpiPoolSession is not the expected class; " + "reuse disabled for this seam", + flush=True, + ) self._patched.add(name)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_common/session_reuse.py` around lines 234 - 267, The patch guard in install_pool_factory_if_loaded can fail silently when a target module’s MpiPoolSession is no longer the expected real_cls, and the seam is still added to self._patched. Update the logic so failed guards are reported with a log message, and only add a module to self._patched after a successful replacement; if the guard fails, leave it pending so the seam can be retried later. Reference the install_pool_factory_if_loaded path and the _patched bookkeeping to keep the reuse hook observable and recoverable.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/test_common/session_reuse.py`:
- Around line 352-363: The drain() shutdown path in session_reuse can block
forever because it starts non-daemon threads and waits with an unbounded join,
unlike _retire() which uses fire-and-forget behavior. Update drain() to bound
the thread wait with a timeout (and handle any still-running shutdown threads
without blocking the caller), using the drain() and _retire() behavior as the
reference points so rpc_factory and pytest_sessionfinish cannot hang on a stuck
MpiPoolSession.shutdown.
---
Nitpick comments:
In `@tests/test_common/session_reuse.py`:
- Around line 79-346: Several functions in SessionReuseCache and its helpers are
missing type annotations, creating inconsistency with the rest of the file. Add
explicit parameter and return type hints to the unannotated helpers and methods
such as _spawn_snapshot, _describe_mismatch, SessionReuseCache.__init__,
acquire, _spawn_fresh, _release, and _forget, and use built-in collection types
where applicable. Keep the existing annotated style used by _visible_gpu_indices
and the other fully typed methods in SessionReuseCache.
- Around line 111-161: The GPU settle barrier in wait_gpu_memory_settle
currently swallows unexpected exceptions in the main polling path and shutdown
path, which makes failures completely silent. Update the broad except blocks
around the NVML polling logic and nvmlShutdown cleanup to emit a concise
diagnostic message (with the exception details) before returning or continuing,
so issues in _free_total, _visible_gpu_indices, or NVML API changes are visible
in test output.
- Around line 234-267: The patch guard in install_pool_factory_if_loaded can
fail silently when a target module’s MpiPoolSession is no longer the expected
real_cls, and the seam is still added to self._patched. Update the logic so
failed guards are reported with a log message, and only add a module to
self._patched after a successful replacement; if the guard fails, leave it
pending so the seam can be retried later. Reference the
install_pool_factory_if_loaded path and the _patched bookkeeping to keep the
reuse hook observable and recoverable.
In `@tests/unittest/llmapi/test_session_reuse.py`:
- Around line 14-16: The file has repeated local imports that should be hoisted
to module scope for consistency and to avoid re-importing on each call. Move the
`os` and `time` imports used by `_FakePool.__init__`, `_wait`, and
`test_weight_cache_env_scoped_to_spawn` to the top-level imports, then remove
the inner import statements while keeping the existing references to those
symbols unchanged.
- Around line 25-35: Add explicit return type annotations to every fixture and
test function in this file, using -> None for functions that do not return a
value. Update the fixture and test definitions such as reuse_cache and the other
test helpers/functions in the module so they comply with the coding guideline.
Keep the behavior unchanged; this is a typing-only cleanup across all affected
functions in the file.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 88028f50-2ed7-4200-a522-7d764e1aee3d
📒 Files selected for processing (7)
tests/conftest.pytests/integration/defs/pytest.initests/test_common/grouped_test_utils.pytests/test_common/session_reuse.pytests/test_common/session_reuse_hooks.pytests/unittest/llmapi/test_session_reuse.pytests/unittest/pytest.ini
Comment on lines
+352
to
+363
| def drain(self) -> None: | ||
| """Shut down all cached pools in parallel (frees GPU/CPU footprint).""" | ||
| with self._lock: | ||
| pools, self._pools = list(self._pools.values()), {} | ||
| if not pools: | ||
| return | ||
| threads = [threading.Thread(target=p.shutdown, name="session-reuse-drain") for p in pools] | ||
| for t in threads: | ||
| t.start() | ||
| for t in threads: | ||
| t.join() | ||
| print(f"[session-reuse] drained {len(pools)} cached pool(s)", flush=True) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
drain() can hang indefinitely on a stuck MPI shutdown.
t.join() has no timeout and the threads aren't daemons. _retire() deliberately uses a fire-and-forget daemon thread specifically to avoid blocking on a hung MpiPoolSession.shutdown(wait=True) (per the upstream MpiPoolSession.shutdown contract). drain() reintroduces exactly that blocking risk, and it's invoked synchronously from rpc_factory (blocking RPC pool creation) and from pytest_sessionfinish (blocking end-of-session teardown) — a single hung pool shutdown would stall the whole test run.
🔧 Proposed fix: bound the join with a timeout
+_DRAIN_JOIN_TIMEOUT_S = 30.0
+
+
class SessionReuseCache:
...
def drain(self) -> None:
"""Shut down all cached pools in parallel (frees GPU/CPU footprint)."""
with self._lock:
pools, self._pools = list(self._pools.values()), {}
if not pools:
return
- threads = [threading.Thread(target=p.shutdown, name="session-reuse-drain") for p in pools]
+ threads = [
+ threading.Thread(target=p.shutdown, name="session-reuse-drain", daemon=True)
+ for p in pools
+ ]
for t in threads:
t.start()
for t in threads:
- t.join()
+ t.join(timeout=_DRAIN_JOIN_TIMEOUT_S)
+ stuck = sum(t.is_alive() for t in threads)
+ if stuck:
+ print(f"[session-reuse] {stuck} pool(s) still shutting down after drain timeout", flush=True)
print(f"[session-reuse] drained {len(pools)} cached pool(s)", flush=True)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def drain(self) -> None: | |
| """Shut down all cached pools in parallel (frees GPU/CPU footprint).""" | |
| with self._lock: | |
| pools, self._pools = list(self._pools.values()), {} | |
| if not pools: | |
| return | |
| threads = [threading.Thread(target=p.shutdown, name="session-reuse-drain") for p in pools] | |
| for t in threads: | |
| t.start() | |
| for t in threads: | |
| t.join() | |
| print(f"[session-reuse] drained {len(pools)} cached pool(s)", flush=True) | |
| _DRAIN_JOIN_TIMEOUT_S = 30.0 | |
| class SessionReuseCache: | |
| ... | |
| def drain(self) -> None: | |
| """Shut down all cached pools in parallel (frees GPU/CPU footprint).""" | |
| with self._lock: | |
| pools, self._pools = list(self._pools.values()), {} | |
| if not pools: | |
| return | |
| threads = [ | |
| threading.Thread(target=p.shutdown, name="session-reuse-drain", daemon=True) | |
| for p in pools | |
| ] | |
| for t in threads: | |
| t.start() | |
| for t in threads: | |
| t.join(timeout=_DRAIN_JOIN_TIMEOUT_S) | |
| stuck = sum(t.is_alive() for t in threads) | |
| if stuck: | |
| print( | |
| f"[session-reuse] {stuck} pool(s) still shutting down after drain timeout", | |
| flush=True, | |
| ) | |
| print(f"[session-reuse] drained {len(pools)} cached pool(s)", flush=True) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_common/session_reuse.py` around lines 352 - 363, The drain()
shutdown path in session_reuse can block forever because it starts non-daemon
threads and waits with an unbounded join, unlike _retire() which uses
fire-and-forget behavior. Update drain() to bound the thread wait with a timeout
(and handle any still-running shutdown threads without blocking the caller),
using the drain() and _retire() behavior as the reference points so rpc_factory
and pytest_sessionfinish cannot hang on a stuck MpiPoolSession.shutdown.
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58239 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58239 [ run ] completed with state
|
jieli-matrix
approved these changes
Jul 10, 2026
…infra Eligible tests automatically reuse a shared MpiPoolSession with zero test changes - no signature edits, no wrapper functions. A pytest plugin (loaded via -p in both test trees' pytest.ini, plus a tests/-level fallback conftest for ini-less dirs) lazily patches the pool-construction seams; at LLM shutdown the pool is returned to a per-size cache and handed to the next bare LLM(...) of the same size, saving the ~50-65s worker spawn+import each time. Eligibility is automatic; tests with special needs keep their private lifecycle: RPC executors get a drain-then-build factory at their own seam (private, never-cached pools); tests passing their own _mpi_session never reach the patched seam; @pytest.mark.private_mpi_session opts out explicitly; pools are retired when env or sys.path changed since spawn (workers freeze both), after a use-count cap (default 16), or when the per-handout health probe fails - every failure path degrades to a fresh synchronous build. Disabled under pytest-xdist. TRTLLM_TEST_REUSE_SESSION=0 turns everything off. Between handouts every worker runs a rank-verified torch._dynamo reset and the handover waits for the previous worker's GPU memory release (NVML settle barrier). Cache-managed pool spawns also enable the worker-side HF weight cache (companion weight-loader PR) scoped to those pools only. Validated on B200: 12 pure-logic unit tests; unmodified single-GPU tests reuse a 1-worker pool; multi-GPU suite 15 passed with 10 handovers; DeepSeek-V3-Lite accuracy tp4/ep4/tp2pp2/pp4 4 passed on one shared pool, GSM8K at reference; plus a ~2.5h 550-case soak (42 handovers, 2 full retire/rebuild cycles, 0 failures). Split out of NVIDIA#15777 per review (test-side only; works with or without the companion MPI-session-ownership and weight-cache PRs, degrading gracefully). Signed-off-by: qgai <qgai@nvidia.com>
- A wrapper whose pool was already released to the cache ignores late shutdown_abort calls (executor error paths can fire after shutdown; the pool may belong to the NEXT test by then); shutdown is idempotent; plain attributes replace the __dict__ indirection. - acquire() consults the TRTLLM_TEST_REUSE_SESSION kill switch on every call, so flipping it off after the seams were patched really restores private-pool behavior. - drain() joins each pool shutdown with a 60s bound and warns instead of hanging the suite on one wedged pool; the factory bypass branch logs instead of silently disabling reuse; wait_gpu_memory_settle drops its never-passed timeout parameter. Two new unit tests cover the released-wrapper abort guard and the mid-suite kill switch (14 pass). Signed-off-by: qgai <qgai@nvidia.com>
The XQA JIT cubin registry is process-global (DecoderXQARunner::getResourceGlobal) and its lookup key omits compile context fields q_seq_len and is_spec_dec_tree. Running the dynamic-tree and non-dynamic-tree variants of this test inside one reused MPI worker makes the second variant hit the cubin compiled for the first variant's q_seq_len (both fold into the same m_tilesize bucket), which fails at launch with CUDA_ERROR_INVALID_VALUE on Hopper. Reproduced deterministically on H100 with the exact CI wheel: the cross-config pair fails in both orders at pool use #2, while same-config pairs, reuse-disabled runs, and fresh-pool single runs all pass. On B200 the same handover passes because SM100 does not use this XQA JIT path. Opt the test out via the private_mpi_session marker until the registry key is fixed upstream to include the full JIT compile context. Signed-off-by: qgai <qgai@nvidia.com>
… reuse On a reused MPI worker (pool use #2) this test deterministically fails with 'The expanded size of the tensor (128) must match the existing size (256)': AutoDeploy worker-process state sized for the previous test's config leaks into the next engine. The same case passes on every recent build without session reuse. Same failure class as the eagle3 XQA JIT opt-out; keep this test on a private session until the state source is root-caused. Signed-off-by: qgai <qgai@nvidia.com>
Three AutoDeploy tests have now failed on reused workers with the same failure class (tensor expand-size mismatch such as 'expanded size of the tensor (128) must match the existing size (256)', at pool use #2): registry accuracy, guided decoding, and eagle3 one-model. AutoDeploy's executor adapter keeps worker-process state sized for the previous engine's config, so per-test markers are whack-a-mole. Replace the individual marker with a plugin-level node-id pattern opt-out (autodeploy / auto_deploy / test_ad_) in session_reuse_hooks and revert the now-redundant marker on the registry accuracy test. Remove the patterns once AutoDeploy re-initializes that state per engine. Signed-off-by: qgai <qgai@nvidia.com>
Two robustness gaps let a single broken pool hang the whole suite,
observed when a test's executor dies mid-shutdown ('Failed to send
object') and poisons the shared pool:
- The health probe called submit_sync, whose future.result() has no
timeout. Workers that are alive but wedged in a collective never
complete the probe task, so the next acquire blocked forever. Bound
each probe round with concurrent.futures.wait(round_timeout) and
treat a timeout as probe failure (retire + fresh spawn).
- _retire disposed pools with a graceful shutdown(), which blocks
indefinitely on wedged workers and leaks their GPU memory into
subsequent tests. Prefer shutdown_abort (bounded grace, then kill)
with shutdown(wait=False) as fallback.
Signed-off-by: qgai <qgai@nvidia.com>
The previous hardening disposed of broken pools via shutdown_abort, which calls MPI_COMM_WORLD.Abort and kills the parent test process along with the workers (observed as 'MPI_ABORT was invoked on rank 0' taking down the whole pytest session right after a pool rebuild). Record the worker PIDs at spawn (best-effort submit_sync rounds on the fresh, healthy pool) and have _retire SIGKILL them before reaping the client side with a plain shutdown. Retired pools are discarded, so nothing needs a graceful stop, and the driver reclaims GPU memory on process death. Missing PIDs simply fall back to graceful shutdown. Signed-off-by: qgai <qgai@nvidia.com>
Healthy retires (lifetime cap, stale env snapshot, duplicate cache slot) go back to a graceful background shutdown: those workers are idle and exit cleanly, and abnormal termination of MPI-spawned children can upset the MPI runtime in the parent test process. The SIGKILL path is kept only for pools that failed the health probe, where workers may be wedged in a collective and a graceful shutdown would block forever. Signed-off-by: qgai <qgai@nvidia.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
sunnyqgg
force-pushed
the
test_auto_reuse
branch
from
July 10, 2026 15:31
2864458 to
ae405f8
Compare
Collaborator
|
PR_Github #58679 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58679 [ run ] completed with state
|
Collaborator
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Test-side changes split out of #15777 per review: eligible tests automatically reuse a shared MPI session with zero test changes — no per-test signature edits, no wrapper functions, no individual test modifications.
A pytest plugin (loaded via
-pin both test trees' pytest.ini, plus atests/-level fallback conftest for ini-less dirs) lazily patches the pool-construction seams; at LLM shutdown the pool is returned to a per-size cache and handed to the next bareLLM(...)of the same size, saving the ~50-65s worker spawn+import each time.Eligibility is automatic; special tests keep their private lifecycle
_mpi_session: never reach the patched seam@pytest.mark.private_mpi_session: explicit opt-outTRTLLM_TEST_REUSE_SESSION=0turns everything offSafety between handouts
Rank-verified per-worker
torch._dynamoreset (recompile-limit accumulation), NVML settle barrier for the previous worker's GPU memory release, and retire-reason logging (lifetime cap / changed env vars by name / sys.path delta). Cache-managed pool spawns also enable the worker-side HF weight cache (companion PR) scoped to those pools only.Validation (B200)
test_llm_multi_gpu_pytorch.py(15 tests)Companion PRs
Recommended merge order: #16053 first. Reused worker processes re-enter
init_pp_commwhen a second pp>1 LLM runs on them; #16053's communicator change (reuse the PPCommNCCL for an unchanged world size) fixes an observed ncclCommDestroy deadlock on that path, and all validation of this PR ran with it. #16054 is optional either way: without it the weight-cache env this PR sets for its managed pools is a harmless no-op (pool reuse still saves the spawn).Test plan
Summary by CodeRabbit
New Features
Bug Fixes
Tests