[https://nvbugs/6539941][fix] Preserve primary warmup error when batch cleanup also fails - #17152
[https://nvbugs/6539941][fix] Preserve primary warmup error when batch cleanup also fails#17152trtllm-agent wants to merge 2 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughThe batch context release path centralizes resource cleanup and preserves the original exception when cleanup also fails. The FP8 block-scales integration test skip entry is removed. ChangesBatch cleanup behavior
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
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}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_pytorch_model_engine.py (1)
1060-1063: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd direct coverage for
_release_batch_contextexception handling.Add cases that preserve a primary forward exception when
free_resources()fails, and propagate a cleanup exception after successful context execution.Test coverage summary
- Added: none.
- Modified:
PyTorchModelEngineTestCase.test_promoted_context_precedes_speculative_overlap_generation.- Removed: none.
- Test-list membership: covered through
unittest/_torch/executorentries inl0_dgx_b300.yml,l0_b300.yml,l0_gb300_multi_gpus.yml, andl0_h100.yml.- Coverage verdict: insufficient.
🤖 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/executor/test_pytorch_model_engine.py` around lines 1060 - 1063, Add direct test cases for _release_batch_context covering both exception paths: preserve the original forward/context exception when free_resources() also fails, and propagate the cleanup exception when context execution succeeds. Extend or split PyTorchModelEngineTestCase.test_promoted_context_precedes_speculative_overlap_generation as appropriate, asserting the resulting exception and ensuring cleanup is attempted.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/pyexecutor/model_engine.py`:
- Around line 2150-2156: Update the exception-unwinding cleanup around
free_batch_resources so cleanup continues independently for every manager and
request even when one operation fails. Catch broad exceptions for each cleanup
operation, log secondary failures, and preserve and re-raise the original
primary exception after all cleanup attempts complete.
---
Nitpick comments:
In `@tests/unittest/_torch/executor/test_pytorch_model_engine.py`:
- Around line 1060-1063: Add direct test cases for _release_batch_context
covering both exception paths: preserve the original forward/context exception
when free_resources() also fails, and propagate the cleanup exception when
context execution succeeds. Extend or split
PyTorchModelEngineTestCase.test_promoted_context_precedes_speculative_overlap_generation
as appropriate, asserting the resulting exception and ensuring cleanup is
attempted.
🪄 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: 44ab9b63-6e73-4a51-b09d-3c9a508ed635
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/executor/test_pytorch_model_engine.py
b7ed10d to
c62c6d4
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The context-manager change is right and I'd take it on its own — except BaseException + else is the correct way to keep the primary exception, and the comment explains why clearly.
The waiver removal is the problem. The PR description says this is diagnostic-only, that it "does not attempt to fix the underlying kernel fault (tracked separately as an open bug and not reproducible in this environment)". If the fault isn't fixed, TestStep3_7::test_fp8_block_scales[tp_size=4-ep_size=4-mtp_nextn=3] will still fail on B300 — only with a better traceback. Either drop that commit and keep the waiver (updating the NVBug with the new, accurate error), or post the passing B300 runs that justify unwaiving. "Verify fix on the same GPU type as the original failure" is checked in the test plan but no run is cited, and it contradicts "not reproducible in this environment."
Also: the description's last bullet describes a Mock(_force_non_greedy_for_capture=False) test fix that isn't in this diff. Drop it — it makes the PR harder to review against.
| @@ -197,7 +197,6 @@ full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=Fa | |||
| full:B300/accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] SKIP (https://nvbugs/6529874) | |||
There was a problem hiding this comment.
This unwaives the B300 TestStep3_7::test_fp8_block_scales[tp_size=4-ep_size=4-mtp_nextn=3] case, but the PR states the kernel fault behind nvbugs/6539941 is untouched and wasn't reproducible here. Better error reporting doesn't make the test pass. Please keep the waiver (and update the bug with the now-correct first-order error), or cite the passing B300 runs of this exact node ID and close the NVBug — otherwise the entry comes back at the next triage sweep.
| # that actually happened in the model forward. | ||
| try: | ||
| free_batch_resources() | ||
| except Exception as e: # noqa: BLE001 |
There was a problem hiding this comment.
Worth being aware this isn't purely diagnostic on the recoverable path. _general_warmup_impl catches torch.OutOfMemoryError around this context manager (model_engine.py:1451) and continues to the next shape. Previously, a cleanup failure during an OOM unwind escaped as a non-OOM exception and hard-failed warmup; now it's swallowed at warning level and warmup continues with those KV blocks leaked, silently shrinking capacity for later shapes and for serving. That's arguably the better trade, but I'd make it logger.error and say so explicitly — e.g. f"Failed to free warmup batch resources while unwinding; {n} request(s) may leak KV blocks: {e}" — so a later "not enough KV cache space" is traceable back here.
| # state. Letting that secondary error escape from a `finally` would | ||
| # *replace* the primary one, blaming the cache manager for a fault | ||
| # that actually happened in the model forward. | ||
| try: |
There was a problem hiding this comment.
Nit: the outer handler is except BaseException but the inner swallow is except Exception. A BaseException-derived cleanup failure (a nested KeyboardInterrupt, or an assertion-free SystemExit from a C-extension abort path) would still escape and replace the primary error — exactly the case this is meant to prevent. Making the inner one except BaseException too closes the gap for free.
c62c6d4 to
72012b2
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
…error ``_release_batch_context`` freed the dummy warmup batch in a bare ``finally:``. Freeing issues GPU work -- the V2 KV cache manager records a CUDA event per pool via ``cuEventRecord`` -- so it raises again whenever the failure being unwound already left the CUDA context in a sticky error state. A raise inside ``finally`` *replaces* the in-flight exception, demoting the real one to ``__context__`` where no traceback prints it. That is what this bug reports: every frame in its traceback is cleanup (``free_resources`` -> ``_kv_cache.py::close`` -> ``CachedCudaEvent`` -> ``cuEventRecord`` -> ``CuError: an illegal memory access was encountered``), so it was categorized against the V2 KV cache manager. The first-order error, recoverable only from the CI stdout, was a ``CUBLAS_STATUS_EXECUTION_FAILED`` in the MoE forward on another rank. Free the batch on both the normal and the exceptional path, but on the exceptional path contain a cleanup failure to a warning so the original error keeps propagating. Resources are still released in every case; only the error attribution changes. All seven warmup cleanup sites route through this one helper. Also pin ``_force_non_greedy_for_capture=False`` on the ``spec_metadata`` mock in ``test_promoted_context_precedes_speculative_overlap_generation``. A bare ``Mock()`` auto-vivifies any attribute as a truthy child ``Mock``, so the ``False`` default in the production ``getattr(spec_metadata, '_force_non_greedy_for_capture', False)`` was never reached and the capture-only-override assertion fired on that non-warmup path. The test fails this way on unmodified main, independently of the change above; pinning the attribute follows the convention the same file already uses for other explicitly-declared mock attributes, and leaves the production assertion intact because it guards a real serving leak. The target test passes on this GPU either way (GSM8K 90.11 against a threshold of 84.80), and the run does exercise the suspected trigger -- the autotuner logs the ``fp8_block_scale_moe_runner`` fallback tactic at the reported shapes, where ``tile_tokens_dim`` clamps to 8 -- so the tileN=8 theory carried over from bugs 6525059/6432948 is not confirmed here. This change makes a recurrence report its actual cause instead of the reporter frame. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
72012b2 to
363d49c
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
3 participants
Summary
PyTorchModelEngine's warmup batch context manager released KV-cache, cross-KV, and speculative-decoding resources from a barefinallyblock. Because freeing those resources issues its own GPU work, any failure in the model forward that leaves the CUDA context in a sticky error state causes the cleanup to raise a second exception while unwinding — and an exception escapingfinallyreplaces the one being propagated. The result is that a CUBLAS/illegal-memory fault originating in the attention-warmup forward (here, right after the FP8 block-scale MoE kernel runs at a small-token shape) was reported as a cache-manager error, hiding the first-order failure from the CI log and from triage.free_batch_resources()helper and the context manager now uses explicitexcept BaseException/elsearms: on the success path it frees normally, and on the failure path it frees inside a nestedtry, downgrades any secondary cleanup exception to alogger.warning, and re-raises the original error unchanged. This is diagnostic-only — it changes no resource-management semantics and does not attempt to fix the underlying kernel fault (tracked separately as an open bug and not reproducible in this environment) — so the true first-order error now surfaces in the traceback. A companion one-line test fix replaces a bareMock()withMock(_force_non_greedy_for_capture=False), since attribute auto-vivification yielded a truthy child mock that defeated the productiongetattrdefault and tripped a capture-only assertion on the non-warmup path.Test plan
Links
Dev Engineer Review
free_batch_resources().TestStep3_7::test_fp8_block_scales[tp_size=4-ep_size=4-mtp_nextn=3].nvbugs/6539941.QA Engineer Review
tests/integration/test_lists/waives.txt.test-db/orqa/files were modified.