[TRTLLM-17083][feat] Cancel paired GEN on CTX failure or cancellation

[chienchunhung]

Summary

Stop the paired generation request when context fails or returns finish_reason="cancelled" in generation-first serving, instead of leaving GEN waiting for its own timeout. Follow-up to the early-cancellation discussion in #20459; this targets main independently.

Scope

• Serving orchestration closes streaming GEN before exposing buffered chunks when CTX returns cancellation, and cancels/drains sibling HTTP tasks on non-streaming failure or caller cancellation.
• Preserve prompt non-streaming GEN-error propagation, GEN-only streaming, and HTTP-client ownership of routing/completion hooks.
• Use the existing worker disconnect → request-abort path. This delivers cancellation intent; it does not authorize KV reuse, change transport ownership/deadlines, coordinate retry attempts, or add a backend abort protocol.

Verification

• Adds 14 CPU test cases: ten orchestration cases and four real-HTTP disconnect-to-worker-abort cases. The latter assert the production abort path reaches the mapped engine request ID before test teardown; model execution is mocked.
• Both test files are included by the existing CPU test-list directory entry.
• Scoped pre-commit checks and git diff --check: passed. Full self-review: no remaining findings at 4df975b8c3.
• Current-head full CI requested; launcher passed authorization and is in preflight. Runtime results remain pending.
• Test-only first commit 35d0b9f9a8; fix in 4df975b8c3. Before-fix RED evidence requires separate execution; current-head CI will verify the fix, not both revisions. No physical-quiescence qualification claimed.

Dev Engineer Review

Generation-first orchestration now closes the paired GEN stream before returning a CTX cancellation response. For non-streaming requests, it cancels and drains pending sibling tasks on failure or caller cancellation. It preserves prompt GEN-error reporting while CTX is pending, and leaves GEN-only streaming ownership with the caller.

The cancellation path removes disaggregated parameters from the CTX response. It does not claim physical transfer quiescence or KV-memory reuse. Current review-finding counts and test execution results are unavailable.

QA Engineer Review

The changed tests cover stream closure, non-streaming sibling cleanup, CTX cancellation, caller cancellation, GEN-only streaming ownership, and HTTP disconnect delivery to the worker abort path. The HTTP test covers streaming and non-streaming requests with CTX errors and cancellation; model execution is mocked.

Both test modules are covered by the unittest/disaggregated directory entry in tests/integration/test_lists/test-db/l0_cpu.yml. Coverage verdict: sufficient for the described request-cancellation behavior. Runtime test results are not provided.

Per-File QA Perspective

• tensorrt_llm/serve/openai_disagg_service.py: Verify that CTX cancellation closes GEN before any buffered chunks reach the caller, and that non-streaming failures and caller cancellation drain both HTTP tasks. Also verify GEN-only streaming and GEN-error propagation remain unchanged.
• tests/unittest/disaggregated/test_openai_disagg_service.py: Covers orchestration cleanup, cancellation outcomes, and stream ownership. Its containing test directory is listed in the CPU CI test list.
• tests/unittest/disaggregated/test_paired_request_cancellation.py: Covers real HTTP disconnect propagation into the worker abort path for streaming and non-streaming requests. Its containing test directory is listed in the CPU CI test list.
• tests/integration/test_lists/test-db/l0_cpu.yml: The existing unittest/disaggregated entry includes both test modules in CPU CI selection.

[chienchunhung]

/bot run

[tensorrt-cicd]

PR_Github #78549 [ run ] triggered by Bot. Commit: 4df975b Link to invocation

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 8b342fd5-ba5e-4bde-b749-b5eb9485563b

📥 Commits

Reviewing files that changed from the base of the PR and between c226e2d and 4df975b.

📒 Files selected for processing: 3

• tensorrt_llm/serve/openai_disagg_service.py
• tests/unittest/disaggregated/test_openai_disagg_service.py
• tests/unittest/disaggregated/test_paired_request_cancellation.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

The generation-first service path now handles context cancellation, coordinates non-streaming requests concurrently, and cleans up pending work. Tests cover streaming and non-streaming cancellation, request failures, caller cancellation, generation-only streaming, and worker abort handling.

Changes

Paired request handling

Layer / File(s)	Summary
Request flow and cancellation cleanup tensorrt_llm/serve/openai_disagg_service.py, tests/unittest/disaggregated/test_openai_disagg_service.py	Generation-only requests return directly. Paired requests handle context cancellation and clean up pending generation work. Tests cover cancellation, failures, caller cancellation, and lazy generation-only streaming.
Worker abort integration coverage tests/unittest/disaggregated/test_paired_request_cancellation.py	The integration test checks worker cancellation and promise abortion when a context request errors or returns a cancellation response.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix

Merge Risk: ⚪ Minimal · up to 4df97

No merge-blocking issue is established; the pending CI run remains a normal completion check.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 45.83% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 24 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly identifies the ticket, feature type, and primary change: cancelling the paired GEN request when CTX fails or is cancelled.
Description check		The description clearly explains the problem, solution, scope, test coverage, and current CI status. It does not reproduce the template's PR Checklist section, but the required information is otherwis…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Description check

Explanation

The description clearly explains the problem, solution, scope, test coverage, and current CI status. It does not reproduce the template's PR Checklist section, but the required information is otherwise substantially complete.

• Fix all pre-merge checks with AI

✨ Finishing Touches: 💡 1

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

[tensorrt-cicd]

PR_Github #78549 [ run ] completed with state FAILURE. Commit: 4df975b
/LLM/main/L0_MergeRequest_PR pipeline #64727 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[Shixiaowei02]

When CTX returns cancelled before GEN completes, canceling GEN prevents its on_resp_done, and this path invokes none. The response loses requested metrics headers and configured JSONL records; finalize once.