[None][test] Add E2E test for cancelled disagg gen request with overlap scheduler

[Tabrizian]

Summary by CodeRabbit

• New Features

  • Added request cancellation capability for distributed inference processing, enabling workers to gracefully cancel in-flight requests.

• Tests

  • Introduced test coverage for cancellation semantics in disaggregated single-GPU scenarios.

Description

Add an integration test that exercises the cancel path for disaggregated generation-only requests when the overlap scheduler is enabled. This test catches the bug where a cancelled gen-only request with py_decoding_iter == 0 gets permanently stuck in active_requests due to the overlap scheduler re-adding it without checking request.is_finished.

PR Checklist

• Please check this after reviewing the above items as appropriate for this PR.

[coderabbitai]

📝 Walkthrough

Walkthrough

Introduces MPI-based cancellation support for disaggregated worker processes. New MPI_CANCEL and MPI_STARTED constants enable workers to receive cancellation signals. Updated run_worker and worker_entry_point signatures include optional support_cancel flag, with branching logic to handle cancellation or standard completion paths. Adds corresponding test coverage.

Changes

Cohort / File(s)	Summary
MPI Cancellation Support tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py	Added MPI_CANCEL and MPI_STARTED constants; updated run_worker() and worker_entry_point() to accept support_cancel parameter; implemented dual code paths for cancellation-aware orchestration (signaling generation start, receiving cancellation indices, aborting futures) versus standard future awaiting; enhanced error handling with logging; added test_disaggregated_cancel_gen_requests() test.
Test Configuration tests/integration/test_lists/test-db/l0_h100.yml	Added test_disaggregated_cancel_gen_requests[TinyLlama-1.1B-Chat-v1.0] test reference to PyTorch pre_merge sections.

Sequence Diagram

sequenceDiagram
    participant Orchestrator
    participant Worker1
    participant Worker2 as Worker2<br/>(support_cancel=True)
    participant MPI

    Orchestrator->>Worker1: Submit request indices
    Orchestrator->>Worker2: Submit request indices
    Worker1->>Worker1: Process request (no cancellation)
    Worker2->>MPI: Signal MPI_STARTED
    Orchestrator->>MPI: Send MPI_CANCEL indices
    MPI->>Worker2: Receive cancellation signal
    Worker2->>Worker2: Abort specified futures
    Worker1->>Orchestrator: Send results (normal path)
    Worker2->>Orchestrator: Send results (cancellation path)

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 14.29% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	❓ Inconclusive	The PR description is minimal and lacks several required template sections including Test Coverage and most PR Checklist items.	Expand the description to include detailed Test Coverage section, verify all PR Checklist items, and clarify the technical changes being made to the codebase.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: adding an E2E test for cancelled disaggregated generation requests with the overlap scheduler.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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

[coderabbitai]

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py`:
- Line 183: Replace the two occurrences of "raise e" in the test's except blocks
with a bare "raise" so the original exception traceback and failure location are
preserved; locate the except handlers in
tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py that
currently re-raise the caught exception with "raise e" and change them to
"raise" without altering any surrounding logic or variable names.
- Around line 682-683: The test relies on a default for
disable_overlap_scheduler; make it explicit by adding the
disable_overlap_scheduler key to the generation worker config appended to
worker_pytorch_configs (i.e., include disable_overlap_scheduler: False or True
as appropriate in the dict added to worker_pytorch_configs where
cuda_graph_config is set to None) so the cancel-path coverage is stable and the
test intent is clear.
- Around line 766-768: The except block in the test (around the
try/except/finally in test_disaggregated_single_gpu.py) is catching Exception,
printing it, and swallowing it which hides assertion failures; change this so
you either remove the broad try/except and let the test framework surface
failures, or catch only expected specific exceptions and after logging re-raise
the exception (i.e., log the error in the except and then raise) to ensure
assertions and other errors fail the test while keeping the finally cleanup
intact.

---

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between ab7a20a and 2f24552.

📒 Files selected for processing (2)

• tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py
• tests/integration/test_lists/test-db/l0_h100.yml

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #37125 [ run ] triggered by Bot. Commit: 685d069 Link to invocation

[tensorrt-cicd]

PR_Github #37125 [ run ] completed with state SUCCESS. Commit: 685d069
/LLM/main/L0_MergeRequest_PR pipeline #28742 completed with status: 'SUCCESS'

Link to invocation