[TRTLLM-13168][feat] test best ucx env for cache transceiver

[chuangz0]

Summary by CodeRabbit

Release Notes

• New Features
  • Added complete KV cache transceiver bandwidth testing harness for distributed 2-node evaluation
  • Enabled Slurm-based configuration, submission, and execution of bandwidth tests with UCX transport tuning
  • Added support for test matrix definition including multiple backend/runtime combinations, request patterns, and KV cache configurations
  • Introduced automated result aggregation and JSON reporting with transport attribution and performance metrics

Description

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

[coderabbitai]

📝 Walkthrough

Walkthrough

A 2-node disaggregated KV cache transceiver bandwidth test harness is introduced, providing configuration-driven SLURM orchestration, validation, and results analysis. The harness encompasses a YAML configuration schema, Python submission CLI with validation, SLURM batch orchestration script, test driver combining MPI and ZMQ for synchronized KV transfer between context and generation roles, and an aggregation pipeline that parses per-rank bandwidth samples and UCX transport logs to rank optimal transports and generate JSON reports.

Changes

KV Cache Transceiver Test Harness

Layer / File(s)	Summary
Configuration Schema and Job Submission config.yaml, submit.py	YAML configuration schema covers SLURM job parameters, hardware layout, test matrix, KV cache sizing, and UCX environment sweep definitions. Python CLI validates required fields, TP/PP symmetry, test matrix bounds, request-id encoding limits, and runtime/backend compatibility before writing resolved config and submitting sbatch.
Report Utilities - Launch Variable and UCX Environment Emission report.py (config loading, case enumeration, launch vars)	Report script loads YAML/JSON config, enumerates flattened test cases from matrix (filtering incompatible V2/C++ combos), and emits shell-ready launch variables and per-sweep UCX export lines to support SLURM orchestration.
SLURM Launch Orchestration launch.slurm	Batch script initializes node/ZMQ port selection, performs one-time in-container TensorRT-LLM install, loops over UCX sweeps applying environment exports and launching paired context/generation srun tasks with timeout wrapping and per-rank log capture, then aggregates results via report.py.
Report Aggregation Pipeline report.py (CSV parsing, transport attribution, aggregation)	Aggregation parses per-rank C++/Python bandwidth CSVs into request-id bucketed samples, extracts UCX transport selections from proto-info logs via sweep-level and case-based parsing with timestamp correlation, merges per-sweep status records by severity ordering, computes per-GPU median/p90 bandwidth and aggregate metrics per combination/sweep, ranks best PASS sweep per combination by bandwidth, and writes results.json plus tabular summary.
Test Driver - KV Cache Manager and Request Setup run_cache_transceiver_test.py (intro, helpers, KV manager, request construction)	Driver initializes KV cache manager (V1 or V2) distributed across pipeline-parallel ranks, defines deterministic per-request seeding and per-layer block-level KV fill/verify helpers, and constructs context-only and generation-only LlmRequest objects with runtime-specific disaggregation parameter handling.
Test Driver - Transfer Orchestration and Main Loop run_cache_transceiver_test.py (transfer, CSV preservation, main)	Driver executes per-request ZMQ lockstep handshake between context and generation leader ranks, transceiver async send/receive with deterministic KV verify, MPI consensus for rank synchronization on failures, per-cell timeout via signal.alarm and HangDetector, C++ CSV preservation, JSONL status recording, and end-to-end main loop across cases/lengths with cleanup barriers.

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is incomplete; it contains only the template boilerplate with placeholders and lacks actual implementation details, design rationale, test coverage information, or filled-in description sections.	Fill in the Description section explaining what the cache transceiver test harness does and why; document test coverage for the new scripts and configuration; confirm all relevant checklist items are addressed.
Docstring Coverage	⚠️ Warning	Docstring coverage is 40.91% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically summarizes the main change: adding a test harness to find the best UCX environment configuration for cache transceiver operations.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[chuangz0]

/bot run

[tensorrt-cicd]

