[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 5 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 |
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 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. |
| for task in self.kv_tasks: | ||
| if not task.wait(timeout=self._timeout_s): | ||
| return WaitResult.TIMEOUT | ||
| while not task.wait(timeout=wait_slice_s): |
There was a problem hiding this comment.
This turns blocking=True from bounded into unbounded, for every caller — not just the precheck.
Before, if not task.wait(timeout=self._timeout_s): return WaitResult.TIMEOUT gave the caller control back within one slice. Now the only exits are task completion and has_failed(), which is ERROR/CANCELLED only. The sender worker marks the task TRANSFERRING and then waits on NIXL's status with no timeout, so a transfer wedged in PENDING/PROCESSING never calls complete() or fail() — nothing sets a terminal status, and this loop spins until the process dies. The old code degraded to a TIMEOUT the caller could act on; this one has no escape.
transceiver.py:698 is the concrete one: with at_least_request_num=None it used to get TIMEOUT, keep the session, and poll again. I agree the ordinary PyExecutor path polls with 0/1 and is unaffected — but block-all is a supported path, and blocking=True is still this method's default.
The ownership fix itself is right and well argued; it's the removal of the escape hatch that I'd want separated. Keeping a real deadline — the request-level kv_transfer_timeout_ms you correctly distinguish from the polling interval — would give you both: retry across slices so blockAll can't return early, and still fail rather than hang when the peer never quiesces.
| elif result == WaitResult.TIMEOUT: | ||
| logger.warning( | ||
| f"TxSession rid={session.disagg_request_id} timed out after {self._sender_future_timeout_ms}ms" | ||
| logger.debug( |
There was a problem hiding this comment.
Worth checking before you re-word this: after the change above, nothing in tensorrt_llm/ returns WaitResult.TIMEOUT any more. git grep -n 'WaitResult.TIMEOUT' -- tensorrt_llm/ at this head matches only this consumer. RxSession.wait_complete maps its timeout to FAILED, and TxSession.wait_complete no longer produces one on either path.
So this elif is dead, and the new "keeping it in progress" wording describes behaviour that can't be reached — the session is now kept in progress by blocking inside wait_complete, not by returning here. Either drop the branch and the timed_out bookkeeping it feeds, or keep a bounded return alive so it means something again. Reads to me as a symptom of the first point rather than a separate bug.
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>
f9ca117 to
f807789
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. |
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
tests/unittest/disaggregated/test_cache_transceiver_harness.py (1)
61-110: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the shared driver loader.
_load_driver_ownership_helpersand_load_driver_request_flowdiffer only in the selected names, the injected stubs, and the module name. The parsing, selection, compile, exec, and missing-name assertion are identical. Duplicating them makes the two name sets drift, which is exactly the failure described in the previous comment.♻️ Proposed refactor
+def _load_driver_subset( + module_name: str, + selected_names: set[str], + stubs: dict, +) -> types.ModuleType: + """Execute selected driver top-level definitions in an isolated module.""" + source = Path(DRIVER_SCRIPT).read_text() + tree = ast.parse(source, filename=DRIVER_SCRIPT) + selected = [ + node + for node in tree.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name in selected_names + ] + module = types.ModuleType(module_name) + module.__dict__.update(stubs) + exec( + compile(ast.Module(body=selected, type_ignores=[]), DRIVER_SCRIPT, "exec"), + module.__dict__, + ) + missing = selected_names - set(module.__dict__) + assert not missing, f"{module_name}: driver definitions not loaded: {sorted(missing)}" + return moduleAlso applies to: 113-172
🤖 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_cache_transceiver_harness.py` around lines 61 - 110, Extract the shared AST-based loading logic from _load_driver_ownership_helpers and _load_driver_request_flow into a reusable loader helper. Parameterize it with the selected symbols, injected module globals/stubs, and module name, while preserving each caller’s missing-name validation and behavior. Update both loaders to delegate to this helper so their name sets and execution paths cannot drift.
🤖 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 `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Around line 1378-1437: Update the wait_complete blocking path around
wait_slice_s and _deadline_monotonic_s so an explicit None tx_overall_timeout_s
cannot cause an unbounded wait. Use a finite fallback overall deadline (or
validate and reject None) while preserving the existing bounded wait-slice and
terminal-state handling in wait_for_task.
In `@tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py`:
- Around line 196-201: Update the enabled guard in the precheck configuration
flow to reject both None and empty llm_models_root values before inserting the
LLM_MODELS_ROOT export. Preserve the existing ValueError message and valid
non-empty path behavior.
In `@tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py`:
- Around line 389-397: Update the auto-setting path around
model_cls.get_model_defaults to avoid passing None; provide a valid TorchLlmArgs
instance containing the required fields, including tensor_parallel_size, before
invoking the hook. Preserve the existing RuntimeError wrapping for genuine hook
failures.
In `@tests/unittest/disaggregated/test_cache_transceiver_harness.py`:
- Around line 63-76: Add "_Timeout" to the selected_names set used by
cache_transceiver_harness_ownership and add the corresponding _Timeout entry to
its stub dictionary, keeping both ownership helper symbol sets consistent with
_load_driver_request_flow.
---
Nitpick comments:
In `@tests/unittest/disaggregated/test_cache_transceiver_harness.py`:
- Around line 61-110: Extract the shared AST-based loading logic from
_load_driver_ownership_helpers and _load_driver_request_flow into a reusable
loader helper. Parameterize it with the selected symbols, injected module
globals/stubs, and module name, while preserving each caller’s missing-name
validation and behavior. Update both loaders to delegate to this helper so their
name sets and execution paths cannot drift.
🪄 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: 4feb90e5-8309-4154-bc46-b76e62aa9189
📒 Files selected for processing (16)
examples/disaggregated/slurm/cache_transceiver_test/README.mdexamples/disaggregated/slurm/cache_transceiver_test/config.yamlexamples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.pyjenkins/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_harness.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
🚧 Files skipped from review as they are similar to previous changes (4)
- jenkins/scripts/perf/local/submit.py
- tests/unittest/scripts/test_perf_submit.py
- jenkins/scripts/perf/submit.py
- tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py
| # ``_timeout_s`` bounds one scheduler wait slice. The separate absolute | ||
| # deadline is shared by every KV task and aux; it is never reset by a | ||
| # later wait_complete() call. | ||
| wait_slice_s = self._timeout_s | ||
| if wait_slice_s is None or wait_slice_s <= 0: | ||
| wait_slice_s = _FALLBACK_TX_WAIT_SLICE_S | ||
|
|
||
| def wait_for_task(task: SendTaskBase) -> WaitResult: | ||
| while True: | ||
| # A task/session terminal state observed at the deadline | ||
| # boundary takes precedence over TIMEOUT. | ||
| if self.has_failed(): | ||
| return WaitResult.FAILED | ||
| if task.status == TaskStatus.TRANSFERRED: | ||
| return WaitResult.COMPLETED | ||
|
|
||
| remaining_s = None | ||
| if self._deadline_monotonic_s is not None: | ||
| remaining_s = self._deadline_monotonic_s - time.monotonic() | ||
| if remaining_s <= 0: | ||
| # The worker can publish terminal state between the | ||
| # checks above and the clock read. Preserve boundary | ||
| # precedence before classifying this as a timeout. | ||
| if self.has_failed(): | ||
| return WaitResult.FAILED | ||
| if task.status == TaskStatus.TRANSFERRED: | ||
| return WaitResult.COMPLETED | ||
| return WaitResult.TIMEOUT | ||
| timeout_s = wait_slice_s if remaining_s is None else min(wait_slice_s, remaining_s) | ||
| task.wait(timeout=timeout_s) | ||
|
|
||
| # A bounded slice keeps cancellation and sibling failure observable. | ||
| for task in self.kv_tasks: | ||
| if not task.wait(timeout=self._timeout_s): | ||
| return WaitResult.TIMEOUT | ||
| if task.status == TaskStatus.ERROR: | ||
| return WaitResult.FAILED | ||
| if self._need_aux and self.aux_task is not None: | ||
| if not self.aux_task.wait(timeout=self._timeout_s): | ||
| return WaitResult.TIMEOUT | ||
| if self.aux_task.status == TaskStatus.ERROR: | ||
| result = wait_for_task(task) | ||
| if result != WaitResult.COMPLETED: | ||
| return result | ||
| if self._need_aux: | ||
| if self.aux_task is None: | ||
| # _finalize_send() installs the aux task synchronously before | ||
| # publishing the request to _send_reqs. Once every KV task is | ||
| # terminal, a missing required aux task is an invariant error, | ||
| # not an asynchronously pending transfer. | ||
| with self.lock: | ||
| if self._terminal_status not in ( | ||
| SessionStatus.ERROR, | ||
| SessionStatus.CANCELLED, | ||
| ): | ||
| self._exception = RuntimeError( | ||
| "required auxiliary transfer was not dispatched" | ||
| ) | ||
| self._terminal_status = SessionStatus.ERROR | ||
| return WaitResult.FAILED | ||
| return WaitResult.COMPLETED | ||
| result = wait_for_task(self.aux_task) | ||
| if result != WaitResult.COMPLETED: | ||
| return result | ||
| return ( | ||
| WaitResult.FAILED | ||
| if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED) | ||
| else WaitResult.COMPLETED | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Find the default value and validation of kv_transfer_timeout_ms.
rg -nP -C4 'kv_transfer_timeout_ms' --type=py --type=cpp --type=cc
rg -nP -C4 'kv_transfer_timeout_ms' -g '*.h' -g '*.hpp' -g '*.pyi'Repository: NVIDIA/TensorRT-LLM
Length of output: 188
🏁 Script executed:
#!/bin/bash
set -eu
echo '--- matching files ---'
rg -n --glob '*.py' --glob '*.cpp' --glob '*.cc' --glob '*.h' --glob '*.hpp' \
'kv_transfer_timeout_ms|tx_overall_timeout_s|_overall_timeout_s|_deadline_monotonic_s' .
echo '--- transfer.py structure and target region ---'
ast-grep outline tensorrt_llm/_torch/disaggregation/native/transfer.py
sed -n '1280,1460p' tensorrt_llm/_torch/disaggregation/native/transfer.py
echo '--- repository configuration references ---'
rg -n -C 5 --glob '*.yaml' --glob '*.yml' --glob '*.json' --glob '*.toml' --glob '*.py' \
'kv_transfer_timeout_ms|tx_overall_timeout_s' .Repository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -eu
echo '--- configuration definitions ---'
sed -n '4190,4275p' tensorrt_llm/llmapi/llm_args.py
sed -n '35,130p' tensorrt_llm/_torch/disaggregation/transceiver.py
sed -n '1200,1285p' tensorrt_llm/_torch/disaggregation/native/transfer.py
sed -n '2290,2370p' tensorrt_llm/_torch/disaggregation/native/transfer.py
sed -n '455,505p' cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp
sed -n '395,425p' cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
echo '--- focused bounded-polling tests ---'
sed -n '110,180p' tests/unittest/disaggregated/test_transceiver_bounded_polling.py
sed -n '620,760p' tests/unittest/disaggregated/test_transceiver_bounded_polling.py
echo '--- explicit unset/default construction sites ---'
sed -n '430,465p' tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py
sed -n '610,665p' tests/unittest/others/test_kv_cache_transceiver.py
echo '--- read-only static verifier ---'
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("tensorrt_llm/llmapi/llm_args.py")
tree = ast.parse(path.read_text())
for node in ast.walk(tree):
if isinstance(node, ast.AnnAssign) and getattr(node.target, "id", None) == "kv_transfer_timeout_ms":
print("kv_transfer_timeout_ms annotation:", ast.unparse(node.annotation))
print("kv_transfer_timeout_ms value:", ast.unparse(node.value))
break
path = Path("tensorrt_llm/_torch/disaggregation/transceiver.py")
text = path.read_text()
print("passes None when config value is None:",
"else None" in text[text.index("transfer_timeout_s"):text.index("transfer_timeout_s") + 300])
print("passes transfer_timeout_s as tx_overall_timeout_s:",
"tx_overall_timeout_s=transfer_timeout_s" in text)
path = Path("tensorrt_llm/_torch/disaggregation/native/transfer.py")
text = path.read_text()
wait_start = text.index(" def wait_complete(self, blocking: bool = True)")
wait_end = text.index(" def set_exception", wait_start)
wait_body = text[wait_start:wait_end]
print("wait slice fallback exists:", "_FALLBACK_TX_WAIT_SLICE_S" in wait_body)
print("deadline check conditional:", "if self._deadline_monotonic_s is not None:" in wait_body)
print("unbounded branch uses task.wait:", "task.wait(timeout=timeout_s)" in wait_body)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 31639
🏁 Script executed:
#!/bin/bash
set -eu
echo '--- current change for the reviewed implementation ---'
git diff -- tensorrt_llm/_torch/disaggregation/native/transfer.py \
tensorrt_llm/_torch/disaggregation/transceiver.py \
tensorrt_llm/llmapi/llm_args.py \
tests/unittest/disaggregated/test_transceiver_bounded_polling.py
echo '--- wait_complete-focused tests ---'
rg -n -C 8 'wait_complete|WaitResult\.TIMEOUT|deadline_monotonic|no overall|overall timeout|unbounded' \
tests/unittest/disaggregated/test_transceiver_bounded_polling.py \
tests/unittest -g '*.py'
echo '--- all construction paths for the Python V2 transceiver ---'
rg -n -C 8 'KvCacheTransceiverV2|CacheTransceiverConfig\(' \
tensorrt_llm/_torch/disaggregation tensorrt_llm/_torch/pyexecutor tests/unittest/disaggregated \
tests/unittest/_torch/executor -g '*.py'Repository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -eu
echo '--- task wait semantics ---'
sed -n '210,280p' tensorrt_llm/_torch/disaggregation/native/transfer.py
echo '--- exact deadline and fallback tests ---'
sed -n '450,625p' tests/unittest/disaggregated/test_transceiver_bounded_polling.py
sed -n '800,885p' tests/unittest/disaggregated/test_transceiver_bounded_polling.py
echo '--- focused diff summary ---'
git diff --stat -- tensorrt_llm/_torch/disaggregation/native/transfer.py \
tensorrt_llm/_torch/disaggregation/transceiver.py \
tensorrt_llm/llmapi/llm_args.py \
tests/unittest/disaggregated/test_transceiver_bounded_polling.py
git diff -U12 -- tensorrt_llm/_torch/disaggregation/native/transfer.py | sed -n '1,260p'Repository: NVIDIA/TensorRT-LLM
Length of output: 10984
Bound wait_complete(blocking=True) when kv_transfer_timeout_ms=None.
CacheTransceiverConfig defaults to 60000, but callers can explicitly set None, and KvCacheTransceiverV2 forwards it as tx_overall_timeout_s=None. A stalled peer can then block the caller indefinitely. Add a finite fallback deadline or reject None for this path.
🤖 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 1378 -
1437, Update the wait_complete blocking path around wait_slice_s and
_deadline_monotonic_s so an explicit None tx_overall_timeout_s cannot cause an
unbounded wait. Use a finite fallback overall deadline (or validate and reject
None) while preserving the existing bounded wait-slice and terminal-state
handling in wait_for_task.
| if enabled: | ||
| if llm_models_root is None: | ||
| raise ValueError("enabled cache-transceiver precheck requires LLM_MODELS_ROOT") | ||
| # Keep this as a top-level assignment. shlex.quote() is not safe when | ||
| # nested inside the double-quoted pytestCommand exports below. | ||
| lines.insert(0, f"export LLM_MODELS_ROOT={shlex.quote(llm_models_root)}") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject an empty llm_models_root too.
The guard only rejects None. An empty string passes and produces export LLM_MODELS_ROOT=''. The precheck then resolves no model directory and silently falls back to FALLBACK_KV_SHAPE, which defeats the fail-fast intent of this check.
🛡️ Proposed fix
if enabled:
- if llm_models_root is None:
+ if not llm_models_root:
raise ValueError("enabled cache-transceiver precheck requires LLM_MODELS_ROOT")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if enabled: | |
| if llm_models_root is None: | |
| raise ValueError("enabled cache-transceiver precheck requires LLM_MODELS_ROOT") | |
| # Keep this as a top-level assignment. shlex.quote() is not safe when | |
| # nested inside the double-quoted pytestCommand exports below. | |
| lines.insert(0, f"export LLM_MODELS_ROOT={shlex.quote(llm_models_root)}") | |
| if enabled: | |
| if not llm_models_root: | |
| raise ValueError("enabled cache-transceiver precheck requires LLM_MODELS_ROOT") | |
| # Keep this as a top-level assignment. shlex.quote() is not safe when | |
| # nested inside the double-quoted pytestCommand exports below. | |
| lines.insert(0, f"export LLM_MODELS_ROOT={shlex.quote(llm_models_root)}") |
🤖 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/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py`
around lines 196 - 201, Update the enabled guard in the precheck configuration
flow to reject both None and empty llm_models_root values before inserting the
LLM_MODELS_ROOT export. Preserve the existing ValueError message and valid
non-empty path behavior.
| if setting == "auto": | ||
| defaults = {} | ||
| if model_cls is not None: | ||
| if hasattr(model_cls, "get_model_defaults"): | ||
| try: | ||
| defaults = model_cls.get_model_defaults(None) or {} | ||
| except Exception as e: # noqa: BLE001 - model hooks may need llm_args | ||
| print( | ||
| f"[precheck] WARNING: get_model_defaults failed ({e!r}); assuming V1", | ||
| flush=True, | ||
| ) | ||
| except Exception as e: # noqa: BLE001 - model hooks are third-party extension points | ||
| raise RuntimeError( | ||
| f"get_model_defaults failed for {model_cls.__name__}; refusing to assume V1" | ||
| ) from e |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Inspect every get_model_defaults override for llm_args usage.
ast-grep run --pattern $'def get_model_defaults($$$) {
$$$
}' --lang python tensorrt_llm 2>/dev/null
rg -nP -A15 'def get_model_defaults\s*\(' --type=py tensorrt_llmRepository: NVIDIA/TensorRT-LLM
Length of output: 14053
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- precheck call site ---'
sed -n '350,420p' tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
printf '%s\n' '--- all definitions and references ---'
rg -n -P 'def get_model_defaults\s*\(|get_model_defaults\s*\(' tensorrt_llm tests/scripts/perf-sanity/cache_transceiver_precheck --type py
printf '%s\n' '--- registered model resolution and Whisper registration ---'
rg -n -P 'Whisper|MODEL|model_cls|ModelLoader|AutoModel|model_type' tensorrt_llm/_torch tests/scripts/perf-sanity/cache_transceiver_precheck --type py | head -300
printf '%s\n' '--- static analysis of every definition ---'
python3 - <<'PY'
import ast
from pathlib import Path
root = Path("tensorrt_llm")
for path in sorted(root.rglob("*.py")):
try:
tree = ast.parse(path.read_text())
except (OSError, SyntaxError):
continue
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "get_model_defaults":
args = [a.arg for a in node.args.args]
llm_arg = args[-1] if args else None
uses = []
for child in ast.walk(node):
if child is node:
continue
if isinstance(child, ast.Attribute) and isinstance(child.value, ast.Name) and child.value.id == llm_arg:
uses.append(f"attribute:{child.attr}")
elif isinstance(child, ast.Subscript) and isinstance(child.value, ast.Name) and child.value.id == llm_arg:
uses.append("subscript")
elif isinstance(child, ast.Call) and isinstance(child.func, ast.Name) and child.func.id == llm_arg:
uses.append("call")
print(f"{path}:{node.lineno} arg={llm_arg!r} uses={sorted(set(uses))}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 42694
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- Whisper class registration and mapping ---'
sed -n '820,875p' tensorrt_llm/_torch/models/modeling_whisper.py
sed -n '835,865p' tensorrt_llm/_torch/models/modeling_utils.py
rg -n -P 'register_model|MODEL_CLASS_MAPPING|WhisperForConditionalGeneration' tensorrt_llm/_torch/models/modeling_whisper.py tensorrt_llm/_torch/models/__init__.py tensorrt_llm/_torch/models/modeling_utils.py
printf '%s\n' '--- model loader argument construction ---'
sed -n '400,440p' tensorrt_llm/_torch/pyexecutor/model_loader.py
rg -n -P 'class TorchLlmArgs|TorchLlmArgs\(' tensorrt_llm/_torch --type py | head -80Repository: NVIDIA/TensorRT-LLM
Length of output: 6184
Do not call get_model_defaults() with None.
The registered WhisperForConditionalGeneration override reads llm_args.tensor_parallel_size. This raises AttributeError and aborts the precheck when Whisper uses "auto". Pass a valid TorchLlmArgs or make the hook accept None.
🤖 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/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py` around
lines 389 - 397, Update the auto-setting path around
model_cls.get_model_defaults to avoid passing None; provide a valid TorchLlmArgs
instance containing the required fields, including tensor_parallel_size, before
invoking the hook. Preserve the existing RuntimeError wrapping for genuine hook
failures.
| selected_names = { | ||
| "_TransferError", | ||
| "_FatalTransferError", | ||
| "_request_ids", | ||
| "_context_completion_error", | ||
| "_gen_completion_error", | ||
| "_can_release_sequence", | ||
| "_release_sequence_if_safe", | ||
| "_validate_context_completion", | ||
| "_validate_python_gen_completion", | ||
| "_first_reason", | ||
| "_exchange_release_decision", | ||
| "_hard_abort_process", | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Add _Timeout to the ownership helper name set.
_exchange_release_decision contains except _Timeout: raise. _Timeout is neither in selected_names nor in the stub dict at lines 85-101, so the name is undefined in cache_transceiver_harness_ownership. The current tests pass because no exception is raised inside that try. If the handshake body ever raises, Python evaluates the except clause and fails with NameError instead of the intended handling, which hides the real failure.
_load_driver_request_flow already selects _Timeout; keep the two sets consistent.
🛠️ Proposed fix
selected_names = {
+ "_Timeout",
"_TransferError",
"_FatalTransferError",📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| selected_names = { | |
| "_TransferError", | |
| "_FatalTransferError", | |
| "_request_ids", | |
| "_context_completion_error", | |
| "_gen_completion_error", | |
| "_can_release_sequence", | |
| "_release_sequence_if_safe", | |
| "_validate_context_completion", | |
| "_validate_python_gen_completion", | |
| "_first_reason", | |
| "_exchange_release_decision", | |
| "_hard_abort_process", | |
| } | |
| selected_names = { | |
| "_Timeout", | |
| "_TransferError", | |
| "_FatalTransferError", | |
| "_request_ids", | |
| "_context_completion_error", | |
| "_gen_completion_error", | |
| "_can_release_sequence", | |
| "_release_sequence_if_safe", | |
| "_validate_context_completion", | |
| "_validate_python_gen_completion", | |
| "_first_reason", | |
| "_exchange_release_decision", | |
| "_hard_abort_process", | |
| } |
🤖 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_cache_transceiver_harness.py` around lines
63 - 76, Add "_Timeout" to the selected_names set used by
cache_transceiver_harness_ownership and add the corresponding _Timeout entry to
its stub dictionary, keeping both ownership helper symbol sets consistent with
_load_driver_request_flow.
|
PR_Github #64468 [ run ] triggered by Bot. Commit: |
|
PR_Github #64468 [ run ] completed with state
|
5 participants
Summary
Prevent Python/NIXL
blockAllcallers from releasing or deregistering KV pages before transfer quiescence is proven, while still enforcing the request-level transfer deadline.Impact
kv_transfer_sender_future_timeout_msis a polling slice, not the request deadline. The old precheck treated one expired slice as completion and recycled source pages while NIXL could still be reading them, which can produce timing-dependent payload mismatches.Premature page reuse is a serious ownership violation. The demonstrated in-tree impact is limited to the disaggregated perf-sanity precheck and standalone cache-transceiver harness, which run in short-lived test processes before serving. Normal PyExecutor finite status polling is unchanged, and there is no evidence from this path of corrupted user output. An out-of-tree caller of this internal Python
blockAllAPI that recycles storage from its return alone could have the same hazard.Fix
kv_transfer_timeout_ms; keep the 1-second future timeout only as a polling/cancellation slice.TIMEOUTat true expiry and retain the session/pages. A terminal result observed at the boundary still wins.Validation
GB300-44_GPUs-11_Nodes-PyTorch-Disagg-PerfSanity-CTX3-NODE1-GPU4-GEN1-NODE8-GPU32-Post-Merge-2with test reuse disabled.The original 8-CTX, concurrency-1760 NVBug workload remains a separate stress-validation requirement.
Related PRs