Skip to content

[None][feat] support NIXL cache transceiver with Ray - #17295

Open
chuangz0 wants to merge 7 commits into
NVIDIA:mainfrom
chuangz0:ray_support_nixl_cache_transceiver
Open

[None][feat] support NIXL cache transceiver with Ray#17295
chuangz0 wants to merge 7 commits into
NVIDIA:mainfrom
chuangz0:ray_support_nixl_cache_transceiver

Conversation

@chuangz0

@chuangz0 chuangz0 commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • support NIXL cache-transceiver rank discovery and metadata exchange with the Torch process group used by Ray
  • support both CPP and PYTHON cache-transceiver runtimes in the Ray disaggregated-serving example
  • re-enable the existing tp1/tp2 Ray disaggregated-serving coverage and add matching Python-runtime coverage
  • preserve access to internal libraries when packaging the NIXL wrapper so transfer-agent stub generation works without environment-specific library paths
  • share one AwaitResponseHelper.process_responses() implementation between regular and Ray RPC workers

Why

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 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 raw None to IPC, where response.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 before responses_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

  • signed-off commit hooks passed, including formatting, lint, YAML/test-list validation, DCO, pre-push pre-commit update, and confidentiality scan
  • clean full Release SM86 build from a new build directory completed, including the main bindings, libtensorrt_llm.so, NIXL wrapper, transfer-agent binding, executor, and UCX wrapper
  • nanobind stub generation completed as part of the clean build
  • a 581 MB wheel was built and installed into an isolated temporary directory with --no-deps
  • the installed wheel's NIXL wrapper has RUNPATH=$ORIGIN:$ORIGIN/nixl/, and ldd resolves libpg_utils.so from the same wheel
  • bare transfer-agent binding import and package import from the isolated wheel passed
  • isolated-wheel NIXL transfer-agent binding tests: 5 passed
  • isolated-wheel single-process Python cache-transceiver test: 1 passed
  • isolated-wheel NIXL initialization through a singleton Torch ProcessGroup passed with MPI disabled
  • Ray RPC preprocessing and event-loop error-broadcast tests: 7 passed
  • local final-code validation: CPP-runtime and PYTHON-runtime tp1 disaggregated-serving E2E each completed five concurrent requests; every request returned 32 completion tokens, and the existing disaggregated-test accuracy checks (The capital of Germany is Berlin, Asyncio is a Python library, and no Berlin Berlin) passed, with no NoneType.has_error, RPCStreamingError, or ray_executor_main_loop failure; context/generation shared one physical RTX A6000 through a temporary fractional placement adjustment that is not part of this PR

The B200/B300 tp1/tp2 CPP/PYTHON E2E cases should be rerun in CI with the final narrowed diff.

Dev Engineer Review

  • Adds Ray/Torch process-group rank and world-size discovery for NIXL cache agents.
  • Adds optional rank and world_size fields to BaseAgentConfig and both Python and C++ NIXL agent wrappers.
  • Adds MPI and single-process fallbacks.
  • Adds IPv6 address parsing and formatting.
  • Updates Ray examples to support C++ and Python transceiver runtimes.
  • Shares response preprocessing between regular workers and Ray RPC workers.
  • Review points remain for address-format versioning, mixed-version upgrades, IPv6 edge cases, discovery failures, and isolated libtorch packaging.
  • The new local-address selection should remain limited to the Ray/Torch process-group path unless the MPI behavior change is intentional.
  • The test-list entries and waiver removal match the added Ray coverage.

QA Engineer Review

  • Added test_ray_disaggregated_serving_python.
  • Refactored test_ray_disaggregated_serving to use runtime parameters.
  • Added concurrent HTTP response validation for Ray disaggregated serving.
  • Added RpcWorkerMixin.fetch_responses unit coverage.
  • Expanded BaseAgentConfig binding tests.
  • Added Python-agent topology tests.
  • Added C++ binding topology-forwarding tests.
  • Added the RPC worker unit test to l0_cpu.yml.
  • Added Python Ray coverage to l0_dgx_b200.yml and l0_dgx_h100.yml.
  • Removed the skipped tp2 Ray test from waives.txt.
  • Coverage is represented in CI test lists. Verdict: sufficient.

@chuangz0
chuangz0 force-pushed the ray_support_nixl_cache_transceiver branch 3 times, most recently from 0b89b30 to dd29ae5 Compare August 5, 2026 09:49
@chuangz0
chuangz0 marked this pull request as ready for review August 5, 2026 09:51
@chuangz0
chuangz0 requested review from a team as code owners August 5, 2026 09:51
@chuangz0

