[None][feat] support NIXL cache transceiver with Ray - #17295
Conversation
0b89b30 to
dd29ae5
Compare
|
/bot run --disable-fail-fast |
WalkthroughThe change adds Torch process-group support to cache agents, switches Ray disaggregated serving to selectable NIXL runtimes, updates NIXL packaging paths, centralizes RPC response processing, and adds 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 RayServingScript
AgentConnectionManager->>PgHelper: Discover rank and world size
AgentConnectionManager->>AgentConnectionManager: Exchange serialized agent state
NixlTransferAgent->>NixlTransferAgent: Build rank-aware IP#port address
RayServingScript->>NixlTransferAgent: Configure selected CPP or PYTHON runtime
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
| while not self.temp_error_responses.empty(): | ||
| responses.append(self.temp_error_responses.get()) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
rg -n -C 4 \
'AwaitResponseHelper\(|process_responses\(|fetch_responses\(|temp_error_responses' \
tensorrt_llm testsRepository: NVIDIA/TensorRT-LLM
Length of output: 11158
🏁 Script executed:
# Check the class hierarchy and whether RpcWorkerMixin instances use both code paths
ast-grep outline tensorrt_llm/executor/rpc_worker_mixin.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 514
🏁 Script executed:
# Check the class hierarchy for workers that use RpcWorkerMixin
rg -n 'class.*RpcWorkerMixin' tensorrt_llmRepository: NVIDIA/TensorRT-LLM
Length of output: 495
🏁 Script executed:
# Check if await_responses() and fetch_responses() are both called on the same instance
rg -n 'await_responses|fetch_responses' tensorrt_llm/executor/*.py --type py -B2 -A2 | head -80Repository: NVIDIA/TensorRT-LLM
Length of output: 6557
🏁 Script executed:
# Check the event loop and threading model - look for async/threading context
rg -n 'Thread|asyncio|run_event_loop|threading' tensorrt_llm/executor/base_worker.py -A2 -B2 | head -100Repository: NVIDIA/TensorRT-LLM
Length of output: 772
🏁 Script executed:
# Understand the ManagedThread context and how AwaitResponseHelper is used
rg -n 'ManagedThread|setup_mainloop' tensorrt_llm/executor/rpc_worker_mixin.py -B3 -A8Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Check how AwaitResponseHelper.__call__ is invoked through the event loop
rg -n '_await_response_helper' tensorrt_llm/executor/base_worker.py -B2 -A2 | head -60Repository: NVIDIA/TensorRT-LLM
Length of output: 549
🏁 Script executed:
# For RPC workers, check if fetch_responses runs on a separate thread from the event loop
rg -n 'fetch_responses_loop_async|setup_engine_remote' tensorrt_llm/executor/rpc_worker_mixin.py -B4 -A10Repository: NVIDIA/TensorRT-LLM
Length of output: 1647
🏁 Script executed:
# Check which code path is active for RpcWorker instances
# Look at RpcWorker initialization and which methods are called
ast-grep outline tensorrt_llm/executor/rpc_worker.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 387
🏁 Script executed:
# Understand the setup_mainloop and what tasks run in it
rg -n 'setup_mainloop\|_fetch_responses_loop_async' tensorrt_llm/executor/base_worker.py -B5 -A15 | head -80Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Check if RpcWorkerMixin instances use both __call__ (ManagedThread) and fetch_responses concurrently
rg -n 'GenerationExecutorWorker\|RayGPUWorker\|RpcWorker' tensorrt_llm/executor/ -A20 | grep -E 'class |setup_mainloop|__call__|fetch_responses' | head -40Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Look at how RpcWorker is instantiated and what setup it does
cat -n tensorrt_llm/executor/rpc_worker.py | head -200Repository: NVIDIA/TensorRT-LLM
Length of output: 8322
🏁 Script executed:
# Search for AwaitResponseHelper.__call__ usage and which worker types use it
rg -n 'AwaitResponseHelper' tensorrt_llm/executor/ --type py -B3 -A3Repository: NVIDIA/TensorRT-LLM
Length of output: 2624
🏁 Script executed:
# Check BaseWorker to see if it has ManagedThread setup
rg -n 'ManagedThread\|_mainloop\|AwaitResponseHelper.*__call__' tensorrt_llm/executor/base_worker.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Check GenerationExecutorWorker to see if it uses both __call__ and fetch_responses
cat -n tensorrt_llm/executor/worker.py | head -150Repository: NVIDIA/TensorRT-LLM
Length of output: 7434
🏁 Script executed:
# Look for where AwaitResponseHelper.__call__ is actually invoked
rg -n '__call__' tensorrt_llm/executor/base_worker.py -B5 -A5Repository: NVIDIA/TensorRT-LLM
Length of output: 651
🏁 Script executed:
# Check if GenerationExecutorWorker has a main loop that calls AwaitResponseHelper
rg -n 'class GenerationExecutorWorker\|def main\|ManagedThread\|setup_mainloop\|_await_response' tensorrt_llm/executor/worker.py -B2 -A4Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Check if RPC server is started in GenerationExecutorWorker
cat -n tensorrt_llm/executor/worker.py | sed -n '67,90p'Repository: NVIDIA/TensorRT-LLM
Length of output: 1260
🏁 Script executed:
# Verify the concurrent access issue - check if fetch_responses can be called via RPC on GenerationExecutorWorker
rg -n 'start_rpc_server\|fetch_responses' tensorrt_llm/executor/worker.py -B2 -A2Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Check RpcWorkerMixin.start_rpc_server to confirm it exposes fetch_responses
rg -n 'start_rpc_server\|def fetch_responses' tensorrt_llm/executor/rpc_worker_mixin.py -B3 -A8Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Search for RPCServer implementation and what methods it exposes
rg -n 'class RPCServer\|def start_rpc_server' tensorrt_llm/executor/ -A15 | head -80Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Check what methods RpcWorkerMixin exposes that can be called via RPC
ast-grep outline tensorrt_llm/executor/rpc_worker_mixin.py --view expandedRepository: NVIDIA/TensorRT-LLM
Length of output: 1090
🏁 Script executed:
# Verify: is fetch_responses actually callable via RPC on a GenerationExecutorWorker with rpc_addr set?
rg -n 'fetch_responses' tensorrt_llm/executor/rpc_worker_mixin.py -B2 -A8Repository: NVIDIA/TensorRT-LLM
Length of output: 2716
🏁 Script executed:
# Final confirmation: Does GenerationExecutorWorker actually have both code paths active simultaneously?
# Check if start_rpc_server() is ever called with rpc_addr set in practice
rg -n 'GenerationExecutorWorker.*rpc_addr\|rpc_addr=' tensorrt_llm/ -B3 -A3Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Verify the test case to confirm the race condition scenario
cat -n tests/unittest/executor/test_rpc_worker_mixin.py | head -100Repository: NVIDIA/TensorRT-LLM
Length of output: 2953
🏁 Script executed:
# Verify the actual synchronization mechanism - check if there's any locking around temp_error_responses
rg -n 'Lock\|RLock\|Condition\|Semaphore\|temp_error_responses' tensorrt_llm/executor/base_worker.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Determine if the two code paths can truly run concurrently by checking Queue.get_nowait usage elsewhere
rg -n 'get_nowait\|queue\.Empty' tensorrt_llm/executor/ -B2 -A2Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
# Check the inheritance chain to verify both code paths exist in a single worker instance
cat -n tensorrt_llm/executor/worker.py | sed -n '37,95p'Repository: NVIDIA/TensorRT-LLM
Length of output: 2804
🏁 Script executed:
# Verify: is Queue.empty() followed by blocking get() a recognized race condition pattern in this codebase?
rg -n 'empty()' tensorrt_llm/executor/ --type py -B1 -A1Repository: NVIDIA/TensorRT-LLM
Length of output: 7205
🏁 Script executed:
# Confirm the race condition concern is documented elsewhere in the code
cat -n tensorrt_llm/executor/executor.py | sed -n '310,320p'Repository: NVIDIA/TensorRT-LLM
Length of output: 748
🏁 Script executed:
# Verify that temp_error_responses.get() without nowait can block indefinitely
python3 -c "
from queue import Queue
import time
import threading
q = Queue()
def reader():
# Simulate the pattern: check empty() then get()
if not q.empty():
item = q.get() # This will block if empty between check and get
print(f'Got: {item}')
# Check if Queue.get() blocks when empty
t = threading.Thread(target=reader)
t.daemon = True
t.start()
t.join(timeout=0.5)
if t.is_alive():
print('RACE CONDITION: get() blocked even though empty() was checked')
else:
print('No blocking observed')
"Repository: NVIDIA/TensorRT-LLM
Length of output: 178
Use get_nowait() with exception handling to drain temp_error_responses safely.
process_responses() is called from both AwaitResponseHelper.__call__() (which runs in a ManagedThread) and RpcWorkerMixin.fetch_responses() (which can be called via RPC from concurrent clients). When GenerationExecutorWorker is configured with an RPC server (rpc_addr provided), both paths are active on the same _await_response_helper instance.
The current drain pattern at lines 1125–1126 uses Queue.empty() followed by blocking get(). If one thread checks empty() and finds the queue non-empty, but another thread consumes the item before the first thread calls get(), the first thread blocks indefinitely.
The codebase documents the correct pattern in executor.py:310–320: use get_nowait() and catch queue.Empty. Apply the same approach here.
Current unsafe pattern
while not self.temp_error_responses.empty():
responses.append(self.temp_error_responses.get())🤖 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/executor/base_worker.py` around lines 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.
|
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. |
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>
760f0b9 to
7a7ac50
Compare
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/integration/defs/examples/test_ray.py (1)
72-197: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd QA list coverage for both Ray serving tests.
Changed tests:
test_ray_disaggregated_servingandtest_ray_disaggregated_serving_python. No test functions were removed. CI lists covertp2on B200 andtp1on H100. No QA list contains either test. Add both tests totests/integration/test_lists/qa/llm_function_core.txtor the designated functional QA list. SetLLM_MODELS_ROOTbefore execution. Coverage verdict: insufficient.🤖 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 72 - 197, Add both test_ray_disaggregated_serving and test_ray_disaggregated_serving_python to the designated functional QA list, such as tests/integration/test_lists/qa/llm_function_core.txt. Ensure the QA entries configure LLM_MODELS_ROOT before execution and cover the existing tp2 B200 and tp1 H100 CI cases.Source: Path instructions
🧹 Nitpick comments (1)
cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp (1)
447-455: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDeclare immutable local variables as
const.
localAgentState,str, and bothrecvBufferSizevariables are not reassigned. Declare themconstto meet the C++ coding guideline.Proposed change
- AgentState localAgentState{mAgentName, m_Agent->getLocalConnectionInfo()}; + AgentState const localAgentState{mAgentName, m_Agent->getLocalConnectionInfo()}; ... - auto str = oStream.str(); + auto const str = oStream.str(); ... - SizeType32 recvBufferSize = std::accumulate(sizeofBuffer.begin(), sizeofBuffer.end(), 0); + SizeType32 const recvBufferSize = std::accumulate(sizeofBuffer.begin(), sizeofBuffer.end(), 0); ... - SizeType32 recvBufferSize = std::accumulate(sizeofBuffer.begin(), sizeofBuffer.end(), 0); + SizeType32 const recvBufferSize = std::accumulate(sizeofBuffer.begin(), sizeofBuffer.end(), 0);Also applies to: 464-464, 490-490
🤖 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 `@cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp` around lines 447 - 455, Update the connection setup code around the localAgentState serialization flow to declare localAgentState and str as const, and apply the same const qualification to both recvBufferSize variables at the referenced locations. Leave reassigned variables unchanged.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.
Outside diff comments:
In `@tests/integration/defs/examples/test_ray.py`:
- Around line 72-197: Add both test_ray_disaggregated_serving and
test_ray_disaggregated_serving_python to the designated functional QA list, such
as tests/integration/test_lists/qa/llm_function_core.txt. Ensure the QA entries
configure LLM_MODELS_ROOT before execution and cover the existing tp2 B200 and
tp1 H100 CI cases.
---
Nitpick comments:
In `@cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp`:
- Around line 447-455: Update the connection setup code around the
localAgentState serialization flow to declare localAgentState and str as const,
and apply the same const qualification to both recvBufferSize variables at the
referenced locations. Leave reassigned variables unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d2d0cb22-f5c1-470e-87d5-60ccf9909c1f
📒 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
🚧 Files skipped from review as they are similar to previous changes (12)
- tests/integration/test_lists/test-db/l0_dgx_b200.yml
- cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h
- tests/integration/test_lists/test-db/l0_cpu.yml
- tensorrt_llm/executor/rpc_worker_mixin.py
- cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h
- cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt
- tensorrt_llm/executor/base_worker.py
- tests/integration/test_lists/test-db/l0_dgx_h100.yml
- tests/unittest/executor/test_rpc_worker_mixin.py
- scripts/build_wheel.py
- examples/ray_orchestrator/disaggregated/disagg_serving_local.sh
- cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp
|
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.
🧹 Nitpick comments (2)
cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp (2)
447-509: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDeclare unchanged local values as
const.
localAgentState,str, and bothrecvBufferSizevalues do not change after initialization. Declare themconst.As per coding guidelines: “declare unmodified variables as
const.”Proposed change
- AgentState localAgentState{mAgentName, m_Agent->getLocalConnectionInfo()}; + AgentState const localAgentState{mAgentName, m_Agent->getLocalConnectionInfo()}; ... - auto str = oStream.str(); + auto const str = oStream.str(); ... - SizeType32 recvBufferSize = std::accumulate(sizeofBuffer.begin(), sizeofBuffer.end(), 0); + SizeType32 const recvBufferSize = std::accumulate(sizeofBuffer.begin(), sizeofBuffer.end(), 0); ... - SizeType32 recvBufferSize = std::accumulate(sizeofBuffer.begin(), sizeofBuffer.end(), 0); + SizeType32 const recvBufferSize = std::accumulate(sizeofBuffer.begin(), sizeofBuffer.end(), 0);🤖 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 `@cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp` around lines 447 - 509, Declare the unchanged local values as const in the agent-state gathering flow: localAgentState, the serialized str value, and both recvBufferSize declarations. Leave mutable buffers, vectors, and loop variables unchanged.Source: Coding guidelines
449-449: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueName the single-process world-size value.
Replace
1with a named constant. This value defines the collective communication boundary.As per coding guidelines: “Avoid unexplained literals other than
0,nullptr,true, andfalse; assign other literals to named constants.”Proposed change
+ int constexpr kSingleProcessWorldSize = 1; AgentState const localAgentState{mAgentName, m_Agent->getLocalConnectionInfo()}; std::vector<AgentState> agentStates(mWorldSize); - if (mWorldSize > 1) + if (mWorldSize > kSingleProcessWorldSize)🤖 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 `@cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp` at line 449, Update the mWorldSize boundary check in the connection logic to replace the hardcoded 1 with a named constant that represents the single-process world size, using the existing connection-related symbols in connection.cpp. Keep the condition behavior unchanged for multi-process cases, but express the collective communication cutoff through a descriptive constant instead of an unexplained literal.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.
Nitpick comments:
In `@cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp`:
- Around line 447-509: Declare the unchanged local values as const in the
agent-state gathering flow: localAgentState, the serialized str value, and both
recvBufferSize declarations. Leave mutable buffers, vectors, and loop variables
unchanged.
- Line 449: Update the mWorldSize boundary check in the connection logic to
replace the hardcoded 1 with a named constant that represents the single-process
world size, using the existing connection-related symbols in connection.cpp.
Keep the condition behavior unchanged for multi-process cases, but express the
collective communication cutoff through a descriptive constant instead of an
unexplained literal.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 03477eb7-c8b3-41e5-8df6-9a031caa31e3
📒 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
🚧 Files skipped from review as they are similar to previous changes (13)
- tests/integration/test_lists/test-db/l0_dgx_b200.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
- tests/integration/test_lists/test-db/l0_cpu.yml
- cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt
- tensorrt_llm/executor/rpc_worker_mixin.py
- scripts/build_wheel.py
- cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp
- cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h
- tests/integration/test_lists/test-db/l0_dgx_h100.yml
- tests/integration/defs/examples/test_ray.py
- tensorrt_llm/executor/base_worker.py
|
The address information exchanged between context and generation servers changes format here without a version marker, so during a rolling upgrade a mixed-version pair would fail to connect, and likewise making the new local-address selection opt-in, since it applies beyond the Ray path. |
|
PR_Github #64023 [ run ] completed with state
|
3 participants
Summary
AwaitResponseHelper.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 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
CPPandPYTHONtransceiver runtimes.QA Engineer Review
test_ray_disaggregated_serving_python(ray_example_root, llm_venv, tp_size).unittest/executor/test_rpc_worker_mixin.pywith response and callback coverage.tests/integration/test_lists/test-db/l0_cpu.yml.tests/integration/test_lists/test-db/l0_dgx_b200.ymltests/integration/test_lists/test-db/l0_dgx_h100.ymltests/integration/test_lists/waives.txt.