[None][test] Add opt-in background prefetch of test MPI sessions and model page cache - #15908
Merged
Conversation
Collaborator
Author
|
/bot run --disable-fail-fast |
Contributor
📝 WalkthroughWalkthroughAdds a new test utility module implementing optional background session prefetching for multi-GPU LLM-API tests, enabled via an environment variable. It prebuilds MPI worker pools and warms OS page cache for model weight files, wires into pytest hooks, and includes a pure-logic test suite. ChangesSession Prefetcher
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant PytestHook
participant SessionPrefetcher
participant MpiPoolSession
participant PageCache
PytestHook->>SessionPrefetcher: on_collection(items)
PytestHook->>SessionPrefetcher: on_test_setup(item)
SessionPrefetcher->>SessionPrefetcher: derive spec and model_dir
alt spec boundary reached
SessionPrefetcher->>MpiPoolSession: build pool in background thread
end
alt model dir changed
SessionPrefetcher->>PageCache: warm_page_cache(model_dir) in background thread
end
PytestHook->>SessionPrefetcher: take(spec)
SessionPrefetcher->>SessionPrefetcher: join background thread
alt spec matches
SessionPrefetcher-->>PytestHook: return prefetched session
else mismatch
SessionPrefetcher->>MpiPoolSession: shutdown discarded session
SessionPrefetcher-->>PytestHook: return None
end
🚥 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 (3)
tests/unittest/llmapi/test_session_prefetcher.py (3)
79-79: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove
import timeto module top-level.Importing inline inside a test function is non-idiomatic; place it with the other imports at the top of the file.
🤖 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_prefetcher.py` at line 79, The test file has an inline import of time inside a test rather than at module scope. Move the import into the top-level imports alongside the other module imports in test_session_prefetcher so the test stays idiomatic and the import location is consistent.Source: Coding guidelines
71-90: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valuePolling-with-sleep for async assertion; consider tightening.
Waiting up to 5s (
100 * 0.05s) for the background warm call to land works but adds latency risk in CI if the thread is ever delayed, and is slower than necessary on the happy path teardown. Since the warm thread handle isn't exposed bySessionPrefetcher.on_test_setup(per the core module context), polling is the only option available from this test file without touching the core module. Giventests/**coverage guidance, this is acceptable but worth a lighter-weight signal (e.g., athreading.Eventset inside the monkeypatched_warmlambda) for a more deterministic and faster test.♻️ Suggested refactor using an Event instead of polling
def test_model_switch_triggers_warm_even_within_block(prefetcher): + warmed_event = threading.Event() + orig_warmed = prefetcher.warmed + def _record(d): + orig_warmed.append(d) + warmed_event.set() items = [ _FakeItem(spec=2, model_dir="/models/a"), _FakeItem(spec=2, model_dir="/models/b"), ] prefetcher.on_collection(items) prefetcher.on_test_setup(items[0]) - # Warm threads are fire-and-forget; poll briefly for the recorded call. - import time - - for _ in range(100): - if prefetcher.warmed: - break - time.sleep(0.05) + assert warmed_event.wait(timeout=5) assert prefetcher.warmed == ["/models/b"]🤖 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_prefetcher.py` around lines 71 - 90, This test waits on background warming with a long polling loop, which makes the assertion slower and less deterministic. Update test_model_switch_triggers_warm_even_within_block to use a tighter synchronization signal, such as a threading.Event set by the monkeypatched _warm path, so the test can wait directly for the warm call instead of polling prefetcher.warmed. Keep the existing SessionPrefetcher.on_collection and SessionPrefetcher.on_test_setup coverage, but replace the sleep loop with a faster explicit wait and preserve the deduplication assertion.
47-47: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd return type annotations to test functions.
None of the test/helper functions in this file annotate a return type. As per coding guidelines, "Always annotate functions with return types (use
Noneif no return)."Also applies to: 55-55, 62-62, 71-71, 92-92, 98-98, 107-107
🤖 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_prefetcher.py` at line 47, Add return type annotations to the test/helper functions in test_session_prefetcher.py so each function explicitly declares a return of None. Update the definitions for the affected test functions, including test_disabled_is_noop and the other referenced tests in this file, to follow the coding guideline of always annotating return types even when no value is returned.Source: Coding guidelines
🤖 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/unittest/llmapi/test_session_prefetcher.py`:
- Around line 92-95: Add coverage for the happy path in
SessionPrefetcher.take(): the current test only exercises the spec-mismatch
branch by setting _built_spec and _built_session directly. Add a test around
take() that sets matching built state, verifies it returns the built session,
and then verifies the internal state is cleared so a second take() with the same
spec returns None; use the existing SessionPrefetcher fixture and the take
method to validate the consume-once behavior.
---
Nitpick comments:
In `@tests/unittest/llmapi/test_session_prefetcher.py`:
- Line 79: The test file has an inline import of time inside a test rather than
at module scope. Move the import into the top-level imports alongside the other
module imports in test_session_prefetcher so the test stays idiomatic and the
import location is consistent.
- Around line 71-90: This test waits on background warming with a long polling
loop, which makes the assertion slower and less deterministic. Update
test_model_switch_triggers_warm_even_within_block to use a tighter
synchronization signal, such as a threading.Event set by the monkeypatched _warm
path, so the test can wait directly for the warm call instead of polling
prefetcher.warmed. Keep the existing SessionPrefetcher.on_collection and
SessionPrefetcher.on_test_setup coverage, but replace the sleep loop with a
faster explicit wait and preserve the deduplication assertion.
- Line 47: Add return type annotations to the test/helper functions in
test_session_prefetcher.py so each function explicitly declares a return of
None. Update the definitions for the affected test functions, including
test_disabled_is_noop and the other referenced tests in this file, to follow the
coding guideline of always annotating return types even when no value is
returned.
🪄 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: c80a006d-33b2-4126-81c5-bb510a689648
📒 Files selected for processing (3)
tests/test_common/session_prefetcher.pytests/unittest/llmapi/conftest.pytests/unittest/llmapi/test_session_prefetcher.py
Collaborator
|
PR_Github #57386 [ run ] triggered by Bot. Commit: |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57394 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57386 [ run ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57402 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57394 [ run ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57407 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57402 [ run ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast |
1 similar comment
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57418 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57420 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57418 [ run ] completed with state |
Collaborator
|
PR_Github #57407 [ run ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast |
1 similar comment
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57434 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57436 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57434 [ run ] completed with state |
Collaborator
|
PR_Github #57420 [ run ] completed with state |
Collaborator
|
PR_Github #57436 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57952 [ run ] triggered by Bot. Commit: |
Collaborator
Author
|
/bot run --disable-fail-fast |
Both layers hand a live MpiPoolSession to a test that did not spawn it, so they share the same invariants — and duplicated the code for them: the spawn-time env+sys.path snapshot, the GPU-memory settle poll, and CUDA_VISIBLE_DEVICES resolution. Extract those into tests/test_common/_session_utils.py; policy (what to do on mismatch or a still-busy GPU) stays in each layer's thin wrapper. Side effects of sharing one implementation: - session_reuse's settle barrier now understands the UUID/MIG forms of CUDA_VISIBLE_DEVICES (previously fell back to polling ALL GPUs on shared CI nodes) and treats an explicitly empty value as no GPUs. - SessionReuseCache gains a public is_active() probe, replacing the prefetcher's reach into REUSE.enabled — seam-yield coordination now goes through a declared API (suspension via private_mpi_session deliberately does not affect seam ownership). Signed-off-by: qgai <qgai@nvidia.com>
MPIPoolExecutor.shutdown returns at disconnect, but a worker's GPU memory is only released when its process actually exits. Anything that starts new GPU work right after shutdown() races that release — the test layers that hand a live pool to the next test (session reuse and prefetch) hit it as 'Executor creation failed due to insufficient GPU memory' in pre-merge CI, and back-to-back LLM() creations in one test survive today only because the second ~50s spawn happens to outlast the release. Root fix: MpiPoolSession(n, wait_shutdown=True) records each worker's (pid, start_time) right after spawn (barrier-pinned one task per worker) and blocks shutdown() until those processes are gone (bounded by a 30s timeout; identity collection is best-effort and fails open to the old behavior). Production is untouched by construction: the default constructs and shuts down exactly as before — no identity collection, no wait. Both test layers now build their pools with the flag, making handover safety deterministic. The NVML settle barrier stays as defense-in-depth for GPU memory the factories never see (pools tests construct themselves, in-process allocations); with factory pools now draining at shutdown it degenerates to a single NVML query on a free GPU. Verified end to end: shutdown(wait_shutdown=True) blocked 1.7s until both spawned workers' PIDs were gone; the default path collects nothing and returns at disconnect as before. Signed-off-by: qgai <qgai@nvidia.com>
…stead With MpiPoolSession(wait_shutdown=True) as the deterministic protection (a pool's shutdown blocks until its workers exited and released their GPU memory), the NVML settle heuristic at handover is redundant for every pool the prefetch/reuse factories build. The remaining corpse producers — the three tests that construct their own MpiPoolSession directly (AutoDeploy multigpu x2, test_base_worker) — now build with wait_shutdown=True as well, closing the same race at the source instead of detecting it downstream. Deleted: wait_gpu_memory_settle in both layers, the 0.5 refuse threshold and busy-discard path in take(), the settle poll and GPU handle resolution in _session_utils (snapshot is all that is shared now), and the settle unit tests. Signed-off-by: qgai <qgai@nvidia.com>
MpiPoolSession(wait_shutdown=True) already records every worker's (pid, start_time) at spawn, so session_reuse's own barrier-pinned collection was a second round-trip gathering the same data. Drop _proc_start_time/_get_worker_pid/_collect_worker_pids and read the library's _worker_identities instead (getattr with () fallback keeps the SIGKILL retire path best-effort as before); the PID-recycling guard in _retire now lazy-imports the library's _process_start_time. Signed-off-by: qgai <qgai@nvidia.com>
…dover Correctness review finding on the settle-barrier removal: the handover comment claimed wait_shutdown=True makes instant handovers safe, but the flag blocks whichever THREAD calls shutdown() — and reuse retires pools in a background daemon, so the test hot path never waited. A duplicate pool retired by _release (two same-size LLMs in one test, both released) holds full model GPU memory until its workers exit; the next cached-pool handover is instant and raced that release — the exact window the deleted NVML settle barrier used to cover. Fix: extract drain()'s bounded retire-thread reap into _reap_retires() and call it before the cached handover. Deterministic (the retire thread itself blocks on worker exit) and a no-op on the common path; every non-cached path already spawns fresh (~50s), outlasting the release. Also name the poisoned-pool consequence in the identity-collection timeout warning (workers can be stuck in the collection barrier if one died during spawn), so a later hang is attributable. Signed-off-by: qgai <qgai@nvidia.com>
…sses With the prefetcher yielding the pool seams to session reuse, reuse's cache misses (first pool of a size, post-drain rebuild after a failure fence, post-retire replacement) still paid the full ~50s synchronous spawn. Close that gap by composing the layers: reuse's miss path (_spawn_fresh) now takes a pre-spawned shadow from the prefetcher when one is armed, and restocks one shadow for the next miss of that size. Coordination mirrors the existing seam-yield probe (sys.modules, lazy, fail-open both ways) so neither layer imports the other at load time. Shadows for reuse are armed with its worker-side weight-cache env. To freeze that into the workers from a BACKGROUND build thread without mutating the parent process env (a set-around-spawn from a thread would race the running test), MpiPoolSession gains env_overrides: extra vars merged into the worker env the library already passes to MPIPoolExecutor, parent env untouched. reuse's own synchronous spawn switches from its set-then-restore dance to the same channel; an explicit user setting of a cache var still wins (overrides only carry keys absent from the parent env). Verified end to end: workers report the override values while the parent env never contains them. Signed-off-by: qgai <qgai@nvidia.com>
…r size Review finding on the shadow-vending path: take() joined ANY in-flight build (up to 180s) before checking the spec, so a cache miss for size A could wait most of a spawn for a size-B shadow only to discard it — slower than no prefetch at all (wait one spawn, then spawn again). The factory path had the same exposure on size changes; vending to reuse's misses makes mixed sizes common enough to matter. Track the in-flight build's spec and return None immediately on a mismatch: the caller falls back to the synchronous spawn right away and the undisturbed build still lands for a later take of its own size. Counted as pools_skipped_size_in_flight in the session summary. Signed-off-by: qgai <qgai@nvidia.com>
…params Replace the explicit prefetch_model_dir markers on the Nemotron-H tests with one more automatic discovery rule: a model_folder/model_dir/ model_path test parameter holding a directory name under LLM_MODELS_ROOT (or an absolute path). The modeling unit-test files share that convention, so they all get weight warming with zero test changes — restoring the PR's no-test-edits property for warming (test_modeling_nemotron_h.py reverts to upstream). Guess-tolerant as before: a value that is not a real weight directory makes warming a silent no-op. Signed-off-by: qgai <qgai@nvidia.com>
The prefetcher was wired only where tests/unittest and tests/integration/defs load it via their own conftest; the repo-root fallback conftest (covering microbenchmarks/, scripts/, torch/ and any future dir under tests/) carried session reuse only. Dispatch to BOTH plugins there — the hooks share names, so importing both modules' functions directly would silently shadow one set, and pytest_plugins is not allowed in a non-rootdir conftest. Also wires reuse's failure-fence hooks (pytest_runtest_logreport/logfinish) into the fallback dirs, which the direct-import list had left out. Signed-off-by: qgai <qgai@nvidia.com>
pytest forbids pytest_plugins in a non-top-level conftest: a repository-root invocation (pytest tests) loads tests/unittest/conftest.py and tests/integration/defs/conftest.py as NESTED conftests and fails the whole collection with 'Defining pytest_plugins in a non-top-level conftest is no longer supported' (review finding; reproduced verbatim). Normal per-tree CI runs never hit it because each tree is its own rootdir there. Replace the declaration with explicit dispatch, mirroring tests/conftest.py: forward to the plugin from the existing pytest_configure and add pytest_runtest_setup/pytest_sessionfinish wrappers. Make dispose() idempotent since a repo-root run dispatches sessionfinish from both the repo-root and the subtree conftest (the other hooks are naturally idempotent). Signed-off-by: qgai <qgai@nvidia.com>
…omplete Review finding: returning () on a collection timeout let a wait_shutdown=True pool stay in service with the contract silently stripped — futures_wait does not cancel pending tasks, so on a slow-but-healthy bootstrap the identity tasks complete later, the import probe succeeds, and the session gets published; at shutdown _wait_workers_exit iterates an empty tuple and returns immediately, reducing wait_shutdown=True to the old non-waiting behavior. The next prefetched pool could then race the old workers' GPU-memory release — the exact OOM the flag exists to prevent, and the justification for having removed the NVML settle barrier. Fail closed instead: require exactly n_workers unique (pid, start_time) identities; on any shortfall tear the pool down (shutdown(wait=False) — a graceful stop could hang on barrier-stuck workers — plus best-effort SIGKILL of the identities we did collect, with the pid-recycling guard) and raise. The test layers keep their resilience explicitly and loudly: session_reuse retries one fresh spawn before propagating, and the prefetcher factory retries once then degrades to a plain pool (pre-prefetch semantics), counted as pools_spawned_degraded in the session summary. Signed-off-by: qgai <qgai@nvidia.com>
Reviewer request: limit initial blast radius by enabling the plugin on
representative stages first and observing before a full rollout. With
no explicit TRTLLM_TEST_PREFETCH_SESSION setting, the plugin is now
active only on two canary stage GROUPS:
- multi-GPU LLM-dense: DGX_H100-4_GPUs-PyTorch-DeepSeek
- single-GPU LLM-dense: A10-PyTorch (both shards: the numbered shard
suffix is assigned by dynamic load balancing and cannot be targeted
individually, so the gate matches group prefixes)
CI already exports the stage name as stageName to the pytest process on
both execution paths, so no Jenkins changes are needed. Explicit
TRTLLM_TEST_PREFETCH_SESSION=1/0 overrides the gate in either direction
(manual opt-in anywhere; the kill switch beats the canary). Runs
without a stageName (local dev) default to off during the canary.
Session reuse (NVIDIA#15777) is unaffected: on non-canary stages its cache
misses simply fall back to the synchronous spawn, exactly the
pre-prefetch behavior. Graduation: after a week of clean canary runs
(zero plugin-attributed failures, healthy session-summary counters)
the gate is removed and the plugin returns to enabled-by-default.
Signed-off-by: qgai <qgai@nvidia.com>
The pool-creation seams held a plain FUNCTION in place of MpiPoolSession. Library code doing 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 in the consumer. Remove the hazard class at its source: the seam now holds a real class whose metaclass routes construction to the reuse factory and instance/subclass checks to the real MpiPoolSession, so both usage patterns keep working — including isinstance consumers added after this layer was written. The exclusion-based check in _register_worker_processes can be reverted to the natural isinstance(self.mpi_session, MpiPoolSession) form at leisure. Signed-off-by: qgai <qgai@nvidia.com>
…etcher Move _isinstance_transparent_shim to _session_utils and use it at the prefetcher's seam install too: when the prefetcher owns the seams (reuse disabled or manual opt-in), it had the same hazard the shim removes for reuse — a bare function at the seam makes any isinstance(x, MpiPoolSession) against the patched module attribute raise TypeError (the NVIDIA#16338 breakage class). Adds an install-side test asserting the patched attribute stays a real type that answers for the real class. Signed-off-by: qgai <qgai@nvidia.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
sunnyqgg
force-pushed
the
test_overlap_prefectch
branch
from
July 17, 2026 04:21
6125955 to
54f4ae1
Compare
2 tasks
Collaborator
|
PR_Github #59871 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59845 [ run ] completed with state |
Collaborator
|
PR_Github #59871 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #59944 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59944 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #60162 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #60162 [ run ] completed with state |
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
7 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
Multi-GPU LLM-API tests pay ~50-65s per test to spawn an
MpiPoolSessionwhose workers importtensorrt_llm, plus a cold read of the model weights. Both are pure CPU/IO work, so they run in background threads while the PREVIOUS test still owns the GPUs — hiding the cost behind its runtime.Enabled by default (
TRTLLM_TEST_PREFETCH_SESSION=0to disable). No production code is changed; everything lives undertests/.Three ways tests benefit
install_pool_factory()patches the threeMpiPoolSessionconstruction seams (executor/proxy.py,executor/rpc_proxy.py,llmapi/llm.py, all in theirmpi_session is Nonebranches) so bareLLM(...)tests consume prefetched pools automatically. Whenever a multi-GPU pool is built or handed over, a spare pool of the same size is built in the background for the next test. A size miss is discarded — the sync build is no slower than without prefetch. Tests passing their own_mpi_session(shared/grouped pools) are never intercepted.@pytest.mark.prefetch_session(N)+ theprefetched_mpi_sessionfixture, for suites that know their boundaries.@pytest.mark.prefetch_model_dir(path)pre-reads the NEXT test's weight files (fires even between tests sharing a session).Safety
TRTLLM*/TLLM*/NCCL*/CUDA*/UCX*/OMPI*vars changed since it was built.pytest_sessionfinishdisposes any unconsumed shadow pool.Measured benefit (4xB200, DeepSeek-V3-Lite nvfp4 / Llama-3.1-8B+EAGLE3)
Savings per boundary = min(pool build time, previous test duration).
Changes
tests/test_common/session_prefetcher.py— prefetcher, shadow-mode factory, env guard, page-cache warming (new)tests/unittest/llmapi/conftest.py— hook wiring, factory install, marker registration, fixture (new)tests/unittest/llmapi/test_session_prefetcher.py— pure-logic unit tests, no MPI/GPU needed (new)Test plan
pytest tests/unittest/llmapi/test_session_prefetcher.py— 14 passed (container, B200 node)