chuangz0 commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@coderabbitai

coderabbitai Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

Walkthrough

The 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.

Changes

Cache transport runtime support

Layer / File(s) Summary
Process-group rank and agent-state exchange
cpp/include/tensorrt_llm/executor/transferAgent.h, cpp/tensorrt_llm/executor/cache_transmission/agent_utils/*
Agent connections store rank and world size, discover topology through MPI or Torch process groups, exchange agent state through the selected collective, and use stored rank values for logging.
NIXL agent initialization and addressing
cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/*, tensorrt_llm/_torch/disaggregation/native/transfer.py, tensorrt_llm/_torch/disaggregation/nixl/*, tests/unittest/bindings/*, tests/unittest/disaggregated/test_agent.py
NIXL configuration and agents accept topology values. Initialization validates them, derives defaults, allocates rank-aware ports, discovers local addresses, and parses IPv4 or bracketed IPv6 endpoints.
Ray runtime selection and integration coverage
examples/ray_orchestrator/disaggregated/*, tests/integration/defs/examples/test_ray.py, tests/integration/test_lists/*
The serving script selects CPP or PYTHON NIXL runtimes. Documentation and integration tests cover both runtimes with concurrent completion validation.

RPC response processing

Layer / File(s) Summary
Shared response processing
tensorrt_llm/executor/base_worker.py, tensorrt_llm/executor/rpc_worker_mixin.py
AwaitResponseHelper.process_responses() centralizes callback filtering and nonblocking temporary-error draining. fetch_responses uses the helper before invoking the response handler.
RPC response validation
tests/unittest/executor/test_rpc_worker_mixin.py, tests/integration/test_lists/test-db/l0_cpu.yml
CPU-only tests verify timeout propagation, callback handling, temporary errors, handler output, and returned responses.

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
Loading

Possibly related PRs

Suggested reviewers: bowenfu, brnguyen2, shixiaowei02

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 24.24% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title follows the required format and clearly identifies Ray support for the NIXL cache transceiver.
Description check ✅ Passed The description explains the problem, solution, motivation, and extensive validation, although it does not reproduce every checklist section.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (2)
tests/integration/defs/examples/test_ray.py (1)

71-78: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add annotations and docstrings to the new test functions.

Add precise parameter types and -> None return types. Add Google-style docstrings for test_ray_disaggregated_serving_python and _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 win

Add 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.yml includes 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

📥 Commits

Reviewing files that changed from the base of the PR and between 9564b3b and dd29ae5.

📒 Files selected for processing (15)
  • cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp
  • cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h
  • cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt
  • cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp
  • cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h
  • examples/ray_orchestrator/disaggregated/disagg_serving_local.sh
  • scripts/build_wheel.py
  • tensorrt_llm/executor/base_worker.py
  • tensorrt_llm/executor/rpc_worker_mixin.py
  • tests/integration/defs/examples/test_ray.py
  • 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
  • tests/integration/test_lists/waives.txt
  • tests/unittest/executor/test_rpc_worker_mixin.py
💤 Files with no reviewable changes (1)
  • tests/integration/test_lists/waives.txt

Comment thread tensorrt_llm/executor/base_worker.py Outdated
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64023 [ run ] triggered by Bot. Commit: dd29ae5 Link to invocation

@chuangz0
chuangz0 force-pushed the ray_support_nixl_cache_transceiver branch from dd29ae5 to 760f0b9 Compare August 5, 2026 10:11
@coderabbitai

coderabbitai Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

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.

@chuangz0
chuangz0 force-pushed the ray_support_nixl_cache_transceiver branch from 760f0b9 to 7a7ac50 Compare August 5, 2026 10:19

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 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`:
- Line 448: Update the address construction and corresponding parsing/connection
logic around mAddress to remain compatible with both legacy ip#port and new
ip:port/[ipv6]:port formats during rolling upgrades. Distinguish IPv4, bracketed
IPv6, and legacy unbracketed IPv6 addresses without misparsing either format,
and ensure old and new instances can connect without requiring synchronized
deployment.
🪄 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: ebcd318b-3630-433f-9216-dbc290cbd380

📥 Commits

Reviewing files that changed from the base of the PR and between 63eb095 and 8b133ca.

📒 Files selected for processing (16)
  • cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp
  • cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h
  • cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt
  • cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp
  • cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h
  • examples/ray_orchestrator/disaggregated/README.md
  • examples/ray_orchestrator/disaggregated/disagg_serving_local.sh
  • scripts/build_wheel.py
  • tensorrt_llm/executor/base_worker.py
  • tensorrt_llm/executor/rpc_worker_mixin.py
  • tests/integration/defs/examples/test_ray.py
  • 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
  • tests/integration/test_lists/waives.txt
  • tests/unittest/executor/test_rpc_worker_mixin.py
💤 Files with no reviewable changes (1)
  • tests/integration/test_lists/waives.txt
🚧 Files skipped from review as they are similar to previous changes (12)
  • scripts/build_wheel.py
  • tests/integration/test_lists/test-db/l0_dgx_h100.yml
  • cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h
  • tests/unittest/executor/test_rpc_worker_mixin.py
  • examples/ray_orchestrator/disaggregated/disagg_serving_local.sh
  • tensorrt_llm/executor/rpc_worker_mixin.py
  • cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt
  • tensorrt_llm/executor/base_worker.py
  • tests/integration/test_lists/test-db/l0_cpu.yml
  • cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h
  • tests/integration/test_lists/test-db/l0_dgx_b200.yml
  • tests/integration/defs/examples/test_ray.py

@chuangz0

chuangz0 commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64186 [ run ] triggered by Bot. Commit: 8b133ca Link to invocation

@BowenFu

BowenFu commented Aug 6, 2026

Copy link
Copy Markdown

Sharpening @Shixiaowei02's point with the concrete mechanism, because the local-address change is not Ray-scoped:

getAvailableIP() (removed) walked getifaddrs, skipped lo/docker*, and took the first AF_INET address. common::getLocalIp("") with no TRTLLM_NIXL_INTERFACE set skips the NIC path entirely and falls into getLocalIpByRemoteOrHostName, which UDP-connects to 8.8.8.8 and reads back getsockname() — i.e. it returns the default-route interface, not the first enumerated one. On a multi-NIC node those are frequently different addresses, so an existing MPI NIXL disagg deployment that never set TRTLLM_NIXL_INTERFACE can silently start advertising a different (e.g. management) NIC for KV transfer. Two follow-on consequences:

  • getLocalIpByRemoteOrHostName falls back to IPv6, so mAddress can become [v6]:port. A peer on the previous build parses with find(":"), so a mixed-version pair breaks — the rolling-upgrade case already raised.
  • When neither the route probe nor hostname resolution yields an address, getLocalIp throws, where the old code returned "UNKNOWN IP" and logged. Better in isolation, but it turns a degraded start into a hard failure at agent construction.

Making the new selection opt-in (or preserving the old first-non-loopback-NIC scan as the default and using getLocalIp only when the Torch PG path is active) would keep this contained to Ray.

Separately: nixl_utils/CMakeLists.txt now links ${TORCH_LIBRARIES} and pg_utils into libtensorrt_llm_nixl_wrapper.so, and build_wheel.py widens its rpath to $ORIGIN:$ORIGIN/nixl/. That makes the deliberately-isolated NIXL wrapper depend on libtorch — worth confirming it still loads in the packaging configurations where the wrapper is dlopened.

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64186 [ run ] completed with state FAILURE. Commit: 8b133ca
/LLM/main/L0_MergeRequest_PR pipeline #52097 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@chuangz0
chuangz0 requested a review from a team as a code owner August 6, 2026 11:01
@chuangz0
chuangz0 requested review from Tabrizian and bo-nv August 6, 2026 11:01
@chuangz0
chuangz0 force-pushed the ray_support_nixl_cache_transceiver branch from a1e2713 to 7d06d2b Compare August 6, 2026 11:01
@chuangz0

chuangz0 commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@coderabbitai

coderabbitai Bot commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

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.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (2)
tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py (1)

104-109: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Remove the assertion from input validation.

Line 104 already proves that world_size is not None when rank is set. Delete the redundant assertion. Keep explicit ValueError checks 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 win

Add 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, negative rank, and rank >= world_size. Both tests are covered by tests/integration/test_lists/test-db/l0_cpu.yml through unittest/disaggregated. After adding the cases, run pytest 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

📥 Commits

Reviewing files that changed from the base of the PR and between 9a2ff2e and 7d06d2b.

📒 Files selected for processing (23)
  • cpp/include/tensorrt_llm/executor/transferAgent.h
  • cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp
  • cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h
  • cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt
  • cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp
  • cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp
  • cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h
  • cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp
  • examples/ray_orchestrator/disaggregated/README.md
  • examples/ray_orchestrator/disaggregated/disagg_serving_local.sh
  • tensorrt_llm/_torch/disaggregation/native/transfer.py
  • tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py
  • tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py
  • tensorrt_llm/executor/base_worker.py
  • tensorrt_llm/executor/rpc_worker_mixin.py
  • tests/integration/defs/examples/test_ray.py
  • 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
  • tests/integration/test_lists/waives.txt
  • tests/unittest/bindings/test_transfer_agent_bindings.py
  • tests/unittest/disaggregated/test_agent.py
  • tests/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

Comment thread cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp Outdated
@chuangz0
chuangz0 force-pushed the ray_support_nixl_cache_transceiver branch from 7d06d2b to 4a2499c Compare August 6, 2026 11:42
@chuangz0

chuangz0 commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64325 [ run ] triggered by Bot. Commit: 4a2499c Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64325 [ run ] completed with state SUCCESS. Commit: 4a2499c
/LLM/main/L0_MergeRequest_PR pipeline #52220 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@brnguyen2 brnguyen2 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_PORT setenv from cacheTransceiverTest.cpp isn'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 (no TRTLLM_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.

}
auto envPort = common::getEnvNixlPort();
uint16_t port = envPort > 0 ? getIncrmentPort(envPort) : getAvailablePort();
uint16_t port = envPort > 0 ? getIncrmentPort(envPort, mRank, mWorldSize) : getAvailablePort();

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Behavior change worth flagging: the fixed-port derivation previously used MpiComm::world() rank/size, which is unique per process within one launch. It now uses the config/session rank, which is not unique across sessions split from one world — context rank 0 and generation rank 0 on the same host both compute basePort + 0, and their static times counters advance in lockstep, so the collision persists across retries and the second agent fails to bind. That's exactly the topology of cacheTransceiverTest, which is presumably why the TRTLLM_NIXL_PORT setenv had to be removed there. If the fixed-port path is meant to stay supported (getEnvNixlPort), the offset needs a host-unique namespace (e.g. world rank when available, or an instance-derived offset); if it's test-only as the comment below says, consider saying so where TRTLLM_NIXL_PORT is documented.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed the fixed-port override entirely. NIXL now always selects an available port, so session-rank port collisions no longer apply.

@@ -994,10 +994,6 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam<AsymmetricTestPara
}
else if (isNixl)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removing the TRTLLM_NIXL_PORT setenv makes this test silently stop covering the fixed-port (getIncrmentPort) path — every run now takes getAvailablePort(). Since this removal is what keeps the test passing after the port derivation switched from world rank to session rank (see comment in transferAgent.cpp), please either note the rationale here in a comment or in the PR description; otherwise a future reader will assume the env-var path is still exercised.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The fixed-port path has been removed, so this test now intentionally covers the only supported dynamic-port behavior.

: mName{config.mName}
{
char const* disableMpi = std::getenv("TLLM_DISABLE_MPI");
bool const mpiEnabled = disableMpi == nullptr || std::atoi(disableMpi) == 0;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This hand-rolls the TLLM_DISABLE_MPI parse with different semantics from the canonical useMPI() in pgUtils.h (which disables MPI only on the exact string "1"): here atoi makes any nonzero-leading value ("2", "01", " 1") disable MPI while useMPI() would not. I assume pgUtils.h can't be included because this wrapper doesn't link torch — if so, a small torch-free helper in common/envUtils shared by both would prevent the two parses from drifting further.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed by using common::getBoolEnv("TLLM_DISABLE_MPI"), matching the canonical exact-"1" semantics without adding a torch dependency.

}
else
{
TLLM_LOG_WARNING("Torch process group is not initialized; cache transceiver defaults to one process");

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When MPI is disabled and the PG isn't initialized, this silently builds a world-size-1 CommState; with an actual multi-rank instance the metadata exchange is skipped and cache transfer fails later in a hard-to-attribute way (peer lookup / connect). Since a warning is easy to lose in rank logs, consider including the remediation in the message (initialize the process group before constructing the cache transceiver, or unset TLLM_DISABLE_MPI) so the failure is diagnosable from the log line alone.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated the warning with remediation: initialize the Torch process group before constructing the cache transceiver for multi-rank execution, or unset TLLM_DISABLE_MPI to use MPI.

@chuangz0
chuangz0 force-pushed the ray_support_nixl_cache_transceiver branch from 4a2499c to 23984ff Compare August 7, 2026 00:16
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>
@chuangz0
chuangz0 force-pushed the ray_support_nixl_cache_transceiver branch from 23984ff to e6b8da3 Compare August 7, 2026 04:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants