[https://nvbugs/6480621][fix] Preserve KV ownership in disaggregated precheck - #17223
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 |
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: 5 (Critical) | ~90 minutes Sequence Diagram(s)sequenceDiagram
participant Precheck
participant ContextTransfer
participant GenerationTransfer
participant KVMemory
Precheck->>ContextTransfer: Dispatch and validate context transfer
Precheck->>GenerationTransfer: Dispatch and validate generation transfer
ContextTransfer->>GenerationTransfer: Coordinate ownership release
GenerationTransfer->>KVMemory: Release pages after synchronization
Precheck->>KVMemory: Retain pages and hard-abort when quiescence is unproven
Possibly related PRs
Suggested labels: 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. |
|
The receive side did not get the same fallback though, so a null |
|
Thanks for the effort! Btw, #17137 is #17223 plus one yaml, and #17121 carries the same block-verification fix as #17223 while adding |
|
PR_Github #65120 [ 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>
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>
b880173 to
cb03f67
Compare
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/others/test_cache_transceiver_precheck_run.py (1)
671-681: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the post-ownership verification path.
_gen_run_wave_runnerfixesplanto{"verify_data": False, "warmup_requests": 0}, so no test reachesgen_run_wavelines 1174-1188. That block changed in this PR: verification now runs after the ownership proof,_free_allruns before the verification consensus, and a verification failure raises_TransferErrorrather than_FatalTransferError. Add a case withverify_data: Truethat asserts the free-then-raise ordering.♻️ Proposed helper change to enable the case
- runner.plan = {"verify_data": False, "warmup_requests": 0} + runner.plan = {"verify_data": verify_data, "warmup_requests": 0}🤖 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/others/test_cache_transceiver_precheck_run.py` around lines 671 - 681, Extend the test coverage around _gen_run_wave_runner and gen_run_wave with a verify_data: True case that exercises post-ownership verification. Assert that _free_all executes before the verification consensus and that a verification failure raises _TransferError rather than _FatalTransferError, while preserving the existing helper behavior for other tests.
🤖 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/others/test_cache_transceiver_precheck_run.py`:
- Around line 671-681: Extend the test coverage around _gen_run_wave_runner and
gen_run_wave with a verify_data: True case that exercises post-ownership
verification. Assert that _free_all executes before the verification consensus
and that a verification failure raises _TransferError rather than
_FatalTransferError, while preserving the existing helper behavior for other
tests.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 10c11ef8-b0e9-4dac-b26d-f05bc63cd1da
📒 Files selected for processing (9)
tensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/transceiver.pytensorrt_llm/llmapi/llm_args.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_transceiver_bounded_polling.pytests/unittest/others/test_cache_transceiver_precheck_config.pytests/unittest/others/test_cache_transceiver_precheck_run.py
🚧 Files skipped from review as they are similar to previous changes (6)
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/_torch/disaggregation/transceiver.py
- tests/scripts/perf-sanity/cache_transceiver_precheck/README.md
- tensorrt_llm/_torch/disaggregation/native/transfer.py
- tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py
- tests/unittest/disaggregated/test_transceiver_bounded_polling.py
|
PR_Github #65362 [ run ] triggered by Bot. Commit: |
Addressed. Python V2 now rejects an unset
I suggest we land #17223 first to get the canonical ownership/deadline fix. Then #17121 could take #17223’s versions of the shared block-verification, model/runtime-resolution, and cc: @Shixiaowei02 , @chuangz0 |
|
PR_Github #65362 [ run ] completed with state
|
dpitman-nvda
left a comment
dpitman-nvda
left a comment
There was a problem hiding this comment.
Approval for infra
|
/bot run --disable-fail-fast |
|
PR_Github #65464 [ run ] triggered by Bot. Commit: |
|
PR_Github #65464 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65616 [ run ] triggered by Bot. Commit: |
|
PR_Github #65616 [ run ] completed with state |
11 participants
Summary
Fix the Python/NIXL block-all ownership contract used by the disaggregated precheck and cache-transceiver harness. An expired polling slice is not proof that transport released the KV pages.
Impact
Before this fix, these short-lived validation processes could free, reuse, or deregister KV pages while transport might still read them, causing timing-dependent payload mismatches or hangs. Normal PyExecutor nonblocking polling is unchanged; this is not evidence of corrupted serving output. An out-of-tree caller of the internal Python block-all API that recycles storage solely from its return could share the hazard.
Fix
kv_transfer_sender_future_timeout_msonly as a bounded sender wait slice andkv_transfer_timeout_msas the absolute sender deadline.TIMEOUTnonterminal and retain the session/pages.kv_transfer_timeout_ms: nullfor Python V2 before worker/CUDA setup. A receive timeout is currently terminal and closes the session, so silently substituting a receive fallback would reintroduce the ownership risk.Validation
cb03f673ee8ebb0f0b2c134d35d1d596d3b60572: 92 CPU-capable precheck tests passed; focused ownership/resolver tests passed 31 with 2 optional runtime skips.b880173d5ran 59,207 tests. All 242 changed-path tests passed; its three failures were one recurring process/infra termination and two transient host-performance regressions seen on unrelated PRs.ci: full pre-merge approvedlabel. The exact 44-GPU proxy runs on stacked PR [https://nvbugs/6480621][test] Revert to 60-second KV transfer timeout for GB300 DeepSeek V4 Pro disaggregated perf-sanity #17137 with reuse disabled.The original 8-CTX, concurrency-1760 NVBug workload remains a separate stress-validation requirement.
Related PRs
llm_models_rootcall-site changes.Dev Engineer Review
TIMEOUT.LLM_MODELS_ROOTfor enabled prechecks and aligned precheck model and runtime resolution with serving.run_cache_transceiver_test.pyto preserve ownership until quiescence is proven.kv_transfer_timeout_msis unset. Confirm that it applies the required finite fallback.#17137and#17121, including duplicatellm_models_rootkeyword handling and landing order.QA Engineer Review
test_precheck_passes_mtp_exact_block_boundary.LLM_MODELS_ROOTquoting, disabled-precheck behavior, and command parsing.LLM_MODELS_ROOTresolution.tests/integration/test_lists/were changed. Test-list coverage data is unavailable.