[None][feat] support NIXL cache transceiver with Ray - #17295
Conversation
0b89b30 to
dd29ae5
Compare
|
/bot run --disable-fail-fast |
|
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 PR adds MPI and Torch process-group topology discovery for NIXL cache agents, supports CPP or PYTHON Ray transceiver runtimes, centralizes RPC response processing, and expands unit and integration coverage. ChangesCache transport runtime support
RPC response processing
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant AgentConnectionManager
participant PgHelper
participant NixlTransferAgent
participant RemoteAgent
AgentConnectionManager->>PgHelper: Discover rank and world size
AgentConnectionManager->>NixlTransferAgent: Initialize with topology values
NixlTransferAgent->>NixlTransferAgent: Allocate port and format local address
NixlTransferAgent->>RemoteAgent: Parse endpoint and connect
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.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/integration/defs/examples/test_ray.py (1)
71-78: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd annotations and docstrings to the new test functions.
Add precise parameter types and
-> Nonereturn types. Add Google-style docstrings fortest_ray_disaggregated_serving_pythonand_run_ray_disaggregated_serving.As per coding guidelines: “Annotate every function” and “use Google-style docstrings for classes and functions.”
🤖 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/examples/test_ray.py` around lines 71 - 78, Add precise type annotations, including -> None, to test_ray_disaggregated_serving_python and _run_ray_disaggregated_serving. Add concise Google-style docstrings to both functions documenting their purpose and each parameter, using the existing parameter names and types.Source: Coding guidelines
tests/unittest/executor/test_rpc_worker_mixin.py (1)
55-64: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd coverage for the empty-response stream path.
Test coverage summary:
- Added:
test_fetch_responses_processes_and_filters_engine_responses.- Modified: none.
- Removed: none.
- CI list:
tests/integration/test_lists/test-db/l0_cpu.ymlincludes this test file.- Manual QA list: not applicable for this CPU-only unit test.
- Verdict: needs follow-up.
The test covers timeout propagation, callback filtering, temporary errors, and queue delivery. It does not verify an empty Python-transceiver response through
fetch_responses_loop_async()and confirm that polling continues until shutdown. Add that regression case.As per path instructions, test-code changes under
tests/**require changed-test, test-list, and coverage reporting. The PR objective identifies empty Python-transceiver responses as the regression target.🤖 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/executor/test_rpc_worker_mixin.py` around lines 55 - 64, Add a regression test alongside test_fetch_responses_processes_and_filters_engine_responses that exercises an empty Python-transceiver response through fetch_responses_loop_async(), verifies polling continues rather than terminating on the empty response, and confirms the loop exits only after shutdown. Update the required changed-test, test-list, and coverage reporting for this tests/** change.Source: Path instructions
🤖 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/executor/base_worker.py`:
- Around line 1125-1126: Update the temp_error_responses drain in
process_responses() to repeatedly call get_nowait() and catch queue.Empty to
terminate the loop. Remove the empty() check and blocking get() pattern,
preserving response collection while remaining safe for concurrent callers.
---
Nitpick comments:
In `@tests/integration/defs/examples/test_ray.py`:
- Around line 71-78: Add precise type annotations, including -> None, to
test_ray_disaggregated_serving_python and _run_ray_disaggregated_serving. Add
concise Google-style docstrings to both functions documenting their purpose and
each parameter, using the existing parameter names and types.
In `@tests/unittest/executor/test_rpc_worker_mixin.py`:
- Around line 55-64: Add a regression test alongside
test_fetch_responses_processes_and_filters_engine_responses that exercises an
empty Python-transceiver response through fetch_responses_loop_async(), verifies
polling continues rather than terminating on the empty response, and confirms
the loop exits only after shutdown. Update the required changed-test, test-list,
and coverage reporting for this tests/** change.
🪄 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: 3191efaa-20ad-4408-9daf-54e33603d284
📒 Files selected for processing (15)
cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cppcpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.hcpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txtcpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cppcpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.hexamples/ray_orchestrator/disaggregated/disagg_serving_local.shscripts/build_wheel.pytensorrt_llm/executor/base_worker.pytensorrt_llm/executor/rpc_worker_mixin.pytests/integration/defs/examples/test_ray.pytests/integration/test_lists/test-db/l0_cpu.ymltests/integration/test_lists/test-db/l0_dgx_b200.ymltests/integration/test_lists/test-db/l0_dgx_h100.ymltests/integration/test_lists/waives.txttests/unittest/executor/test_rpc_worker_mixin.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
|
PR_Github #64023 [ run ] triggered by Bot. Commit: |
dd29ae5 to
760f0b9
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. |
760f0b9 to
7a7ac50
Compare
|
PR_Github #64499 [ run ] completed with state
|
00ddbe9 to
07fdc46
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #64919 [ run ] triggered by Bot. Commit: |
07fdc46 to
e583d1d
Compare
|
PR_Github #64919 [ run ] completed with state
|
df5ee8f to
b31b08d
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #65186 [ run ] triggered by Bot. Commit: |
|
PR_Github #65186 [ run ] completed with state
|
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
…dback Use RFC 3986 bracketed IPv6 literals for the NIXL agent address so IPv4 keeps the legacy ip:port format and stays compatible across versions. Drain deferred error responses with get_nowait() to avoid blocking when the ManagedThread and RPC fetch_responses() race on the same queue. Document the NIXL backend and --transceiver_runtime option in the Ray disaggregated-serving example README. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
b31b08d to
d07e811
Compare
|
/bot run --stage-list "DGX_H100-PyTorch-6, RTXPro6000D-PyTorch-1, GB300-4_GPUs-PyTorch-1" |
|
PR_Github #65439 [ run ] triggered by Bot. Commit: |
|
PR_Github #65439 [ run ] completed with state
|
|
/bot skip --comment "all test have passed" |
|
PR_Github #65502 [ skip ] triggered by Bot. Commit: |
|
PR_Github #65502 [ skip ] completed with state |
8 participants
Summary
TRTLLM_NIXL_PORTfixed-port override; NIXL listeners now always select an available portAwaitResponseHelper.process_responses()implementation between regular and Ray RPC workersWhy
The NIXL transfer agent and agent connection manager assumed MPI for rank discovery and collective metadata exchange. Ray workers disable MPI and use a Torch process group, so NIXL cache transfer could not initialize correctly in that environment. In addition, the tp1 test was skipped under an unrelated bug and tp2 remained waived.
The undocumented
TRTLLM_NIXL_PORToverride was removed intentionally. Its rank-derived fixed-port scheme was not safe across independent context and generation sessions, where session-local ranks can overlap. NIXL listeners now always select an available port and advertise the selected endpoint through exchanged agent metadata.Local-address selection is also an intentional behavior change for all NIXL deployments, not only the Ray/Torch process-group path. All listening NIXL agents now use the shared
common::getLocalIp()selection withTRTLLM_NIXL_INTERFACEas the explicit override; when no interface is configured, the shared default-route/hostname-based selection is used. IPv6 literals are advertised as[address]:portso the address and dynamically selected port remain unambiguous.The Python-transceiver first-token flow can produce an empty response (
None) for an iteration with no client-visible serialized result. The regular worker path already runs engine responses through_engine_response_callback, filters non-client responses, and appends deferred submission errors before IPC handling.RpcWorkerMixin.fetch_responses()bypassed that existing preprocessing and passed the rawNoneto IPC, whereresponse.has_error()terminated the Ray RPC stream.The existing callback/filter/deferred-error logic is factored into
AwaitResponseHelper.process_responses(). Both the regular worker and Ray RPC worker call that shared function beforeresponses_handler(). No PyExecutor or placement-group behavior is changed by this fix.This follows the intent of #10259 while adapting it to the current cache-transceiver and ProcessGroup implementation.
Validation
libtensorrt_llm.so, NIXL wrapper, transfer-agent binding, executor, and UCX wrapper--no-depsRUNPATH=$ORIGIN:$ORIGIN/nixl/, andlddresolveslibpg_utils.sofrom the same wheelThe capital of Germany is Berlin,Asyncio is a Python library, and noBerlin Berlin) passed, with noNoneType.has_error,RPCStreamingError, orray_executor_main_loopfailure; context/generation shared one physical RTX A6000 through a temporary fractional placement adjustment that is not part of this PRThe B200/B300 tp1/tp2 CPP/PYTHON E2E cases should be rerun in CI with the final narrowed diff.
Dev Engineer Review
rankandworld_sizefields toBaseAgentConfigand both Python and C++ NIXL agent wrappers.QA Engineer Review
test_ray_disaggregated_serving_python.test_ray_disaggregated_servingto use runtime parameters.RpcWorkerMixin.fetch_responsesunit coverage.BaseAgentConfigbinding tests.l0_cpu.yml.l0_dgx_b200.ymlandl0_dgx_h100.yml.waives.txt.