[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 #64186 [ run ] completed with state
|
a1e2713 to
7d06d2b
Compare
|
/bot run --disable-fail-fast |
|
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. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py (1)
104-109: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the assertion from input validation.
Line 104 already proves that
world_sizeis notNonewhenrankis set. Delete the redundant assertion. Keep explicitValueErrorchecks for invalid input.🤖 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/nixl/_agent_py.py` around lines 104 - 109, Remove the redundant assert world_size is not None from the rank/world_size validation block, since the preceding condition guarantees both are provided together. Preserve the existing ValueError checks for invalid rank and world_size values.Source: Coding guidelines
tests/unittest/disaggregated/test_agent.py (1)
56-69: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd invalid topology-input tests.
Test coverage summary: insufficient. The two added tests cover only valid
rank=2, world_size=4. Add cases for a partial pair,world_size <= 0, negativerank, andrank >= world_size. Both tests are covered bytests/integration/test_lists/test-db/l0_cpu.ymlthroughunittest/disaggregated. After adding the cases, runpytest tests/unittest/.🤖 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_agent.py` around lines 56 - 69, Expand test_python_agent_accepts_topology_without_forwarding_it_to_nixl with invalid topology cases covering partial rank/world_size input, non-positive world_size, negative rank, and rank greater than or equal to world_size; assert each raises the expected validation error. Keep the existing valid topology assertions and run pytest tests/unittest/ to verify the coverage.Sources: Coding guidelines, 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 `@cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp`:
- Around line 145-148: Update getIncrmentPort to calculate the incremented port
in a wider unsigned type, then validate that the result does not exceed the
maximum valid port range before converting to uint16_t. Preserve the existing
rank, worldSize, and times progression while ensuring overflow is rejected
rather than producing an incorrect port.
---
Nitpick comments:
In `@tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py`:
- Around line 104-109: Remove the redundant assert world_size is not None from
the rank/world_size validation block, since the preceding condition guarantees
both are provided together. Preserve the existing ValueError checks for invalid
rank and world_size values.
In `@tests/unittest/disaggregated/test_agent.py`:
- Around line 56-69: Expand
test_python_agent_accepts_topology_without_forwarding_it_to_nixl with invalid
topology cases covering partial rank/world_size input, non-positive world_size,
negative rank, and rank greater than or equal to world_size; assert each raises
the expected validation error. Keep the existing valid topology assertions and
run pytest tests/unittest/ to verify the coverage.
🪄 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: 645d8890-d3a0-44d6-b549-e7af673a5715
📒 Files selected for processing (23)
cpp/include/tensorrt_llm/executor/transferAgent.hcpp/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/agentBindings.cppcpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cppcpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.hcpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cppexamples/ray_orchestrator/disaggregated/README.mdexamples/ray_orchestrator/disaggregated/disagg_serving_local.shtensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.pytensorrt_llm/_torch/disaggregation/nixl/_agent_py.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/bindings/test_transfer_agent_bindings.pytests/unittest/disaggregated/test_agent.pytests/unittest/executor/test_rpc_worker_mixin.py
💤 Files with no reviewable changes (2)
- tests/integration/test_lists/waives.txt
- cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp
🚧 Files skipped from review as they are similar to previous changes (11)
- tests/integration/test_lists/test-db/l0_cpu.yml
- tests/integration/test_lists/test-db/l0_dgx_b200.yml
- tests/integration/test_lists/test-db/l0_dgx_h100.yml
- cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h
- tensorrt_llm/executor/rpc_worker_mixin.py
- cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h
- examples/ray_orchestrator/disaggregated/disagg_serving_local.sh
- tensorrt_llm/executor/base_worker.py
- examples/ray_orchestrator/disaggregated/README.md
- tests/unittest/executor/test_rpc_worker_mixin.py
- tests/integration/defs/examples/test_ray.py
7d06d2b to
4a2499c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #64325 [ run ] triggered by Bot. Commit: |
|
PR_Github #64325 [ run ] completed with state
|
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
The PG-collective path mirrors the existing ucxCacheCommunicator pattern and the shared process_responses() is a clean fix for the Ray RPC stream crash — the validation writeup is appreciated. A few things beyond the inline comments:
- Ticket: this is a nontrivial feature (new rank-discovery path, config surface, binding fields, new CI coverage) shipping under
[None]. Please attach a JIRA/TRTLLM ticket. - Undescribed test change: the removal of the
TRTLLM_NIXL_PORTsetenv fromcacheTransceiverTest.cppisn't mentioned in the PR description, and it's not incidental — see the inline comment on the port-derivation rank change it's compensating for. - IP selection behavior change: switching to
common::getLocalIp()changes the default (noTRTLLM_NIXL_INTERFACE) selection from "first non-docker/lo IPv4 interface" to the default-route probe / hostname fallback, and failure now throws instead of proceeding with"UNKNOWN IP". Both are improvements and align NIXL with the Mooncake path, but it can pick a different NIC on multi-homed hosts — worth one line in the description so it's findable if someone's deployment changes behavior. - Waiver removal: tp1 is root-caused and locally validated; tp2 (nvbugs/5612502) rests on the CI rerun you already flagged. Please make sure the B200/H100 disagg stages are green on the final diff before merge.
e6b8da3 to
36d82af
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #64499 [ run ] triggered by Bot. Commit: |
|
The local address selection still applies to every NIXL deployment, not only Ray, and |
| } | ||
| auto envPort = common::getEnvNixlPort(); | ||
| uint16_t port = envPort > 0 ? getIncrmentPort(envPort) : getAvailablePort(); | ||
| uint16_t port = getAvailablePort(); |
There was a problem hiding this comment.
This drops TRTLLM_NIXL_PORT support: the PR deletes getEnvNixlPort(), its declaration, and getIncrmentPort(), so the agent now always binds a random free port instead of honoring a pinned base port. I can see why -- the fixed-port scheme does not survive Ray -- and the knob is undocumented (no hits outside cacheTransceiverTest.cpp, which this PR updates). But it is still a user-reachable behavior change for anyone pinning the NIXL port in a firewalled or port-restricted deployment, and nothing in the PR description mentions it. Worth calling out in the description as an intentional removal.
Same hunk, separate note: mAddress moves from getAvailableIP() to common::getLocalIp(common::getEnvNixlInterface(), mRank). Converging on the helper the rest of this file already uses is the right direction, but it can change which local IP the agent advertises on a multi-NIC host when TRTLLM_NIXL_INTERFACE is unset.
|
PR_Github #64499 [ 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>
36d82af to
00ddbe9
Compare
7 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.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.