[None][feat] align time metrics between cpp and python cache transceiver#15871
Conversation
3da38fc to
5171454
Compare
|
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:
WalkthroughIntroduces instance-scoped cache-transfer output, synchronized clock and transfer metrics, generation-side CSV summaries, Python transfer-session tracking, and tests for cross-rank timing aggregation. ChangesKV cache transfer timing and instance-scoped output
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant KvCacheTransceiverV2
participant GenerationRanks
participant PerfLogManager
KvCacheTransceiverV2->>GenerationRanks: allgather request timing and size
GenerationRanks-->>KvCacheTransceiverV2: merged metrics
KvCacheTransceiverV2->>PerfLogManager: log generation summary
Possibly related PRs
Suggested labels: 🚥 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: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp (1)
187-201: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winUse the PG rank fallback here too
getTransferOutputPath()unconditionally readsmpi::MpiComm::world().getRank(), but this codebase has a non-MPI path. That can make thesend/recvtiming CSV names collide when MPI is disabled; mirror theuseMPI() ? ... : pg_utils::get_world_pg()->getRank()branch used inwriteGenTransferSummary.🤖 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/batch_manager/dataTransceiver.cpp` around lines 187 - 201, getTransferOutputPath() always uses mpi::MpiComm::world().getRank(), which breaks the non-MPI path and can cause send/recv timing CSV name collisions. Update the rank lookup in this helper to mirror writeGenTransferSummary by using the PG rank fallback when useMPI() is false, and keep the rest of the prefix/path construction in getTransferOutputPath() unchanged.tensorrt_llm/_torch/disaggregation/transceiver.py (1)
58-108: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winEnv-var gate for the timing-sync collective is read per-process, not synchronized across ranks — deadlock risk.
Both
_sync_transfer_timing(Line 403) and its call site incheck_gen_transfer_status(Line 687) decide whether to invoke the blockingself._gen_allgather(...)collective by independently callingos.getenv("TRTLLM_KVCACHE_TIME_OUTPUT_PATH")on each process.completedis already made consistent across ranks via_gen_consensus_outcome, but the env var itself is never synchronized. If this variable is set on some ranks/processes and not others (e.g. inconsistent env propagation across nodes in a multi-node TP/PP job), some ranks will call the collective while others skip it — a classic mismatched-collective-call hang, since allgather requires every participant to call it in lockstep.The class already has an established pattern for this exact problem — broadcasting a rank-0 decision once at construction (see
_broadcast_instance_name,_broadcast_context_endpoint). Consider doing the same for this flag instead of re-readingos.getenvon every iteration from every rank independently.🔒 Proposed fix
def _init_sync_policy(self): m = self._mapping self._ctx_need_tp_sync = m.tp_size > 1 and not m.enable_attention_dp self._ctx_need_pp_sync = m.pp_size > 1 self._gen_need_sync = not (m.world_size == 1 or (m.enable_attention_dp and m.pp_size == 1)) pp_allgather: Callable = getattr(self._dist, "pp_allgather") self._gen_allgather: Callable = ( pp_allgather if m.enable_attention_dp else self._dist.allgather ) + self._kv_time_output_enabled = self._broadcast_kv_time_output_enabled() + + def _broadcast_kv_time_output_enabled(self) -> bool: + if self._dist.rank == 0: + enabled = bool(os.getenv("TRTLLM_KVCACHE_TIME_OUTPUT_PATH")) + self._dist.broadcast(enabled, 0) + return enabled + return cast(bool, self._dist.broadcast(None, 0))Then replace the per-call
os.getenv("TRTLLM_KVCACHE_TIME_OUTPUT_PATH")checks at Line 403 and Line 687 withself._kv_time_output_enabled.Also applies to: 124-132, 393-406, 686-699
🤖 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/transceiver.py` around lines 58 - 108, The timing-sync collective gate is being decided independently on each rank via os.getenv, which can mismatch across processes and deadlock the allgather. Add a one-time, rank-synchronized flag in KvCacheTransceiverV2 (similar to _broadcast_instance_name and _broadcast_context_endpoint), store it on the instance as _kv_time_output_enabled, and use that cached value in _sync_transfer_timing and check_gen_transfer_status instead of re-reading the environment each time.
🧹 Nitpick comments (4)
cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp (1)
193-200: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDuplicate accessors for the same values.
kv_cache_transfer_start/kv_cache_transfer_endare already exposed as read-only properties (lines 194-195); the newget_kv_cache_transfer_start/get_kv_cache_transfer_endmethods (196-197) bind to the identical getters. Consider standardizing on one style (likely the explicitget_*/set_*method pairing, matching the existing setters at lines 199-200) and dropping the redundant property to avoid two ways to read the same field.🤖 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/nanobind/batch_manager/bindings.cpp` around lines 193 - 200, The batch manager bindings expose the same KV cache transfer timestamps twice, once as read-only properties and again as explicit getters. In bindings.cpp, update the GenLlmReq nanobind definitions to use one consistent API style only: keep the get_kv_cache_transfer_start/get_kv_cache_transfer_end and matching setters, or keep the properties, but remove the redundant duplicate accessor set so the same value is not read through two parallel bindings.cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp (2)
350-364: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueHard-coded
kUuidLen = 36duplicates the format assumption fromgenerateInstanceId().Nothing ties this constant to the actual output width of
generateInstanceId()(lines 75-87). If that formatting ever changes, this broadcast path silently truncates/corrupts the instance ID with no compile-time or runtime guard.🤖 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/batch_manager/cacheTransceiver.cpp` around lines 350 - 364, The PG broadcast path in cacheTransceiver.cpp hard-codes the UUID length separately from generateInstanceId(), so it can silently truncate or corrupt instance IDs if the format changes. Replace the local kUuidLen assumption in the allgather path with a single shared source of truth tied to generateInstanceId()’s output size, and use that same value when allocating sendBuf/recvBuf and reconstructing mInstanceId. Add a guard or assertion in CacheTransceiver to ensure the received instance ID length matches the expected generated length.
755-812: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
batchUpdateKVCacheTransferBWalways performs the allgather/allgatherv, unlike other collectives in this file.
gatherRequestIdsandreduceTransferStatesboth guard collective calls withcomm->getSize() > 1before invoking them; this new function unconditionally callscomm->allgather/allgathervregardless ofworldSize. For single-rank deployments that still enableTRTLLM_KVCACHE_TIME_OUTPUT_PATH, this issues an unnecessary (and, depending on the PG backend, potentially unsupported) size-1 collective.🤖 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/batch_manager/cacheTransceiver.cpp` around lines 755 - 812, batchUpdateKVCacheTransferBW currently unconditionally runs comm->allgather and comm->allgatherv even when there is only one rank. Update this function to mirror gatherRequestIds and reduceTransferStates by guarding the collective path with worldSize > 1, and for the single-rank case skip the distributed merge and just use the local serialized data to populate/update the cache transfer bandwidth state. Locate the change in batchUpdateKVCacheTransferBW and keep the useMPI/non-MPI branches inside the multi-rank-only path.tensorrt_llm/_torch/disaggregation/native/perf_logger.py (1)
369-395: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCatch a narrower exception type than bare
Exception.The try block only does filesystem/logging setup (
os.makedirs,logging.FileHandler, stream writes), which realistically raisesOSError(includingPermissionError). CatchingExceptionbroadly can mask unrelated bugs.♻️ Proposed fix
- except Exception as e: + except OSError as e: sys.stderr.write( f"[KV Transfer] Warning: Failed to create gen summary log file {log_file}: {e}\n" ) returnBased on learnings and coding guidelines: "Avoid broad exception handling — catch specific exceptions, not bare
except:" and "limit the except to the smallest set of errors possible; use specific exception types."🤖 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/native/perf_logger.py` around lines 369 - 395, The exception handling in the gen summary log setup is too broad, so narrow the `except` in `perf_logger.py` around the `try` block that creates the `file_logger` and `file_handler`. Catch only the filesystem/logging errors this code can realistically raise, such as `OSError` (including `PermissionError`), and keep the existing warning/return path in the `self._file_loggers` setup unchanged.Sources: Coding guidelines, Linters/SAST tools
🤖 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/batch_manager/cacheTransceiver.cpp`:
- Around line 1246-1266: `CacheTransceiver::writeGenTransferSummary` opens
`mGenTransferSummaryFile` without verifying success, so add an `is_open()` check
immediately after `open(filePath)` and fail with an explicit
`TLLM_CHECK_WITH_INFO` message if the file cannot be opened. Use the same
pattern already used in `dataTransceiver.cpp` (`release()`/`receiveSync()`) and
keep the fix localized to the `writeGenTransferSummary` path so missing
permissions or bad output paths are surfaced instead of silently dropping
writes.
- Around line 366-401: The steady-clock calibration in CacheTransceiver has a
race and inconsistent synchronization: the
`LlmRequest::sGlobalSteadyClockOffset` check/set in this block is
unsynchronized, and the PG path in `CacheTransceiverComm` skips the
pre-timestamp barrier used by the MPI path. Protect the one-time initialization
with the same kind of mutex/guard used elsewhere in `CacheTransceiver`, and make
the collective calibration path synchronize all ranks immediately before reading
`steady_clock::now()` so both `useMPI()` and non-MPI branches measure from a
consistent point.
In `@tests/unittest/disaggregated/test_perf_logger.py`:
- Around line 109-123: The test
`test_log_gen_transfer_summary_disabled_without_env` is mutating global
environment state by popping `TRTLLM_KVCACHE_TIME_OUTPUT_PATH` without restoring
it. Update the test to scope the env change with `patch.dict` (as in the sibling
env-isolation test) so `PerfLogManager` can be exercised without leaking state
into other tests; keep the call to `log_gen_transfer_summary` unchanged.
In `@tests/unittest/disaggregated/test_sync_transfer_timing.py`:
- Around line 118-131: The test `test_skips_when_no_env` mutates `os.environ`
directly by popping `TRTLLM_KVCACHE_TIME_OUTPUT_PATH` without restoring it,
which can leak state into other tests. Update this test to use the same isolated
environment approach used elsewhere in the class, such as
`@patch.dict("os.environ", {}, clear=False)` with the pop inside the context or
`monkeypatch.delenv`, and keep the assertions around `_sync_transfer_timing` and
`allgather_mock` unchanged.
---
Outside diff comments:
In `@cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp`:
- Around line 187-201: getTransferOutputPath() always uses
mpi::MpiComm::world().getRank(), which breaks the non-MPI path and can cause
send/recv timing CSV name collisions. Update the rank lookup in this helper to
mirror writeGenTransferSummary by using the PG rank fallback when useMPI() is
false, and keep the rest of the prefix/path construction in
getTransferOutputPath() unchanged.
In `@tensorrt_llm/_torch/disaggregation/transceiver.py`:
- Around line 58-108: The timing-sync collective gate is being decided
independently on each rank via os.getenv, which can mismatch across processes
and deadlock the allgather. Add a one-time, rank-synchronized flag in
KvCacheTransceiverV2 (similar to _broadcast_instance_name and
_broadcast_context_endpoint), store it on the instance as
_kv_time_output_enabled, and use that cached value in _sync_transfer_timing and
check_gen_transfer_status instead of re-reading the environment each time.
---
Nitpick comments:
In `@cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp`:
- Around line 350-364: The PG broadcast path in cacheTransceiver.cpp hard-codes
the UUID length separately from generateInstanceId(), so it can silently
truncate or corrupt instance IDs if the format changes. Replace the local
kUuidLen assumption in the allgather path with a single shared source of truth
tied to generateInstanceId()’s output size, and use that same value when
allocating sendBuf/recvBuf and reconstructing mInstanceId. Add a guard or
assertion in CacheTransceiver to ensure the received instance ID length matches
the expected generated length.
- Around line 755-812: batchUpdateKVCacheTransferBW currently unconditionally
runs comm->allgather and comm->allgatherv even when there is only one rank.
Update this function to mirror gatherRequestIds and reduceTransferStates by
guarding the collective path with worldSize > 1, and for the single-rank case
skip the distributed merge and just use the local serialized data to
populate/update the cache transfer bandwidth state. Locate the change in
batchUpdateKVCacheTransferBW and keep the useMPI/non-MPI branches inside the
multi-rank-only path.
In `@cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp`:
- Around line 193-200: The batch manager bindings expose the same KV cache
transfer timestamps twice, once as read-only properties and again as explicit
getters. In bindings.cpp, update the GenLlmReq nanobind definitions to use one
consistent API style only: keep the
get_kv_cache_transfer_start/get_kv_cache_transfer_end and matching setters, or
keep the properties, but remove the redundant duplicate accessor set so the same
value is not read through two parallel bindings.
In `@tensorrt_llm/_torch/disaggregation/native/perf_logger.py`:
- Around line 369-395: The exception handling in the gen summary log setup is
too broad, so narrow the `except` in `perf_logger.py` around the `try` block
that creates the `file_logger` and `file_handler`. Catch only the
filesystem/logging errors this code can realistically raise, such as `OSError`
(including `PermissionError`), and keep the existing warning/return path in the
`self._file_loggers` setup unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 933ba7bb-bd93-47aa-ab78-5cf8e314c058
📒 Files selected for processing (12)
cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.hcpp/tensorrt_llm/batch_manager/cacheTransceiver.cppcpp/tensorrt_llm/batch_manager/dataTransceiver.cppcpp/tensorrt_llm/batch_manager/dataTransceiver.hcpp/tensorrt_llm/nanobind/batch_manager/bindings.cppcpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpptensorrt_llm/_torch/disaggregation/native/perf_logger.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/transceiver.pytests/integration/test_lists/test-db/l0_h100.ymltests/unittest/disaggregated/test_perf_logger.pytests/unittest/disaggregated/test_sync_transfer_timing.py
5171454 to
eb692a9
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57307 [ run ] triggered by Bot. Commit: |
|
PR_Github #57307 [ run ] completed with state
|
eb692a9 to
488b581
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57368 [ run ] triggered by Bot. Commit: |
|
PR_Github #57368 [ run ] completed with state
|
488b581 to
5979c4f
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57424 [ run ] triggered by Bot. Commit: |
|
PR_Github #57424 [ run ] completed with state
|
5979c4f to
cb412f3
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57651 [ run ] triggered by Bot. Commit: |
|
PR_Github #57651 [ run ] completed with state
|
fba6ee9 to
1c80fd3
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #60680 [ run ] completed with state
|
064ae71 to
441d88c
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)
cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp (1)
1319-1324: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winProtect the async-resource registry with one mutex.
cancelRequest()readsmInstanceToAsyncResourcewhilerequestAndReceiveAsyncMultiThreads()inserts into it and~Impl()iterates it, so concurrent access can race onstd::unordered_map. Guard all accesses with the same lock.🤖 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/batch_manager/dataTransceiver.cpp` around lines 1319 - 1324, Protect every access to mInstanceToAsyncResource with one shared mutex: update cancelRequest() around its lookup, requestAndReceiveAsyncMultiThreads() around insertion, and ~Impl() around iteration or clearing. Use the same lock for all three paths and preserve the existing cancellation, registration, and cleanup behavior.
🧹 Nitpick comments (1)
cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp (1)
794-804: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffMove agent-only control-path operations behind the connection interface.
dataTransceiver.cpprepeatedly probesAgentConnectionManager/AgentConnectionwithdynamic_cast, and line 1235 also needsconst_castto call a mutating method. Expose these agent-specific operations on the interface or resolve the backend once during construction so the transceiver can call them directly.
cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp#L794-L804cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp#L1232-L1236cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp#L1401-L1407🤖 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/batch_manager/dataTransceiver.cpp` around lines 794 - 804, Replace the repeated AgentConnectionManager/AgentConnection dynamic_cast checks in dataTransceiver.cpp with agent-specific operations exposed through the connection interface, or resolve the backend once during DataTransceiver construction. Update sites 794-804, 1232-1236, and 1401-1407 to call the interface directly, including a non-const mutating method so the const_cast at 1232-1236 is removed; preserve the existing agent and non-agent send behavior.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 `@cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp`:
- Around line 1319-1324: Protect every access to mInstanceToAsyncResource with
one shared mutex: update cancelRequest() around its lookup,
requestAndReceiveAsyncMultiThreads() around insertion, and ~Impl() around
iteration or clearing. Use the same lock for all three paths and preserve the
existing cancellation, registration, and cleanup behavior.
---
Nitpick comments:
In `@cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp`:
- Around line 794-804: Replace the repeated
AgentConnectionManager/AgentConnection dynamic_cast checks in
dataTransceiver.cpp with agent-specific operations exposed through the
connection interface, or resolve the backend once during DataTransceiver
construction. Update sites 794-804, 1232-1236, and 1401-1407 to call the
interface directly, including a non-const mutating method so the const_cast at
1232-1236 is removed; preserve the existing agent and non-agent send behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 59caa2ba-133b-4dbd-b0b6-14d8fdddfe1c
📒 Files selected for processing (5)
cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.hcpp/include/tensorrt_llm/batch_manager/llmRequest.hcpp/tensorrt_llm/batch_manager/cacheTransceiver.cppcpp/tensorrt_llm/batch_manager/dataTransceiver.cppcpp/tensorrt_llm/batch_manager/dataTransceiver.h
🚧 Files skipped from review as they are similar to previous changes (3)
- cpp/tensorrt_llm/batch_manager/dataTransceiver.h
- cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h
- cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
d4ea8a9 to
016e3de
Compare
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
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/batch_manager/llmRequest.cpp`:
- Around line 25-33: Protect the process-global steady-clock offset with
synchronized accessors: in cpp/tensorrt_llm/batch_manager/llmRequest.cpp:25-33,
add mutex-guarded getter/setter or equivalent write-once atomic publication; in
cpp/include/tensorrt_llm/batch_manager/llmRequest.h:87-94, replace the raw
mutable reference with those accessors; in
cpp/include/tensorrt_llm/batch_manager/llmRequest.h:2398-2401, read a
synchronized snapshot; and in
cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp:484-488, bind the
synchronized accessors instead of the field.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0ec334ea-fa71-4b50-a5bf-aae27a6906e8
📒 Files selected for processing (7)
cpp/include/tensorrt_llm/batch_manager/llmRequest.hcpp/tensorrt_llm/batch_manager/cacheTransceiver.cppcpp/tensorrt_llm/batch_manager/llmRequest.cppcpp/tensorrt_llm/nanobind/batch_manager/bindings.cppcpp/tensorrt_llm/nanobind/bindings.cpptests/unittest/disaggregated/test_perf_logger.pytests/unittest/disaggregated/test_sync_transfer_timing.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tests/unittest/disaggregated/test_perf_logger.py
- tests/unittest/disaggregated/test_sync_transfer_timing.py
- cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
|
PR_Github #61247 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
tests/unittest/disaggregated/test_perf_logger.py (1)
82-107: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover C++-path priority for normal perf logs.
This only exercises
log_gen_transfer_summary, which reads the environment directly; it does not cover the changed initialization and normal file naming paths. Add a test with both C++ and legacy variables set, calllog(), and assert{tmpdir}/test_inst_0.csvis used instead of the legacy filename.As per path instructions, test coverage feedback must be actionable.
🤖 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_perf_logger.py` around lines 82 - 107, Add a test for the normal perf-log path that sets both the C++ output-path variable and the legacy variable, constructs `PerfLogManager`, and calls `log()`. Assert that the C++ path produces `{tmpdir}/test_inst_0.csv` and that the legacy filename is not used, covering initialization priority and normal file naming rather than `log_gen_transfer_summary`.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 `@examples/disaggregated/slurm/cache_transceiver_test/README.md`:
- Line 67: Update the CSV artifact pattern documentation near the C++
transceiver timing entry to include the optional __c<ci> suffix produced by
_preserve_cpp_csvs, while retaining the existing unsuffixed form for
single-combination runs.
In `@tensorrt_llm/_torch/disaggregation/native/perf_logger.py`:
- Around line 391-395: In the gen summary log file creation error handler,
replace the broad Exception catch with filesystem-specific OSError handling,
including PermissionError through its existing subclass relationship. Keep the
warning output and return behavior unchanged, while allowing unexpected
programming errors to propagate.
- Around line 334-341: Annotate every newly introduced function: add -> None to
perf_logger.py’s log_gen_transfer_summary, the new test methods in
test_perf_logger.py, and all test methods in test_sync_transfer_timing.py; add
precise parameter and return types to the fake-request helpers and
_make_transceiver in test_sync_transfer_timing.py. Apply the annotations at all
four identified sites without changing behavior.
---
Nitpick comments:
In `@tests/unittest/disaggregated/test_perf_logger.py`:
- Around line 82-107: Add a test for the normal perf-log path that sets both the
C++ output-path variable and the legacy variable, constructs `PerfLogManager`,
and calls `log()`. Assert that the C++ path produces `{tmpdir}/test_inst_0.csv`
and that the legacy filename is not used, covering initialization priority and
normal file naming rather than `log_gen_transfer_summary`.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d7f0e955-0e0a-4840-9288-5c22f0776a87
📒 Files selected for processing (22)
cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.hcpp/include/tensorrt_llm/batch_manager/llmRequest.hcpp/tensorrt_llm/batch_manager/cacheTransceiver.cppcpp/tensorrt_llm/batch_manager/dataTransceiver.cppcpp/tensorrt_llm/batch_manager/dataTransceiver.hcpp/tensorrt_llm/batch_manager/llmRequest.cppcpp/tensorrt_llm/nanobind/batch_manager/bindings.cppcpp/tensorrt_llm/nanobind/bindings.cppcpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cppexamples/disaggregated/slurm/cache_transceiver_test/README.mdexamples/disaggregated/slurm/cache_transceiver_test/report.pyexamples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.pytensorrt_llm/_torch/disaggregation/native/perf_logger.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/transceiver.pytensorrt_llm/_torch/pyexecutor/py_executor.pytests/integration/defs/disaggregated/test_disaggregated.pytests/integration/test_lists/test-db/l0_h100.ymltests/unittest/disaggregated/test_bounce.pytests/unittest/disaggregated/test_cache_transceiver_harness_report.pytests/unittest/disaggregated/test_perf_logger.pytests/unittest/disaggregated/test_sync_transfer_timing.py
🚧 Files skipped from review as they are similar to previous changes (10)
- tests/integration/test_lists/test-db/l0_h100.yml
- cpp/tensorrt_llm/batch_manager/llmRequest.cpp
- cpp/tensorrt_llm/batch_manager/dataTransceiver.h
- cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp
- cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp
- cpp/tensorrt_llm/nanobind/bindings.cpp
- cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp
- cpp/include/tensorrt_llm/batch_manager/llmRequest.h
- tensorrt_llm/_torch/disaggregation/native/transfer.py
- tensorrt_llm/_torch/disaggregation/transceiver.py
chzblych
left a comment
chzblych
left a comment
There was a problem hiding this comment.
Approved for the testing part changes.
|
PR_Github #61247 [ 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>
…naming The cache transceiver now writes per-transfer timing CSVs as <instanceId>_<rank>_<tag>.csv (instanceId is a runtime UUID that disambiguates instances sharing an output directory). Update all consumers that previously assumed the fixed rank_<rank>_<tag>.csv name: - integration test test_disaggregated_kv_cache_time_output: match the send/recv files by their _<tag>.csv suffix instead of rank_0_*. - report.py: glob *_recv*.csv (still matches legacy rank_* names). - run_cache_transceiver_test.py _preserve_cpp_csvs: rename by the _<rank>_<tag>.csv suffix, keeping the per-rank no-race guarantee. - README: document the <instanceId>_* file names. - harness report unittest: use realistic UUID-named fixtures and keep a legacy rank_ case for backward compatibility. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
…alibration Address PR review comments on the cache-transceiver time-sync feature: - Collapse the global steady-clock offset into a single process-global copy reached through globalSteadyClockOffset() (storage in llmRequest.cpp) instead of an inline-static member that each .so duplicated; nanobind binds to that single storage so libtensorrt_llm.so and the extension module share one value. - Guard the one-time offset calibration in CacheTransceiver with a mutex to avoid a check-and-set race, and add a pre-timestamp pseudo-barrier for the process-group path (the MPI path already barriers). - Reuse a thread_local RNG in generateInstanceId to avoid re-seeding per call. - Fail loudly with TLLM_CHECK_WITH_INFO when the gen transfer summary file cannot be opened. - Scope TRTLLM_KVCACHE_TIME_OUTPUT_PATH with patch.dict in the two unit tests so they no longer leak environment state. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
… CSV suffix doc Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
016e3de to
6771229
Compare
|
/bot run --stage-list "DGX_B200-PyTorch-9, DGX_H100-PyTorch-3" |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
cpp/tensorrt_llm/batch_manager/dataTransceiver.h (1)
271-272: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new
instanceIdparameter.Add an
@param instanceIdentry to both constructor comments describing its role in instance-scoped transfer output naming.As per coding guidelines, public C++ interfaces should use Doxygen comments and document their arguments.
Also applies to: 335-336
🤖 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/batch_manager/dataTransceiver.h` around lines 271 - 272, Add Doxygen `@param` instanceId documentation to both CacheSender and corresponding constructor comments, describing that it controls instance-scoped transfer output naming. Preserve the existing documentation and parameter behavior.Source: Coding guidelines
tests/unittest/disaggregated/test_sync_transfer_timing.py (1)
17-18: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNarrow the override type instead of using
Any.
_make_transceiveronly accepts known override values, butAnydisables static checking for them and for subsequentoverrides.get(...)results. Define a typed override shape and callable alias, or use a narrower protocol, so invalid test doubles are caught during type checking.As per coding guidelines, avoid
Anyand unnecessary type ignores.Also applies to: 54-61
🤖 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_sync_transfer_timing.py` around lines 17 - 18, Replace Any in _make_transceiver and its override handling with a typed override shape and callable alias or narrower protocol that describes the accepted values and get results. Update the related overrides.get usage around lines 54–61 so static checking catches invalid test doubles, without adding type ignores.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/disaggregated/test_perf_logger.py`:
- Around line 82-108: Update the affected perf logger tests, including
test_log_gen_transfer_summary_writes_csv and the test at the referenced second
range, to remove and close every FileHandler before each TemporaryDirectory
context exits. Add a shared teardown/helper that accesses the configured logger,
removes its handlers, closes them, and resets the PerfLogManager singleton as
needed, then invoke it in both tests.
---
Nitpick comments:
In `@cpp/tensorrt_llm/batch_manager/dataTransceiver.h`:
- Around line 271-272: Add Doxygen `@param` instanceId documentation to both
CacheSender and corresponding constructor comments, describing that it controls
instance-scoped transfer output naming. Preserve the existing documentation and
parameter behavior.
In `@tests/unittest/disaggregated/test_sync_transfer_timing.py`:
- Around line 17-18: Replace Any in _make_transceiver and its override handling
with a typed override shape and callable alias or narrower protocol that
describes the accepted values and get results. Update the related overrides.get
usage around lines 54–61 so static checking catches invalid test doubles,
without adding type ignores.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d0a6f39c-cc90-42ad-87ff-a4ab94d73ff5
📒 Files selected for processing (22)
cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.hcpp/include/tensorrt_llm/batch_manager/llmRequest.hcpp/tensorrt_llm/batch_manager/cacheTransceiver.cppcpp/tensorrt_llm/batch_manager/dataTransceiver.cppcpp/tensorrt_llm/batch_manager/dataTransceiver.hcpp/tensorrt_llm/batch_manager/llmRequest.cppcpp/tensorrt_llm/nanobind/batch_manager/bindings.cppcpp/tensorrt_llm/nanobind/bindings.cppcpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cppexamples/disaggregated/slurm/cache_transceiver_test/README.mdexamples/disaggregated/slurm/cache_transceiver_test/report.pyexamples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.pytensorrt_llm/_torch/disaggregation/native/perf_logger.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/transceiver.pytensorrt_llm/_torch/pyexecutor/py_executor.pytests/integration/defs/disaggregated/test_disaggregated.pytests/integration/test_lists/test-db/l0_h100.ymltests/unittest/disaggregated/test_bounce.pytests/unittest/disaggregated/test_cache_transceiver_harness_report.pytests/unittest/disaggregated/test_perf_logger.pytests/unittest/disaggregated/test_sync_transfer_timing.py
🚧 Files skipped from review as they are similar to previous changes (16)
- tests/integration/test_lists/test-db/l0_h100.yml
- cpp/tensorrt_llm/nanobind/bindings.cpp
- cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp
- cpp/tensorrt_llm/batch_manager/llmRequest.cpp
- tensorrt_llm/_torch/pyexecutor/py_executor.py
- examples/disaggregated/slurm/cache_transceiver_test/README.md
- tests/unittest/disaggregated/test_bounce.py
- cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h
- cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp
- tensorrt_llm/_torch/disaggregation/native/perf_logger.py
- cpp/include/tensorrt_llm/batch_manager/llmRequest.h
- cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp
- examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py
- tensorrt_llm/_torch/disaggregation/transceiver.py
- tensorrt_llm/_torch/disaggregation/native/transfer.py
- cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
|
PR_Github #61471 [ run ] triggered by Bot. Commit: |
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
|
PR_Github #61471 [ run ] completed with state |
|
/bot skip --comment "all tests have passed" |
|
PR_Github #61525 [ skip ] triggered by Bot. Commit: |
|
PR_Github #61525 [ skip ] completed with state |
5 participants
Summary by CodeRabbit
New Features
TRTLLM_KVCACHE_TIME_OUTPUT_PATH.instanceIdpropagated through C++ cache sender/receiver to generate per-instance/per-rank CSV filenames.GenericLlmRequest.get_kv_cache_transfer_start()GenericLlmRequest.get_kv_cache_transfer_end()transfer_size, enabling accurate KV-cache size accumulation and logging.Bug Fixes / Correctness Improvements
Tests
_sync_transfer_timingmerge behavior and gating (tests/unittest/disaggregated/test_sync_transfer_timing.py).PerfLogManager.log_gen_transfer_summaryCSV creation and disabled/no-op behavior (tests/unittest/disaggregated/test_perf_logger.py).transfer_sizein the KV result binary prefix (tests/unittest/disaggregated/test_bounce.py).rank_*pattern support (tests/unittest/disaggregated/test_cache_transceiver_harness_report.py).tests/integration/defs/disaggregated/test_disaggregated.py).tests/integration/test_lists/test-db/l0_h100.yml).Dev Engineer Review
Correctness
requestId, and write merged timing/size back into the corresponding local request objects.transfer_sizeis carried end-to-end (sender perf-derived value → receiver unpack →kv_cache_size_bytesaccumulation). Generation-side requestkv_cache_size/transfer_endare stamped from the session state before subsequent processing.TRTLLM_KVCACHE_TIME_OUTPUT_PATH, use deterministic{instanceId}_{rank}naming, and are serialized with a mutex to avoid concurrent write corruption.API / Consistency
instanceIdis consistently threaded through C++ sender/receiver construction and used for output naming, with Python-side logging using the same C++-set naming scheme when enabled.Performance / Concurrency
(requestId, start, end, size)tuples for local completed requestsError Handling / Robustness
TRTLLM_KVCACHE_TIME_OUTPUT_PATHis not set.Potential risks to review
transfer_size: ensure all producers/consumers agree on the extended KV agent result frame and that older behavior correctly defaultstransfer_size(tests cover framing/prefix roundtrip, but multi-version interoperability should be confirmed if relevant).rank_*patterns—verify no remaining call sites assume the old naming.QA Engineer Review
Test changes detected
Unit test code (outside test-list files)
tests/unittest/disaggregated/test_perf_logger.pytest_log_gen_transfer_summary_writes_csvtest_log_gen_transfer_summary_disabled_without_envtest_cpp_env_takes_priority_over_legacy_for_logtests/unittest/disaggregated/test_sync_transfer_timing.py(new)TestSyncTransferTimingvalidating:TRTLLM_KVCACHE_TIME_OUTPUT_PATHis set and sync is required (_gen_need_sync)tests/unittest/disaggregated/test_bounce.pytest_kv_result_prefix_roundtripto include packedsizetest_make_kv_result_msg_uses_binary_frameto assert encodedtransfer_sizetests/unittest/disaggregated/test_cache_transceiver_harness_report.pyrank_*parsing.Integration test code
tests/integration/defs/disaggregated/test_disaggregated.pytest_disaggregated_kv_cache_time_outputto discover UUID-prefixed*_send.csv/*_recv.csvartifacts dynamically.Integration test list
tests/integration/test_lists/test-db/l0_h100.ymlunittest/disaggregated/test_sync_transfer_timing.py.Coverage / Verdict
test_sync_transfer_timing.pyis included intests/integration/test_lists/test-db/l0_h100.yml.Adds cross-rank time synchronization for KV-cache transfer metrics in disaggregated serving, so the per-request transfer timing recorded on different ranks/nodes becomes directly comparable and can be dumped for offline analysis.
Motivation: in disaggregated serving a single KV-cache transfer spans multiple ranks (TP/DP/PP). Each rank records its own transfer_start/transfer_end, but the local steady_clock epochs differ across processes/nodes, so the raw timestamps are not comparable, and the previous per-request all-gather to reconcile them was expensive at high request counts.
What this PR does:
The whole feature is gated by the env var TRTLLM_KVCACHE_TIME_OUTPUT_PATH (set it to an output directory to enable; unset = complete no-op, no effect on inference). perf_logger.py also now recognizes this C++-standard env var as the primary switch, falling back to the existing TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO / TLLM_KV_TRANSFER_PERF_LOG_FILE.
This is a performance-observability feature only — it does not change inference numerics.
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.