[https://nvbugs/6539941][fix] Preserve primary warmup error when batch cleanup also fails

[trtllm-agent]

Summary

• Root cause: PyTorchModelEngine's warmup batch context manager released KV-cache, cross-KV, and speculative-decoding resources from a bare finally block. 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 escaping finally replaces 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.
• Fix: The release logic was extracted into a local free_batch_resources() helper and the context manager now uses explicit except BaseException / else arms: on the success path it frees normally, and on the failure path it frees inside a nested try, downgrades any secondary cleanup exception to a logger.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 bare Mock() with Mock(_force_non_greedy_for_capture=False), since attribute auto-vivification yielded a truthy child mock that defeated the production getattr default and tripped a capture-only assertion on the non-warmup path.
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6539941

Dev Engineer Review

• Centralizes warmup batch cleanup in free_batch_resources().
• Suppresses cleanup failures only when another exception is already active.
• Re-raises the original model-forward exception unchanged.
• Preserves normal cleanup error propagation when no earlier exception exists.
• Logs suppressed cleanup failures as warnings.
• The change does not alter public APIs or resource-release behavior.
• The mock fix prevents unintended Mock attribute auto-vivification from triggering unrelated assertions.
• No configuration or test-list files changed.

QA Engineer Review

• Modified test fixture behavior in tests/unittest/_torch/executor/test_pytorch_model_engine.py.
• The change explicitly sets spec_metadata._force_non_greedy_for_capture=False.
• No test function was added, removed, or renamed.
• Test-list coverage was not identified in the provided changes.
• Verdict: needs follow-up.

[coderabbitai]

Walkthrough

The batch context release path now preserves the original exception when cleanup also fails. Normal cleanup errors still propagate. The speculative-overlap test fixture explicitly sets _force_non_greedy_for_capture to False.

Changes

Batch cleanup behavior

Layer / File(s)	Summary
Exception-safe cleanup and test fixture tensorrt_llm/_torch/pyexecutor/model_engine.py, tests/unittest/_torch/executor/test_pytorch_model_engine.py	_release_batch_context uses a local cleanup helper. Cleanup failures are logged and suppressed during exception unwinding, while normal exits propagate them. The test fixture sets _force_non_greedy_for_capture to False.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#17040: Both changes modify PyTorchModelEngine cleanup and resource-release behavior in model_engine.py.

Suggested reviewers: schetlur-nv, mikeiovine, allisonlim-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title uses the required NVBugs and fix format and clearly states that the primary warmup error is preserved.
Description check	✅ Passed	The description explains the root cause and fix, lists test coverage, and addresses the checklist; it is complete despite using “Test plan” instead of “Test Coverage”.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unittest/_torch/executor/test_pytorch_model_engine.py (1)

1060-1063: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add direct coverage for _release_batch_context exception 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/executor entries in l0_dgx_b300.yml, l0_b300.yml, l0_gb300_multi_gpus.yml, and l0_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

📥 Commits

Reviewing files that changed from the base of the PR and between b136596 and b7ed10d.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/model_engine.py
• tests/unittest/_torch/executor/test_pytorch_model_engine.py

[coderabbitai]

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

rg -n -C 8 -P '^\s*(async\s+)?def\s+free_resources\s*\(' --glob '*.py' .
rg -n -C 4 '\.free_resources\(' --glob '*.py' .

Repository: NVIDIA/TensorRT-LLM

Length of output: 50375

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- cleanup helper and enclosing exception path ---'
rg -n -C 18 'free_batch_resources|warmup batch resources|unwinding' tensorrt_llm/_torch/pyexecutor/model_engine.py

printf '%s\n' '--- helper definition and direct calls ---'
rg -n -C 12 'def free_batch_resources|free_batch_resources\(' tensorrt_llm/_torch/pyexecutor/model_engine.py

printf '%s\n' '--- resource-manager cleanup implementation ---'
sed -n '2580,2632p' tensorrt_llm/_torch/pyexecutor/resource_manager.py

Repository: NVIDIA/TensorRT-LLM

Length of output: 8737

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- all definitions and calls of the cleanup helper ---'
rg -n -C 20 'free_batch_resources|warmup batch resources|unwinding' tensorrt_llm/_torch/pyexecutor/model_engine.py

printf '%s\n' '--- cleanup-related exception classes and CUDA error handling ---'
rg -n -C 5 'CUDA|cuda|torch\.cuda|OutOfMemoryError|RuntimeError|except .*Error' \
  tensorrt_llm/_torch/pyexecutor/model_engine.py \
  tensorrt_llm/_torch/pyexecutor/resource_manager.py

Repository: NVIDIA/TensorRT-LLM

Length of output: 50376

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
import ast
from pathlib import Path

root = Path("tensorrt_llm/_torch")
for path in sorted(root.rglob("*.py")):
    try:
        tree = ast.parse(path.read_text())
    except SyntaxError:
        continue
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "free_resources":
            raises = [
                ast.unparse(child)
                for child in ast.walk(node)
                if isinstance(child, ast.Raise) and child.exc is not None
            ]
            calls = [
                ast.unparse(child)
                for child in ast.walk(node)
                if isinstance(child, ast.Call)
            ]
            print(f"{path}:{node.lineno}")
            print("  raises:", raises or ["<implicit>"])
            print("  calls:", calls[:12])
PY

printf '%s\n' '--- manager selection and warmup request cleanup ---'
rg -n -C 8 '_get_draft_kv_cache_manager|kv_cache_manager_key|SPEC_RESOURCE_MANAGER|CROSS_KV_CACHE_MANAGER' \
  tensorrt_llm/_torch/pyexecutor/model_engine.py

Repository: NVIDIA/TensorRT-LLM

Length of output: 26989

---

Continue cleanup after a secondary failure.

free_batch_resources() stops at the first failing manager or request, so later resources remain allocated. Run each cleanup operation independently, log secondary failures, and re-raise the primary exception. The manager methods have no documented single cleanup exception type to catch.

🤖 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/pyexecutor/model_engine.py` around lines 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.

Source: Coding guidelines