PR_Github #51965 [ run ] triggered by Bot. Commit: cf96b6b Link to invocation

[coderabbitai]

Actionable comments posted: 13

🧹 Nitpick comments (3)

examples/disaggregated/slurm/cache_transceiver_test/report.py (1)

53-57: ⚡ Quick win

Add explicit return annotations to the new public helpers.

This module is introducing a new CLI/helper surface, but every public function here is currently untyped. Adding concrete return types now will make the config/aggregation contracts much easier to check and maintain. As per coding guidelines Static type checking with mypy is opt-in by submodule; always annotate functions with return types (use None if function does not return).

Also applies to: 60-65, 68-88, 94-117, 123-128, 450-542, 591-609

🤖 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 `@examples/disaggregated/slurm/cache_transceiver_test/report.py` around lines
53 - 57, The public helper functions in this module (starting with load_config)
lack return type annotations; update load_config to declare an explicit return
type (e.g., dict[str, Any] or Union[dict, Any] as appropriate) and add explicit
return types to all other newly introduced public helpers in this file (the
functions in the ranges noted: 60-65, 68-88, 94-117, 123-128, 450-542, 591-609).
For each function (including load_config) import typing primitives as needed
(Any, Dict, Optional, Union, List) and annotate the def signatures with the
concrete return types (use None for functions that return nothing) so
mypy/static checks can validate the config/aggregation contracts.

examples/disaggregated/slurm/cache_transceiver_test/submit.py (1)

34-46: ⚡ Quick win

Add return type annotations to the new helpers.

These function definitions are currently unannotated, but this repo requires return types on Python functions.

As per coding guidelines, "Static type checking with mypy is opt-in by submodule; always annotate functions with return types (use None if function does not return)".

Also applies to: 126-126

🤖 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 `@examples/disaggregated/slurm/cache_transceiver_test/submit.py` around lines
34 - 46, Add explicit return type annotations to the newly added helper
functions: annotate parse_args() to return argparse.Namespace, annotate
_is_placeholder(v) to return bool, and annotate validate(cfg) to return the
appropriate type (use None if it only performs validation). Update the function
signatures for parse_args, _is_placeholder, and validate accordingly so they
conform to the repository's mypy/type annotation guideline.

examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py (1)

310-313: ⚡ Quick win

Don’t swallow cleanup failures with blanket except Exception.

These paths currently discard the only signal that free_sequence(), socket reset, or shutdown() also failed, which makes the next poisoned cell/case much harder to diagnose. Catch the expected exception types here, or at least log unexpected cleanup failures before continuing.

As per coding guidelines "Avoid broad exception handling — catch specific exceptions, not bare except:" and "When using try-except blocks, limit the except clause to the smallest set of errors possible".

Also applies to: 351-354, 476-479, 603-606

🤖 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
`@examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py`
around lines 310 - 313, The try/except around free_sequence(kvm, req, kv_handle,
use_v2) is swallowing all exceptions; change it to catch only the expected
cleanup exceptions (e.g., specific socket/error types or a custom CleanupError)
and for any other Exception log the failure before continuing (use logger.error
or logger.exception with context including req/kv_handle/kvm) so unexpected
errors aren’t silently dropped; apply the same pattern to the other cleanup
sites (the socket reset and shutdown calls referenced in the review) to either
catch specific exception classes or log unexpected exceptions.

