[None][test] Optimize LLM test startup overhead#15777
Conversation
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughIntroduces an in-memory LRU cache for HF checkpoint weight loading, controlled via environment variables. Adds MPI session ownership tracking across executor, proxy, and LLM classes to avoid shutting down externally-owned sessions. Updates tests to share MPI sessions and adds a grouped NVFP4 4-GPU premerge test with corresponding test-list changes. ChangesWeight caching and MPI session sharing
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/llmapi/test_llm.py (1)
1846-1925: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMissing
try/finallyaroundllm.shutdown()inllm_return_logprobs_test_harness.
llm.shutdown()is only reached if none of theasserts in the generate loop (or streaming block) raise. Unlikellm_get_stats_test_harness/llm_get_stats_async_test_harness(usewith LLM_CLASS(...) as llm:) and_test_llm_capture_request_error(explicittry/finally), a failing assertion here skips cleanup entirely.This matters more now that this harness is called with a shared, module-scoped
MpiPoolSession(test_llm_return_logprobs_streaming_tp2intest_llm_multi_gpu_pytorch.py). Per mpi4py's docs, workers serve submitted tasks until signaled for completion; an un-shut-down proxy leaves its MPI worker(s) still executing, which can block or delay subsequent tests reusing the same shared session, since tasks queue behind busy workers.🔧 Proposed fix: wrap the harness body in try/finally
- llm = LLM_CLASS( - llama_model_path, - kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4, - **kv_cache_args_extra), - tensor_parallel_size=tp_size, - gather_generation_logits=True, - **llm_args_extra, - **extra_llm_kwargs, - ) - - prompts = ["A B C D E F G H I J K"] - sampling_params = SamplingParams( - logprobs=logprobs, - prompt_logprobs=prompt_logprobs, - return_context_logits=return_context_logits, - return_generation_logits=return_generation_logits) - - for output in llm.generate(prompts, sampling_params): + llm = LLM_CLASS( + llama_model_path, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4, + **kv_cache_args_extra), + tensor_parallel_size=tp_size, + gather_generation_logits=True, + **llm_args_extra, + **extra_llm_kwargs, + ) + try: + prompts = ["A B C D E F G H I J K"] + sampling_params = SamplingParams( + logprobs=logprobs, + prompt_logprobs=prompt_logprobs, + return_context_logits=return_context_logits, + return_generation_logits=return_generation_logits) + + for output in llm.generate(prompts, sampling_params): ... # (unchanged body) - if streaming: + if streaming: ... # (unchanged body, indented one level) - asyncio.run(main()) - llm.shutdown() + asyncio.run(main()) + finally: + llm.shutdown()🤖 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_llm.py` around lines 1846 - 1925, `llm_return_logprobs_test_harness` should always clean up the `LLM_CLASS` instance even when an assertion in the generate/streaming checks fails. Wrap the harness body that creates `llm` and runs `llm.generate` / `llm.generate_async` in a `try/finally`, and move `llm.shutdown()` into the `finally` block so the shared MPI-backed session is not left running. Use the existing `llm` variable and keep the cleanup guaranteed regardless of failures.
🧹 Nitpick comments (3)
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py (1)
74-75: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse Python 3.10 built-in generics for the new annotation.
Line 74 adds
List[str]; preferlist[str]for new Python 3.10+ code.Proposed fix
- def _weight_files_cache_key(weight_files: List[str], + def _weight_files_cache_key(weight_files: list[str], use_consolidated: bool) -> tuple:Based on learnings, Python 3.10+ features are available throughout the codebase. As per coding guidelines, “Prefer using the built-in types
list,dict, andtuple.”🤖 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 `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py` around lines 74 - 75, The new type annotation in _weight_files_cache_key still uses typing.List; update the weight_files parameter to use the Python 3.10 built-in generic list[str] instead. Keep the rest of the signature unchanged and apply the same built-in typing style consistently in this function if any related annotations are added.Sources: Coding guidelines, Learnings
tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py (1)
103-103: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd return annotations to the new test functions.
Line 103 and Line 136 should use
-> None.Proposed fix
-def test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper(tmp_path, monkeypatch): +def test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper(tmp_path, monkeypatch) -> None:-def test_weight_cache_disabled_by_default(tmp_path, monkeypatch): +def test_weight_cache_disabled_by_default(tmp_path, monkeypatch) -> None:As per coding guidelines, “Always annotate functions. Make the return type
Noneif the function does not return anything.”Also applies to: 136-136
🤖 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/_torch/models/checkpoints/hf/test_weight_loader.py` at line 103, Add explicit return type annotations to the new pytest functions in test_weight_loader, including test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper and the other new test at the referenced location; both should be annotated with -> None to match the project’s function annotation guideline. Use the test function names to locate and update their definitions without changing any test logic.Source: Coding guidelines
tests/integration/defs/accuracy/test_llm_api_pytorch.py (1)
2585-2593: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winGuard against silent coverage loss from duplicate case IDs.
Line 2585 builds a dict keyed by
id, so duplicate entries in_NVFP4_4GPU_PREMERGE_CASESwould overwrite earlier cases before the set comparison runs. Add a uniqueness assertion before creatingcases_by_id.Proposed check
+ case_ids = tuple(case["id"] for case in self._NVFP4_4GPU_PREMERGE_CASES) + assert len(case_ids) == len(set(case_ids)), ( + "Duplicate DeepSeek V3 Lite NVFP4 premerge case IDs" + ) cases_by_id = { case["id"]: case for case in self._NVFP4_4GPU_PREMERGE_CASES }As per path instructions,
tests/**: act as a QA engineer reviewing test changes and coverage for TensorRT-LLM.🤖 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/integration/defs/accuracy/test_llm_api_pytorch.py` around lines 2585 - 2593, Add a pre-check in the NVFP4 4GPU premerge coverage test to fail fast on duplicate case IDs before building cases_by_id, since the dict comprehension in the accuracy test can silently overwrite earlier entries. In the test method that uses _NVFP4_4GPU_PREMERGE_CASES and _NVFP4_4GPU_PREMERGE_GROUPS, assert that the list of case["id"] values is unique, then keep the existing grouped_case_ids coverage assertions after cases_by_id is created.Source: Path instructions
🤖 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 `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py`:
- Around line 55-71: The HF weight cache max-entry setting is only enforced on
cache writes, so existing entries in `_WEIGHT_CACHE` can still be reused when
`TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES` is set to 0 or invalid. Update the cache
lookup path in `weight_loader.py` so `_weight_cache_max_entries()` is checked
before returning any cached weights, and treat non-positive/disabled values as a
hard disable for both read and write behavior. Apply the same guard consistently
in the affected load/save flow around `_WEIGHT_CACHE`,
`_is_weight_cache_enabled()`, and `_weight_cache_max_entries()` so cached
entries are never reused when the cache is disabled.
In `@tensorrt_llm/executor/executor.py`:
- Around line 639-645: Guard the caller-supplied MPI session in
GenerationExecutor._create_ipc_executor so a reused MpiPoolSession cannot
mismatch the expected worker count. Add a size check before using mpi_session,
comparing its worker count against model_world_size (or the derived IPC worker
count), and reject or ignore the injected session when the sizes differ. Keep
the validation close to the existing need_spawn_mpi_workers() logic in
executor.py and use the existing symbols mpi_session, model_world_size, and
_create_ipc_executor to locate the change.
In `@tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py`:
- Around line 103-105: The cache tests currently only set TRTLLM_HF_WEIGHT_CACHE
in test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper, so
TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES can leak in from the environment and change
reuse behavior. Update the cache test setup in test_weight_loader.py to
explicitly control both cache env vars, and add a focused regression around
HfWeightLoader’s cache-control path for max entries set to 0 and invalid values
so the reload/reuse behavior is covered. Coverage for this path is currently
insufficient and should be expanded with a dedicated test.
---
Outside diff comments:
In `@tests/unittest/llmapi/test_llm.py`:
- Around line 1846-1925: `llm_return_logprobs_test_harness` should always clean
up the `LLM_CLASS` instance even when an assertion in the generate/streaming
checks fails. Wrap the harness body that creates `llm` and runs `llm.generate` /
`llm.generate_async` in a `try/finally`, and move `llm.shutdown()` into the
`finally` block so the shared MPI-backed session is not left running. Use the
existing `llm` variable and keep the cleanup guaranteed regardless of failures.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py`:
- Around line 74-75: The new type annotation in _weight_files_cache_key still
uses typing.List; update the weight_files parameter to use the Python 3.10
built-in generic list[str] instead. Keep the rest of the signature unchanged and
apply the same built-in typing style consistently in this function if any
related annotations are added.
In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Around line 2585-2593: Add a pre-check in the NVFP4 4GPU premerge coverage
test to fail fast on duplicate case IDs before building cases_by_id, since the
dict comprehension in the accuracy test can silently overwrite earlier entries.
In the test method that uses _NVFP4_4GPU_PREMERGE_CASES and
_NVFP4_4GPU_PREMERGE_GROUPS, assert that the list of case["id"] values is
unique, then keep the existing grouped_case_ids coverage assertions after
cases_by_id is created.
In `@tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py`:
- Line 103: Add explicit return type annotations to the new pytest functions in
test_weight_loader, including
test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper and the other
new test at the referenced location; both should be annotated with -> None to
match the project’s function annotation guideline. Use the test function names
to locate and update their definitions without changing any test logic.
🪄 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: e4c1ec00-9f13-4c67-88a1-4d8ae0f406b8
📒 Files selected for processing (11)
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.pytensorrt_llm/executor/executor.pytensorrt_llm/executor/proxy.pytensorrt_llm/executor/rpc_proxy.pytensorrt_llm/llmapi/llm.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/test-db/l0_dgx_h100.ymltests/integration/test_lists/test-db/l0_gb200_multi_gpus.ymltests/unittest/_torch/models/checkpoints/hf/test_weight_loader.pytests/unittest/llmapi/test_llm.pytests/unittest/llmapi/test_llm_multi_gpu_pytorch.py
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #57095 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #57105 [ run ] triggered by Bot. Commit: |
|
PR_Github #57095 [ run ] completed with state |
|
PR_Github #57105 [ run ] completed with state
|
Superjomn
left a comment
Superjomn
left a comment
There was a problem hiding this comment.
LGTM on the mpi_session related changes
|
/bot run --disable-fail-fast |
|
PR_Github #57341 [ run ] triggered by Bot. Commit: |
|
PR_Github #57341 [ run ] completed with state
|
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>
|
/bot reuse-pipeline |
|
PR_Github #58899 [ reuse-pipeline ] triggered by Bot. Commit: |
|
PR_Github #58899 [ reuse-pipeline ] completed with state |
| if prior is not None and prior is not real: | ||
| self._retire(real) # a pool of this size is already cached | ||
| return | ||
| self._pools[real.n_workers] = real |
There was a problem hiding this comment.
Overall, this is a useful optimization and the normal test isolation semantics remain unchanged: a failed pytest item does not automatically fail other items.
The existing health-probe path should also handle a hard worker failure: if a worker future raises or times out, the cached pool is retired with broken=True and the next test gets a fresh pool.
The remaining gap is a failed test whose pool still appears healthy to the probe. Since shutdown() is overloaded to return the pool to the cache, a test that leaves undetected worker-side state behind could make that pool available to up to 15 subsequent same-size tests (max_uses=16). This is not expected to happen for every test failure, but it could create a difficult-to-diagnose cascade when it does.
Could we add a lightweight pytest-level failure fence?
- Record any failed item from
pytest_runtest_logreport. - After teardown, in
pytest_runtest_logfinish, call
REUSE.drain()if that item failed. - Keep the current reuse behavior unchanged for passing
tests.
For example:
def pytest_runtest_logreport(report):
if report.failed:
_FAILED_ITEMS.add(report.nodeid)
def pytest_runtest_logfinish(nodeid, location):
if nodeid in _FAILED_ITEMS:
_FAILED_ITEMS.remove(nodeid)
REUSE.drain()Draining all cached pools is intentionally conservative and only adds the pool-shutdown cost after a failure; the normal passing-test path has no additional overhead. This should be treated as a fallback for state contamination that the health probe cannot detect, not as a replacement for the existing broken-pool retirement path.
Please also add a regression test showing that after a failed item, the next acquire does not reuse the previously cached pool.
There was a problem hiding this comment.
Good catch, implemented in b3e7a1f exactly as you sketched: pytest_runtest_logreport
records failed items, and pytest_runtest_logfinish drains all cached pools after the
failed item's teardown. Passing tests are unchanged. Also added the two regression tests:
after a failed item the next acquire gets a fresh pool, and after a passing item reuse
stays intact.
A failed test's pool can pass the health probe (workers alive and responsive) while still carrying worker-side state the probe cannot see, and would otherwise serve up to max_uses-1 subsequent same-size tests. Add a pytest-level failure fence per review feedback: record failed items in pytest_runtest_logreport and drain ALL cached pools in pytest_runtest_logfinish once the failed item's teardown has returned its pool to the cache. This is a conservative fallback for contamination the probe cannot detect, not a replacement for the broken-pool retirement path. The passing-test path is unchanged and pays no extra cost; a failure adds one pool-shutdown plus a fresh spawn for the next test. Regression tests: after a failed item the next acquire gets a fresh pool; after a passing item reuse stays intact. Signed-off-by: qgai <qgai@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #58921 [ run ] triggered by Bot. Commit: |
|
|
||
| @property | ||
| def max_uses(self) -> int: | ||
| return int(os.environ.get("TRTLLM_TEST_REUSE_MAX_USES", "16")) |
There was a problem hiding this comment.
Why we need to limit the max uses? What happens if we reuse unlimited times?
A wedged pool shutdown inside a drain thread must not keep the interpreter alive at process exit: the bounded join only protects the drain caller, while a non-daemon thread would make pytest hang until the CI stage timeout. The retire path already uses a daemon thread; apply the same to drain, per review feedback. Signed-off-by: qgai <qgai@nvidia.com>
If a broken pool's worker has already exited, the OS may recycle its PID before the background retire thread runs, and os.kill would then hit an unrelated process owned by the same user - most plausibly a worker of the replacement pool being spawned concurrently. Record (pid, start_time from /proc/pid/stat) pairs at spawn and re-verify the start time right before SIGKILL; a mismatch or missing process means the PID no longer belongs to the recorded worker and is skipped, falling back to the graceful client-side reap. Per review feedback. Signed-off-by: qgai <qgai@nvidia.com>
Retire threads were fire-and-forget: a disposal still in flight when the session ends would be cut off by interpreter exit (daemon threads), leaving a graceful shutdown half-done. Track them in a registry and have drain() - which runs at the natural rendezvous points (failure fence, opt-out, session finish) - join them with the same bounded wait used for pool shutdowns, per review feedback. The per-test hot path stays non-blocking, and daemon=True still guarantees a wedged disposal cannot hang interpreter exit. Signed-off-by: qgai <qgai@nvidia.com>
| expected = mpi_session.n_workers | ||
| seen: set = set() | ||
| for _ in range(max_rounds): | ||
| futures = mpi_session.submit(_run_and_get_worker_rank, fn) |
There was a problem hiding this comment.
Forgive me, I dont fully understand this/we can discuss.
Will this cause the fn to be run multiple times?
There was a problem hiding this comment.
Yes, it could — and you were right to flag it. Fixed in b086d20: each task now
starts with an MPI_Barrier, so a worker that grabs a task blocks until every
worker has one. The n tasks land exactly one per worker, in one round. No more
retry loop, and fn never runs twice.
| try: | ||
| for _ in range(4): | ||
| pids.update(real.submit_sync(_get_worker_pid)) | ||
| if len(pids) >= n_workers: |
There was a problem hiding this comment.
May work, but I dont fully understand this one, why need 4 tries?
Shall we use more reliable way to collect pids if the submit_sync does not work reliably?
There was a problem hiding this comment.
Agreed. The 4 tries were working around the same issue as the probe loop:
MPIPoolExecutor doesn't guarantee one task per worker. b086d20 replaces both
with one barrier-based helper that runs the task exactly once per worker, so
the magic numbers (8 rounds / 4 tries) are gone. Verified on a real 4-GPU
pool: 20/20 runs hit all 4 workers exactly once, and a wedged worker still
fails the probe at the timeout.
MPIPoolExecutor has no one-task-per-worker guarantee, so both the pool health probe (8 resubmit rounds) and worker PID collection (4 tries) compensated with retry-and-dedup loops. Replace both with one exactly-once primitive: each submitted task opens with an MPI_Barrier across the worker world, so a worker holding a task blocks until every worker holds one — by pigeonhole the n_workers tasks land exactly one per worker, in a single round. This also strengthens the probe: completing the barrier proves the worker world's collectives still function (the failure mode wedged pools actually exhibit), not just the task loop. Timeout semantics and the broken-pool retirement path are unchanged. Validated on a real 4-GPU MPIPoolExecutor spawn: 20/20 probes covered exactly 4 workers with worker COMM_WORLD size 4, and a wedged worker turned the probe into a bounded RuntimeError at the timeout. Signed-off-by: qgai <qgai@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #58954 [ run ] triggered by Bot. Commit: |
|
PR_Github #58921 [ run ] completed with state |
|
PR_Github #58954 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58991 [ run ] triggered by Bot. Commit: |
|
PR_Github #58991 [ run ] completed with state |
7 participants
Summary
Changes
Test plan
Summary by CodeRabbit
New Features
Bug Fixes
Tests