[TRTLLM-14010][feat] report KV cache transfer state on executor hangs - #16300
[TRTLLM-14010][feat] report KV cache transfer state on executor hangs#16300bo-nv wants to merge 5 commits into
Conversation
e878d4c to
022bdc4
Compare
|
/bot run --disable-fail-fast --stage-list "PerfSanity" |
|
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:
WalkthroughCacheTransceiver now publishes synchronized transfer status snapshots and exposes diagnostic output through C++ and Python interfaces. PyExecutor registers transceiver diagnostics with HangDetector, which logs provider output when a hang is detected. ChangesCacheTransceiver Hang Diagnostics
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant HangDetector
participant KvCacheTransceiver
participant CacheTransceiver
PyExecutor->>HangDetector: register transceiver status provider
HangDetector->>KvCacheTransceiver: request status dump on hang
KvCacheTransceiver->>CacheTransceiver: delegate bound implementation request
CacheTransceiver-->>KvCacheTransceiver: return formatted transfer snapshot
KvCacheTransceiver-->>HangDetector: return diagnostic string
HangDetector-->>PyExecutor: log status during hang handling
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 |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/integration/defs/perf/test_perf_sanity.py (2)
1507-1534: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant failure-status write.
_write_server_failureruns in theexceptblock (Line 1509) and again infinally(Line 1526), sincefailure_messageis set. Theexceptwrite is immediately superseded by thefinallywrite (which also captures the post-terminationreturn_code). Consider dropping theexcept-block write and keeping only thefinallypath to avoid the double I/O and duplicated first-error scan.🤖 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/integration/defs/perf/test_perf_sanity.py` around lines 1507 - 1534, Remove the _write_server_failure call from the except block and retain the finally-block invocation after server cleanup, preserving failure_message assignment and exception propagation so the single write includes the post-termination return_code.
1521-1524: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
server_proc.wait()afterterminate()has no timeout and can hang the harness.If a server ignores
SIGTERM,wait()blocks indefinitely, stalling the test until the CI-level timeout. Bound the wait and escalate tokill().♻️ Proposed fix
if server_proc is not None: if server_proc.poll() is None: server_proc.terminate() - server_proc.wait() + try: + server_proc.wait(timeout=30) + except subprocess.TimeoutExpired: + server_proc.kill() + server_proc.wait()🤖 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/integration/defs/perf/test_perf_sanity.py` around lines 1521 - 1524, Update the server_proc cleanup block to bound the wait after terminate(), using a timeout and escalating to server_proc.kill() when the process does not exit promptly; then ensure the killed process is reaped without an unbounded wait.
🤖 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/integration/defs/perf/test_perf_sanity.py`:
- Around line 1507-1534: Remove the _write_server_failure call from the except
block and retain the finally-block invocation after server cleanup, preserving
failure_message assignment and exception propagation so the single write
includes the post-termination return_code.
- Around line 1521-1524: Update the server_proc cleanup block to bound the wait
after terminate(), using a timeout and escalating to server_proc.kill() when the
process does not exit promptly; then ensure the killed process is reaped without
an unbounded wait.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 62d726ed-d58f-4f53-8b35-4b0475ec08f1
📒 Files selected for processing (1)
tests/integration/defs/perf/test_perf_sanity.py
|
PR_Github #58900 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast --stage-list "*GB200*PerfSanity*,*GB300*PerfSanity*" |
|
PR_Github #58900 [ run ] completed with state
|
|
PR_Github #58909 Bot args parsing error: usage: /bot [-h] |
|
/bot run --disable-fail-fast --stage-list "GB200PerfSanity*,GB300PerfSanity*" |
|
PR_Github #58914 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast --stage-list "*GB200*PerfSanity*,*GB300*PerfSanity*" |
|
PR_Github #58916 [ run ] triggered by Bot. Commit: |
|
PR_Github #58914 [ run ] completed with state |
|
/bot run --disable-fail-fast --stage-list "*GB200*PerfSanity*,*GB300*PerfSanity*" |
|
PR_Github #58927 [ run ] triggered by Bot. Commit: |
|
PR_Github #58916 [ run ] completed with state |
|
PR_Github #58927 [ run ] completed with state
|
5982080 to
08a66e7
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59170 [ run ] triggered by Bot. Commit: |
|
PR_Github #59170 [ run ] completed with state
|
c7b5be2 to
a310022
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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/transceiver.py`:
- Line 174: Make access to _send_sessions and _recv_sessions thread-safe by
introducing or reusing one shared lock for all mutations and the session
snapshot in get_status_dump(). Acquire that lock while copying sessions.items()
and while executor/transfer code adds, removes, or updates entries, preserving
the diagnostic snapshot without RuntimeError during concurrent changes.
In `@tensorrt_llm/_torch/pyexecutor/hang_detector.py`:
- Around line 117-121: Add explicit -> None return annotations to
register_status_provider and the async _detect_hang method, preserving their
existing behavior and signatures otherwise.
- Around line 123-130: Update the hang-detection flow around the lock and
status-provider loop to snapshot _status_providers while holding self.lock, then
invoke each provider outside the lock with a bounded timeout so blocked or
native calls cannot delay the hard-kill path. Ensure print_all_stacks() and
on_detected() execute independently of provider completion, while preserving
status logging for providers that return within the timeout.
- Around line 131-132: Replace the broad silent handler around provider
diagnostics in tensorrt_llm/_torch/pyexecutor/hang_detector.py:131-132 with
logging of each provider failure while continuing to process remaining
providers. In tensorrt_llm/_torch/disaggregation/transceiver.py:178-192, catch
only the documented property-access exceptions and log unexpected session-state
failures instead of suppressing them.
🪄 Autofix (Beta)
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: cc870216-9cc0-4b1d-9eb4-3bfc68445f4e
📒 Files selected for processing (7)
cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.hcpp/tensorrt_llm/batch_manager/cacheTransceiver.cppcpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpptensorrt_llm/_torch/disaggregation/transceiver.pytensorrt_llm/_torch/pyexecutor/hang_detector.pytensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.pytensorrt_llm/_torch/pyexecutor/py_executor.py
🚧 Files skipped from review as they are similar to previous changes (3)
- cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp
- cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h
- cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
Signed-off-by: Bo Deng <deemod@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #61563 [ run ] triggered by Bot. Commit: |
|
PR_Github #61563 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/hang_detector.py (1)
117-120: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the
providerargument in this public API.Add a Google-style
Argssection describing the nonblocking callable and its string return contract.Proposed docstring update
def register_status_provider(self, provider: Callable[[], str]) -> None: - """Register a nonblocking callable that returns status to dump on hang detection.""" + """Register a nonblocking status provider. + + Args: + provider: Callable returning diagnostic status for hang reporting. + """🤖 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/pyexecutor/hang_detector.py` around lines 117 - 120, Update the docstring of HangDetector.register_status_provider to include a Google-style Args section for provider, stating that it is a nonblocking callable returning a status string for hang-detection dumps.Source: Coding guidelines
🤖 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/pyexecutor/hang_detector.py`:
- Around line 137-140: Update the hang-detection flow around print_all_stacks()
to set _detected before emitting diagnostics, then guarantee on_detected() runs
in a finally block even when stack formatting or logging raises. Also bound the
stack-dumping/logging operation so a stalled logger cannot prevent hard-kill
propagation.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/hang_detector.py`:
- Around line 117-120: Update the docstring of
HangDetector.register_status_provider to include a Google-style Args section for
provider, stating that it is a nonblocking callable returning a status string
for hang-detection dumps.
🪄 Autofix (Beta)
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: b1fd8af5-1f79-47dd-a52b-b8d7c090a8d7
📒 Files selected for processing (1)
tensorrt_llm/_torch/pyexecutor/hang_detector.py
|
PR_Github #61829 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #61836 [ run ] triggered by Bot. Commit: |
|
PR_Github #61829 [ run ] completed with state |
|
PR_Github #61836 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62063 [ run ] triggered by Bot. Commit: |
|
PR_Github #62063 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62188 [ run ] triggered by Bot. Commit: |
|
PR_Github #62188 [ run ] completed with state |
|
/bot run |
|
PR_Github #62379 [ run ] triggered by Bot. Commit: |
|
PR_Github #62379 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62403 [ run ] triggered by Bot. Commit: |
Summary
Status output
V1 / C++ transceiver:
V2 / Python NIXL transceiver:
Dev Engineer Review
CacheTransceivernow maintains a synchronizedStatusSnapshot(mutex-guarded, with an atomic “sync requester active” counter) and exposesgetStatusDump()for a consistent TX/RX/transfer-buffer status view.SyncRequesterStatusGuardpluspublishStatusSnapshot() noexceptto refresh snapshot data at key transfer checkpoints and around consensus/terminal cleanup.noexcept, swallowsstd::system_error) to avoid disrupting transfer flows._wait_reqssize.HangDetectornow supports registration of caller-provided status providers; failures in individual providers are caught and logged so hang handling/stack reporting andon_detectedstill proceed.getStatusDump()is surfaced to Python viaKvCacheTransceiver.get_status_dump()/BindKvCacheTransceiver.get_status_dump(), and then registered withHangDetectorfromPyExecutorwhen a transceiver is present.<nanobind/stl/string.h>and reformats a constructor default.tests/unittest/_torch/executor/test_hang_detector_kill.pyis listed intests/integration/test_lists/test-db/l0_sanity_check.yml(entry at line ~41).QA Engineer Review
tests/unittest/_torch/executor/test_hang_detector_kill.py::test_status_provider_errors_are_loggedtests/integration/test_lists/test-db/l0_sanity_check.ymlentry forunittest/_torch/executor/test_hang_detector_kill.py.Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.