🤖 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/launch.slurm`:
- Around line 46-63: The install logic can trigger scripts/build_wheel.py twice
in parallel because INSTALL_CMD is executed via srun across both tasks; change
the flow so building the wheel (when BUILD_WHEEL="true") happens exactly once
and the install step runs on all nodes. Concretely, when BUILD_WHEEL is true,
separate the BUILD step from the distributed srun: run the BUILD (the command
assigned to BUILD that runs scripts/build_wheel.py --clean) on a single task
(e.g., invoke srun with -N1 or guard with SLURM_PROCID/SLURM_NODEID to run only
on rank 0) to produce the wheel/artifacts in REPO, then run a subsequent srun
across all nodes that performs the INSTALL_CMD (pip install -e .[devel] or pip
install the built wheel) without re-running scripts/build_wheel.py; update
INSTALL_CMD, BUILD, and the srun invocation accordingly so BUILD_WHEEL and REPO
no longer cause duplicate concurrent builds.
- Around line 114-123: The loop currently logs CTX_RC and GEN_RC then discards
them and the final aggregation call (python3 "${HERE}/report.py" ...
--aggregate) is unchecked, so failures don't affect the batch exit code; modify
the script to capture each sweep's CTX_RC and GEN_RC and propagate any non-zero
into a cumulative EXIT_RC (e.g., set EXIT_RC=$((EXIT_RC|CTX_RC|GEN_RC)) or
similar), then run the aggregation command and update EXIT_RC with its exit
status (check the exit code of python3 "${HERE}/report.py" and OR it into
EXIT_RC); finally exit the script with exit $EXIT_RC so any
timeout/transfer/aggregation error fails the Slurm job.

In `@examples/disaggregated/slurm/cache_transceiver_test/report.py`:
- Around line 223-227: The docstrings in several functions (starting with
_is_kv_data_header) violate D205 by missing a blank line after the
single-sentence summary; for each affected multi-line docstring add a single
blank line between the summary sentence and the following description so the
summary is separated from the rest of the paragraph (apply the same fix to the
other multi-line docstrings noted in the review). Locate the functions by their
names (e.g., _is_kv_data_header) and update their triple-quoted docstrings to
include the required blank line immediately after the first sentence, preserving
existing wording and indentation.
- Around line 437-446: When reading sweep JSONL lines in the aggregation loop
(the block that opens path and iterates lines, producing rec and key, updating
merged using severity), wrap the json.loads(line) call in a try/except that
catches json.JSONDecodeError (and optionally ValueError) so a single
truncated/partial last line does not abort aggregation; on decode error, log or
silently skip that line and continue to the next line rather than raising,
preserving the rest of the results.

In
`@examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py`:
- Around line 1-2: Update the file header year in run_cache_transceiver_test.py
from "2024-2025" to include 2026 so the SPDX header reflects the latest
modification year (e.g., change "2024-2025 NVIDIA CORPORATION & AFFILIATES. All
rights reserved." to "2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights
reserved." and update the SPDX-License-Identifier line if needed).
- Around line 627-645: The timeout handler currently just resets the alarm and
socket but continues reusing the same case, which leaks state; after catching
_Timeout in the except block you must tear down the in-progress case and skip
the remainder of that case. Specifically, invoke the case teardown that frees
request state (call free_sequence(...) or the equivalent cleanup for the current
kvm/xcvr), ensure cancel_watchdog() and reset_sock() are called, and then skip
the rest of the loop for this case (e.g., set a flag or break out of the
per-case loop so subsequent req_len iterations do not reuse the tainted case)
before reaching comm.Barrier(); keep the existing record(...) and logging
behavior.
- Around line 81-82: Several multiline docstrings in
run_cache_transceiver_test.py are missing the required blank line after the
one-line summary (causing D205 failures); update the docstrings (including the
KvCacheConfig wrapper docstring that starts with "KvCacheConfig wrapper for
KVCacheManagerV2" and the docstrings around the other noted ranges at 183-184
and 374-383) by inserting a single blank line immediately after the summary line
in each multiline docstring so Ruff stops failing.
- Around line 160-161: Replace the runtime-conditional asserts with explicit
exceptions so failures aren't skipped under python -O: where you call ok =
kv.resume(torch.cuda.current_stream().cuda_stream) and currently do assert ok,
raise a RuntimeError (or ValueError) with a clear message including
req.py_request_id when ok is false; likewise replace assert n == n_cfg with an
explicit check that raises a RuntimeError including the observed n and expected
n_cfg to fail fast on MPI topology mismatches.
- Around line 495-496: Replace the manual PP ownership calculation (pp_rank =
rank // tp) with the Mapping-provided value to avoid breaking when CP
parallelism is enabled: use mapping.pp_rank wherever PP ownership is needed
(e.g., pass mapping.pp_rank into local_layer_count instead of pp_rank) so the
call local_layer_count(cfg["kv_cache"]["num_layers"], pp, pp_rank) becomes
local_layer_count(cfg["kv_cache"]["num_layers"], pp, mapping.pp_rank).

In `@examples/disaggregated/slurm/cache_transceiver_test/submit.py`:
- Around line 46-118: The validate() function currently indexes into cfg without
guarding required sections (e.g., cfg["hardware"], cfg["parallel"],
cfg["kv_cache"], cfg["test_matrix"]) causing exceptions on partial configs;
change direct indexing to safe retrievals (cfg.get("hardware", {}) etc.), verify
presence and types of required subsections before using keys, append appropriate
errors when sections or required fields are missing, and similarly guard usages
of combinations/combo lookup and later accesses around lines 128-129 so missing
or malformed entries produce validation messages instead of tracebacks.
- Around line 54-58: The validation currently checks only
environment.container_image and environment.work_dir; add the same placeholder
check for environment.container_mount so placeholder values like
"<container_mount>" are rejected early. In the same block that uses env =
cfg.get("environment", {}) and _is_placeholder(), verify
env.get("container_mount") and append an error like "environment.container_mount
must be set" to the errors list when missing or a placeholder, ensuring
submissions with the default placeholder fail validation instead of failing
later at launch.
- Line 2: Update the SPDX copyright year range in the file's header by changing
the existing SPDX comment line that currently reads "SPDX-FileCopyrightText:
Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved."
to include 2026 (e.g., "2024-2026") so the header reflects the new file
creation/modification year.
- Around line 169-170: The loop that appends slurm.get("extra_sbatch", []) items
to sbatch_cmd treats each item as a single argv token, breaking multi-token
options; change the logic in the block handling extra_sbatch so that for each
entry from slurm.get("extra_sbatch", []) you split it into arguments (use
shlex.split to respect quoted tokens) and extend sbatch_cmd with the resulting
tokens; also accept entries that are already lists by extending directly rather
than appending the whole object.

---

Nitpick comments:
In `@examples/disaggregated/slurm/cache_transceiver_test/report.py`:
- Around line 53-57: The public helper functions in this module (starting with
load_config) lack return type annotations; update load_config to declare an
explicit return type (e.g., dict[str, Any] or Union[dict, Any] as appropriate)
and add explicit return types to all other newly introduced public helpers in
this file (the functions in the ranges noted: 60-65, 68-88, 94-117, 123-128,
450-542, 591-609). For each function (including load_config) import typing
primitives as needed (Any, Dict, Optional, Union, List) and annotate the def
signatures with the concrete return types (use None for functions that return
nothing) so mypy/static checks can validate the config/aggregation contracts.

In
`@examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py`:
- Around line 310-313: The try/except around free_sequence(kvm, req, kv_handle,
use_v2) is swallowing all exceptions; change it to catch only the expected
cleanup exceptions (e.g., specific socket/error types or a custom CleanupError)
and for any other Exception log the failure before continuing (use logger.error
or logger.exception with context including req/kv_handle/kvm) so unexpected
errors aren’t silently dropped; apply the same pattern to the other cleanup
sites (the socket reset and shutdown calls referenced in the review) to either
catch specific exception classes or log unexpected exceptions.

In `@examples/disaggregated/slurm/cache_transceiver_test/submit.py`:
- Around line 34-46: Add explicit return type annotations to the newly added
helper functions: annotate parse_args() to return argparse.Namespace, annotate
_is_placeholder(v) to return bool, and annotate validate(cfg) to return the
appropriate type (use None if it only performs validation). Update the function
signatures for parse_args, _is_placeholder, and validate accordingly so they
conform to the repository's mypy/type annotation guideline.

🪄 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: c8248b1c-0ec7-4b19-994d-18ea50ca45c2

📥 Commits

Reviewing files that changed from the base of the PR and between 846d0f4 and cf96b6b.

📒 Files selected for processing (6)

• examples/disaggregated/slurm/cache_transceiver_test/README.md
• examples/disaggregated/slurm/cache_transceiver_test/config.yaml
• examples/disaggregated/slurm/cache_transceiver_test/launch.slurm
• examples/disaggregated/slurm/cache_transceiver_test/report.py
• examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py
• examples/disaggregated/slurm/cache_transceiver_test/submit.py

[tensorrt-cicd]

PR_Github #51965 [ run ] completed with state FAILURE. Commit: cf96b6b
/LLM/main/L0_MergeRequest_PR pipeline #41315 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #52049 [ run ] triggered by Bot. Commit: cd49c9e Link to invocation

[tensorrt-cicd]

PR_Github #52049 [ run ] completed with state SUCCESS. Commit: cd49c9e
/LLM/main/L0_MergeRequest_PR pipeline #41382 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #52269 [ run ] triggered by Bot. Commit: 01a4f0f Link to invocation

[tensorrt-cicd]

PR_Github #52269 [ run ] completed with state SUCCESS. Commit: 01a4f0f
/LLM/main/L0_MergeRequest_PR pipeline #41580 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #52621 [ run ] triggered by Bot. Commit: 01a4f0f Link to invocation

[tensorrt-cicd]

PR_Github #52621 [ run ] completed with state SUCCESS. Commit: 01a4f0f
/LLM/main/L0_MergeRequest_PR pipeline #41903 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #56261 [ run ] triggered by Bot. Commit: d99359c Link to invocation

[tensorrt-cicd]

PR_Github #56261 [ run ] completed with state FAILURE. Commit: d99359c
/LLM/main/L0_MergeRequest_PR pipeline #45119 (Partly Tested) 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]

/bot run --stage-list "DGX_B200-PyTorch-7, DGX_H100-PyTorch-6, H100_PCIe-AutoDeploy-1, H100_PCIe-PyTorch-Ray-1"

[tensorrt-cicd]

PR_Github #56318 [ run ] triggered by Bot. Commit: 2dea02b Link to invocation

[tensorrt-cicd]

PR_Github #56318 [ run ] completed with state SUCCESS. Commit: 2dea02b
/LLM/main/L0_MergeRequest_PR pipeline #45168 (Partly Tested) 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]

/bot run --stage-list "DGX_B200-PyTorch-7, DGX_H100-PyTorch-6, H100_PCIe-AutoDeploy-1, H100_PCIe-PyTorch-Ray-1"

[chuangz0]

/bot run --stage-list "DGX_B200-PyTorch-7, DGX_H100-PyTorch-6, H100_PCIe-AutoDeploy-1, H100_PCIe-PyTorch-Ray-1"

[tensorrt-cicd]

PR_Github #56559 [ run ] triggered by Bot. Commit: 9048817 Link to invocation

[tensorrt-cicd]

PR_Github #56559 [ run ] completed with state SUCCESS. Commit: 9048817
/LLM/main/L0_MergeRequest_PR pipeline #45394 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[chuangz0]

/bot skip --comment "all test have passed"

[tensorrt-cicd]

PR_Github #56775 [ skip ] triggered by Bot. Commit: 9048817 Link to invocation

[tensorrt-cicd]

PR_Github #56775 [ skip ] completed with state SUCCESS. Commit: 9048817
Skipping testing for commit 9048817

Link to invocation

[Shixiaowei02]

Great work. This looks impressive. Could you also include a sample output here? For example:

=== Best UCX env set per combination by per-GPU bandwidth (req_len=8000) ===
  NIXL/CPP/V1        -> put_zcopy          85.30 GB/s/GPU  [rc_mlx5]  UCX_TLS=all UCX_RNDV_SCHEME=put_zcopy UCX_CUDA_IPC_ENABLE_MNNVL=n