[https://nvbugs/6480621][fix] Preserve KV ownership in disaggregated precheck - #17223
[https://nvbugs/6480621][fix] Preserve KV ownership in disaggregated precheck#17223chienchunhung wants to merge 4 commits into
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #63581 [ run ] triggered by Bot. Commit: |
|
PR_Github #63581 [ run ] completed with state
|
|
/bot run --disable-fail-fast --stage-list "DGX_H100-PyTorch-4,A30-PyTorch-2,GB300-12_GPUs-3_Nodes-PyTorch-Disagg-PerfSanity-FUNCTIONAL-ONLY-CTX1-NODE1-GPU4-GEN1-NODE2-GPU8-1" |
|
PR_Github #63819 [ run ] triggered by Bot. Commit: |
|
PR_Github #63819 [ run ] completed with state |
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
c42d28b to
f9ca117
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63865 [ run ] triggered by Bot. Commit: |
WalkthroughThe change propagates ChangesCache transceiver precheck
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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.
🧹 Nitpick comments (2)
tests/unittest/disaggregated/test_transceiver_bounded_polling.py (1)
441-448: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a case that fails a sibling between wait slices.
test_tx_session_blocking_wait_detects_failed_sibling_behind_pending_tasknever reaches the loop body.wait_completecallshas_failed()first, and the ERROR sibling is already present, so the method returnsFAILEDbefore anytask.wait()call. The assertionpending_task.wait_calls == []confirms this. The test therefore duplicatestest_tx_session_blocking_wait_treats_task_failure_as_terminalinstead of covering the in-loophas_failed()check attensorrt_llm/_torch/disaggregation/native/transfer.pylines 1372-1375.To cover that check, start with no failed task and flip a sibling to
ERRORfrom inside the firstwait()call.♻️ Proposed test that exercises the in-loop sibling check
def test_tx_session_blocking_wait_detects_failed_sibling_behind_pending_task() -> None: pending_task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False) - failed_task = _FakeTask(TaskStatus.ERROR) + sibling = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False) - session = _make_tx_session([pending_task, failed_task]) + session = _make_tx_session([pending_task, sibling]) + wait = pending_task.wait + + def fail_sibling_during_wait(timeout: Optional[float] = None) -> bool: + result = wait(timeout) + sibling.status = TaskStatus.ERROR + return result + + pending_task.wait = fail_sibling_during_wait assert session.wait_complete(blocking=True) == WaitResult.FAILED - assert pending_task.wait_calls == [] - assert failed_task.wait_calls == [] + assert pending_task.wait_calls == [0.25] + assert sibling.wait_calls == []🤖 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_transceiver_bounded_polling.py` around lines 441 - 448, Update test_tx_session_blocking_wait_detects_failed_sibling_behind_pending_task so no task is initially in ERROR; make the first pending task’s wait() transition the sibling task to ERROR, then assert wait_complete(blocking=True) returns WaitResult.FAILED and verifies the expected wait calls. This must exercise the in-loop has_failed() check rather than the initial pre-loop failure check.tensorrt_llm/_torch/disaggregation/native/transfer.py (1)
1359-1394: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRemove stale
WaitResult.TIMEOUThandling. The executor passesatLeastNumwith a default of0; no caller undertensorrt_llmpassesNone.TxSessionandRxSessionreturn onlyCOMPLETED,FAILED, orNone, so remove the unreachableTIMEOUT,timed_out, and related consensus plumbing.🤖 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/disaggregation/native/transfer.py` around lines 1359 - 1394, Remove the obsolete WaitResult.TIMEOUT and timed_out consensus handling from the blockAll execution path, including any related plumbing in TxSession and RxSession. Preserve the existing COMPLETED, FAILED, and None outcomes, and keep the atLeastNum default behavior without adding None handling.
🤖 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 `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Around line 1359-1394: Remove the obsolete WaitResult.TIMEOUT and timed_out
consensus handling from the blockAll execution path, including any related
plumbing in TxSession and RxSession. Preserve the existing COMPLETED, FAILED,
and None outcomes, and keep the atLeastNum default behavior without adding None
handling.
In `@tests/unittest/disaggregated/test_transceiver_bounded_polling.py`:
- Around line 441-448: Update
test_tx_session_blocking_wait_detects_failed_sibling_behind_pending_task so no
task is initially in ERROR; make the first pending task’s wait() transition the
sibling task to ERROR, then assert wait_complete(blocking=True) returns
WaitResult.FAILED and verifies the expected wait calls. This must exercise the
in-loop has_failed() check rather than the initial pre-loop failure check.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d9c1baee-87b6-4c9b-aee8-8c23bd2307af
📒 Files selected for processing (12)
jenkins/scripts/perf/local/submit.pyjenkins/scripts/perf/submit.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/transceiver.pytests/scripts/perf-sanity/cache_transceiver_precheck/README.mdtests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.pytests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.pytests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.pytests/unittest/disaggregated/test_transceiver_bounded_polling.pytests/unittest/others/test_cache_transceiver_precheck_config.pytests/unittest/others/test_cache_transceiver_precheck_run.pytests/unittest/scripts/test_perf_submit.py
|
Could you keep the blocking path bounded and consistent before merge?
Without these changes, a stalled peer can wait forever or pages can still be released before transfer completion. |
|
PR_Github #63865 [ run ] completed with state
|
|
The blocking wait now has no deadline left, while the receive side still bounds itself. Serving never takes that path, but the harnesses and precheck do, a stalled peer hangs to the stage limit instead of timing out. |
4 participants
Summary
Fix the disaggregated cache-transceiver precheck so synthetic KV pages remain owned until Python/NIXL transfers are proven complete. The change also aligns the precheck's model/runtime selection and verified payload boundaries with the target test.
Issue and impact
kv_transfer_sender_future_timeout_ms(1 second by default) is a polling interval, not the request-levelkv_transfer_timeout_msdeadline. Previously, the Python sender'sblockAllpath could return after one interval while a send remained queued or in flight. The precheck treated that return as completion and released every KV page in the wave. A later wave could reuse and overwrite those pages before the earlier send finished consuming them, providing a concrete mechanism consistent with the observed timing-dependent payload mismatches.The ownership violation is serious, but its demonstrated in-tree impact is confined to the synthetic CI precheck. That gate transfers deterministic patterns in separate, short-lived MPI processes and fails before the real model servers and benchmark start; its KV pool is destroyed before serving begins. The demonstrated impact is precheck/CI failure before benchmark execution; there is no evidence of corrupted user output from this path.
Normal PyExecutor uses finite
0/1status polling and does not combine sender block-all with unconditional release; the C++ block-all path already waits for completion. An in-tree call-site audit found no production caller using the affected sequence. An out-of-tree caller of this internal Python API that releases storage based only on block-all return could have the same ownership hazard. The real request deadline and production finite-poll behavior are unchanged.Changes
LLM_MODELS_ROOTand fail unresolved automatic KV-manager selection so the target resolves to DeepSeek V4 Pro with V2/Python, matching serving.Validation
GB300-44_GPUs-11_Nodes-PyTorch-Disagg-PerfSanity-CTX3-NODE1-GPU4-GEN1-NODE8-GPU32-Post-Merge-2.This validates the 3-CTX, concurrency-180 CI proxy. The original 8-CTX, concurrency-1760 E2E workload remains a separate validation requirement.
Related PRs
Dev Engineer Review
LLM_MODELS_ROOT.QA Engineer Review
Test-code changes include:
test_precheck_passes_mtp_exact_block_boundary.test_resolve_model_prefs_auto_requires_registered_modeltest_resolve_model_prefs_auto_propagates_model_default_failuretest_resolve_model_prefs_auto_propagates_resolver_failuretest_resolve_model_prefs_explicit_v1_does_not_require_modeltest_precheck_commands_export_model_root_safely.mtp_draft_len.No
tests/integration/test_lists/coverage entries are identified for these test functions. Verdict: needs follow-up.