[TRTLLM-15264][fix] Fail only the affected requests on disagg peer-layout mismatch - #17482
[TRTLLM-15264][fix] Fail only the affected requests on disagg peer-layout mismatch#17482brnguyen2 wants to merge 3 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe transfer layer caches incompatible peer endpoints, propagates compatibility diagnostics through task failures, and prevents repeated discovery or registration. Tests verify request isolation and compatible transfers. Three integration test waiver entries were added. ChangesPeer incompatibility handling
Integration test waivers
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Receiver
participant SenderInfo
participant TransferTask
participant CompatiblePeer
Receiver->>SenderInfo: Validate peer recurrent-state layout
SenderInfo-->>Receiver: Cache PeerIncompatibleError for endpoint
Receiver->>TransferTask: Mark affected request as FAILED
Receiver->>SenderInfo: Reuse cached incompatibility on repeated request
Receiver->>CompatiblePeer: Dispatch compatible transfer
CompatiblePeer-->>TransferTask: Complete transfer
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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
🤖 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 `@tests/unittest/disaggregated/test_kv_transfer.py`:
- Around line 1567-1585: Extend
test_incompatible_peer_fails_only_affected_requests to assert
wait_complete(blocking=False) returns WaitResult.FAILED after each failed
receive() for rx1 and rx2, while preserving the existing diagnostics and cache
assertions. If DISAGG_TRANS_ERROR mapping is implemented outside this module,
add corresponding executor-level coverage there, and include the required
coverage summary with the test changes.
🪄 Autofix
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: c7de9f57-3b83-4209-929c-0593f4cf322d
📒 Files selected for processing (2)
tensorrt_llm/_torch/disaggregation/native/transfer.pytests/unittest/disaggregated/test_kv_transfer.py
|
/bot run |
|
PR_Github #65161 [ run ] triggered by Bot. Commit: |
|
PR_Github #65161 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65215 [ run ] triggered by Bot. Commit: |
|
PR_Github #65215 [ run ] completed with state
|
|
/bot run |
7c7b84f to
98d480d
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. |
There was a problem hiding this comment.
🧹 Nitpick comments (3)
tests/unittest/disaggregated/test_kv_transfer.py (3)
1511-1511: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRestore the logger level after the test.
tensorrt_llm.logger.set_level("info")changes global state. The test does not restore the previous level, so later tests in the session inherit it.♻️ Proposed fix using the monkeypatch context already in the test
- tensorrt_llm.logger.set_level("info") + _prev_log_level = tensorrt_llm.logger.logger.level + tensorrt_llm.logger.set_level("info")Then restore it in the existing
finallyblock:finally: + tensorrt_llm.logger.set_level(_prev_log_level) for s in (setup_good, setup_bad):🤖 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/disaggregated/test_kv_transfer.py` at line 1511, Update the test containing tensorrt_llm.logger.set_level("info") to save the logger’s previous level and restore it in the existing finally block, ensuring global logger state is unchanged after the test.
1499-1500: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd nonblocking and executor-state assertions
The test module is listed in
tests/integration/test_lists/test-db/l0_h100.yml.KVSlice.token_rangeis optional, so its omission is valid.Test coverage summary:
- Added:
test_incompatible_peer_fails_only_affected_requests- Modified or removed: none
- Covered: blocking failure isolation, cached incompatibility, diagnostic preservation, and compatible-peer transfer
- Missing:
wait_complete(blocking=False)failure andDISAGG_TRANS_ERRORrequest-state mapping- 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/disaggregated/test_kv_transfer.py` around lines 1499 - 1500, Extend test_incompatible_peer_fails_only_affected_requests to assert wait_complete(blocking=False) behavior for the incompatible transfer and verify that the resulting request state maps to DISAGG_TRANS_ERROR, while preserving the existing blocking failure-isolation, cached incompatibility, diagnostics, and compatible-peer assertions.Source: Path instructions
1594-1599: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRelease the KV cache pools explicitly.
KVCacheManager.shutdown()synchronizes and releases the pools, but thefinallyblock stops only the transfer workers. Callshutdown()for each manager in both setups after stopping the workers.🤖 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/disaggregated/test_kv_transfer.py` around lines 1594 - 1599, Update the test cleanup finally block to call KVCacheManager.shutdown() for each manager in setup_good and setup_bad after stopping their context and generator transfer workers. Ensure both setup managers release their KV cache pools explicitly.
🤖 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.
Nitpick comments:
In `@tests/unittest/disaggregated/test_kv_transfer.py`:
- Line 1511: Update the test containing tensorrt_llm.logger.set_level("info") to
save the logger’s previous level and restore it in the existing finally block,
ensuring global logger state is unchanged after the test.
- Around line 1499-1500: Extend
test_incompatible_peer_fails_only_affected_requests to assert
wait_complete(blocking=False) behavior for the incompatible transfer and verify
that the resulting request state maps to DISAGG_TRANS_ERROR, while preserving
the existing blocking failure-isolation, cached incompatibility, diagnostics,
and compatible-peer assertions.
- Around line 1594-1599: Update the test cleanup finally block to call
KVCacheManager.shutdown() for each manager in setup_good and setup_bad after
stopping their context and generator transfer workers. Ensure both setup
managers release their KV cache pools explicitly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d15373ed-8f1d-4f44-988d-da52fd98db33
📒 Files selected for processing (2)
tensorrt_llm/_torch/disaggregation/native/transfer.pytests/unittest/disaggregated/test_kv_transfer.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/disaggregation/native/transfer.py
|
PR_Github #65245 [ run ] triggered by Bot. Commit: |
|
PR_Github #65245 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65298 [ run ] triggered by Bot. Commit: |
|
PR_Github #65298 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65324 [ run ] triggered by Bot. Commit: |
|
PR_Github #65324 [ run ] completed with state
|
|
/bot run |
2989402 to
4d32bb8
Compare
|
/bot run |
|
PR_Github #66088 [ run ] triggered by Bot. Commit: |
|
PR_Github #66088 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66150 [ run ] triggered by Bot. Commit: |
|
PR_Github #66150 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66212 [ run ] triggered by Bot. Commit: |
|
PR_Github #66212 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66267 [ run ] triggered by Bot. Commit: |
|
/bot run |
be8a9b1 to
f6288c8
Compare
|
PR_Github #66284 [ run ] triggered by Bot. Commit: |
|
PR_Github #66267 [ run ] completed with state |
|
PR_Github #66284 [ run ] completed with state
|
|
/bot run |
…yout mismatch Receiver._get_sender_info raised MambaPolicy.validate_peer_compatible's ValueError on the executor thread with nothing catching it, so one incompatible context instance took down the whole generation worker, including traffic from healthy peers. Route the failure through the existing transfer-error path instead: _get_sender_info converts the validation ValueError into a typed PeerIncompatibleError (still a ValueError subclass) that dispatch_task catches to fail just that request's KVRecvTask. The session then reports WaitResult.FAILED and the request ends in DISAGG_TRANS_ERROR, the same mechanism used for remote agent failures and cancellation, while the worker keeps serving other peers. The check itself is unchanged and still runs before REGISTER_RANK_INFO, so no dealers are connected and no partial registration happens for the rejected peer. The incompatibility is cached per info_endpoint so later requests to the same peer fail fast without another REQUEST_INSTANCE_INFO round-trip or re-validation; the cached diagnostic notes that a generation-worker restart is needed to re-validate (e.g. after redeploying a compatible server on the same endpoint). Add a unit test covering: receive() does not raise and the session fails with the diagnostic preserved; the second request to the same endpoint fails fast without re-validation; and the same Receiver still completes a real transfer from a compatible peer. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…t_nvfp4, Initial failure: `accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 (pre-existing failure, tracking bug pending) Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
test_incompatible_peer_fails_only_affected_requests only checked the blocking wait_complete path. Add wait_complete(blocking=False) == WaitResult.FAILED assertions after each failed receive() to cover the non-blocking polling path returning the terminal failure for an errored task (no None / spurious success). Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
f6288c8 to
6cde495
Compare
|
PR_Github #66352 [ run ] triggered by Bot. Commit: |
|
PR_Github #66352 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66422 [ run ] triggered by Bot. Commit: |
|
PR_Github #66422 [ run ] completed with state
|
4 participants
Description
In the native (Python) disaggregation transceiver,
Receiver._get_sender_inforunsMambaPolicy.validate_peer_compatibleat first contact with a new context peer. TheValueErrorit raises on a recurrent-state layout mismatch propagated up the executor thread, so a single incompatible context peer could take down the whole generation worker (error budget / fatal shutdown), and in the async path it also left an orphaned receive session.This change converts that failure into a per-request error:
PeerIncompatibleError(ValueError)intransfer.py.Receivercaches incompatible peers by info endpoint; later requests to the same peer fail fast without re-validating or re-contacting it. The diagnostic notes that a worker restart is required to re-validate (e.g. after the peer is redeployed with a fixed config).Receiver.dispatch_taskcatches the typed error and fails just that task, which flows through the existing task ERROR -> session ERROR ->WaitResult.FAILED->DISAGG_TRANS_ERRORpath (same mechanism as remote-agent failure and cancellation), in both the sync and async transceiver paths. Healthy peers and in-flight requests are unaffected.Test Coverage
New
test_incompatible_peer_fails_only_affected_requestsintests/unittest/disaggregated/test_kv_transfer.py: monkeypatchesvalidate_peer_compatibleto reject one context worker and asserts (a)receive()does not raise and the session fails with the diagnostic preserved, (b) a second request to the same endpoint fails fast with no re-validation and no registration, and (c) the same receiver then completes a real verified transfer from a compatible context worker.Full file runs on Blackwell hardware:
test_kv_transfer.py46 passed,test_kda_mamba_transfer.py14 passed.PR Checklist
[TRTLLM-15264][fix]conventionDev Engineer Review
PeerIncompatibleErrorto isolate recurrent-state layout mismatches by peer.TestLagunaXS::test_fp8,TestLagunaXS::test_nvfp4, andTestLlama3_1_8BInstruct::test_bfloat16_4gpus.QA Engineer Review
test_incompatible_peer_fails_only_affected_requests().WaitResult.FAILED,DISAGG_TRANS_ERROR, diagnostic preservation, cached fail-fast behavior, no repeated validation or registration, and successful transfers from compatible peers.tests/integration/test_lists/.pr-babysitter waive set-bug NVIDIA/TensorRT-LLM#17482 <bug-url>— the waives entry carries a placeholder until then