Repository navigation
[TRTLLM-11316][feat] Add kvbm disagg tests and support kvbm disagg slurm task - #13400
Draft
reasonsolo wants to merge 8 commits into
Draft
reasonsolo wants to merge 8 commits into
reasonsolo wants to merge 8 commits into
Conversation
reasonsolo
changed the title
reasonsolo
changed the title
reasonsolo
marked this pull request as ready for review
May 12, 2026 08:46
reasonsolo
marked this pull request as draft
May 12, 2026 08:49
Contributor
📝 WalkthroughWalkthroughThis PR integrates KVBM (KV cache backing module) into the disaggregated SLURM benchmarking infrastructure. It adds configuration sections for service discovery (etcd or HTTP-based), updates job submission logic to orchestrate KVBM connector injection and cluster discovery, modifies worker execution scripts to handle new server roles and cluster URIs, improves server health monitoring, adds connector logging, and introduces integration tests for KV cache backing scenarios. ChangesKVBM and Service Discovery for Disaggregated Serving
🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py (1)
1-1:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate SPDX copyright year on this modified file.
Line 1 still ends at 2024; it should reflect the current modification year.
As per coding guidelines, "Include NVIDIA copyright header on all new files; update year on modified files" and "All C++, Python, and other source files must contain NVIDIA copyright header with current modification year".
🤖 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/pyexecutor/kv_cache_connector.py` at line 1, The SPDX copyright header on the top of tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py still shows "2022-2024"; update that header line (the string beginning "# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.") to include the current modification year by replacing "2024" with "2026" so the header reads "... Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved."tests/integration/defs/disaggregated/test_disaggregated.py (1)
1-1:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate SPDX copyright year on this modified file.
The header still ends at 2024 even though this file is modified in this PR.
Suggested fix
-# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.As per coding guidelines, "Include NVIDIA copyright header on all new files; update year on modified files" and "All C++, Python, and other source files must contain NVIDIA copyright header with current modification year".
🤖 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/disaggregated/test_disaggregated.py` at line 1, Update the SPDX copyright header range in the file by changing the year suffix from "2024" to the current modification year (e.g., "2026") in the SPDX line that starts with "# SPDX-FileCopyrightText:" so the header reads "Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved."tests/integration/defs/accuracy/test_disaggregated_serving.py (1)
1-1:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winAdd NVIDIA SPDX header to this Python source file.
This modified source file is missing the required NVIDIA copyright header.
As per coding guidelines, "Include NVIDIA copyright header on all new files; update year on modified files" and "All C++, Python, and other source files must contain NVIDIA copyright header with current modification year".
🤖 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/accuracy/test_disaggregated_serving.py` at line 1, Add the required NVIDIA SPDX copyright/header comment block at the very top of this Python source (above the existing "import concurrent" line) containing the correct SPDX identifier, copyright notice and current modification year; ensure the header matches the project's standard Python header format used elsewhere and update the year if this is a modified file.
🤖 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/benchmark/disaggr_torch_dwdp.slurm`:
- Around line 20-21: The script parses --kvbm-enabled (kvbm_enabled) and
--kvbm-package (kvbm_package) but never uses them; wire these parsed values into
the job environment or launcher invocation so they take effect — e.g., if
kvbm_enabled is non-empty export KVBM_ENABLED="$kvbm_enabled" (or set an
equivalent env var) and export KVBM_PACKAGE="$kvbm_package" before the
srun/sbatch/torchrun call, or append appropriate flags/args to the command that
launches the training process (refer to the variables kvbm_enabled and
kvbm_package where arguments are constructed and ensure they are propagated).
In `@examples/disaggregated/slurm/benchmark/disaggr_torch.slurm`:
- Around line 20-21: The script uses kvbm_enabled and kvbm_package in the
argument parsing but doesn't initialize them, causing failures under set -u;
initialize kvbm_enabled and kvbm_package to sensible defaults (e.g., empty
string or "false"/"" as appropriate) before the option-parsing loop so lines
referencing kvbm_enabled and kvbm_package later (the
--kvbm-enabled/--kvbm-package handling and subsequent uses) won't trigger
unbound variable errors; update the top of the script to set kvbm_enabled="" and
kvbm_package="" (or an appropriate default) prior to the parse loop.
- Around line 147-152: The script dangerously interpolates user-controlled
variables (e.g., kvbm_package) into bash -c double-quoted commands (used with
srun and container execution like the block using container_name and
container_mount), enabling shell injection; change each invocation to pass the
user value as a positional parameter instead (use single-quoted command string
with $1 and supply the variable as an argument to srun) for kvbm_package and the
other variables mentioned (trtllm_wheel_path, trtllm_repo, build_command,
get_env_file) so the container-run bash receives the package/path/command as $1
(or $2...) instead of interpolating into the command string.
In `@examples/disaggregated/slurm/benchmark/submit.py`:
- Around line 396-398: The code currently sets sd_backend =
sd_config.get('backend', 'http') and later treats any non-"etcd" value as HTTP;
instead validate the backend value explicitly: read sd_backend from sd_config,
check that it is one of the supported values (e.g., "etcd" or "http"), and if
not raise a clear ValueError/RuntimeError with the invalid value and allowed
options; update the logic that branches on sd_backend (the places referencing
sd_backend after the sd_config assignment) to assume it's valid rather than
silently defaulting to HTTP.
- Around line 567-580: The etcd background launch in the sd_backend == 'etcd'
block (variables: etcd_port, cmd, start_server_cmds) needs a readiness probe
before server processes that use disagg_cluster_uri are started; after building
the etcd start command and before appending subsequent server startup commands,
add a looped health check (e.g., curl http://$(hostname):{etcd_port}/health or
/version) with short sleeps and a total timeout to verify etcd is responding,
fail/exit on timeout, and only then proceed to append the server startup
commands to start_server_cmds so dependent services won’t race the registry
coming up.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py`:
- Around line 507-511: The log statement in bind_connector_meta is emitting the
full metadata payload at INFO which can leak sensitive request data and add
runtime cost; change the logger call in bind_connector_meta (the site using
mpi_rank() and metadata) to avoid serializing or including repr(metadata) and
instead log only safe summary info (e.g., type(metadata).__name__ and length if
available) or move the full payload to a debug-level log behind
logger.isEnabledFor(logging.DEBUG) so payloads are never logged at INFO; ensure
the modified log uses those symbols (bind_connector_meta, mpi_rank, metadata)
and does not call repr(metadata) unconditionally.
In `@tests/integration/defs/accuracy/test_disaggregated_serving.py`:
- Around line 1514-1521: The extra_env dictionary is using hardcoded ZMQ ports
and a broad host; replace the fixed values for DYN_KVBM_TRTLLM_ZMQ_PORT (and any
other hardcoded ports like DYN_KVBM_TRTLLM_ZMQ_PORT_2 or similar found near
lines referenced) with dynamically allocated ports (e.g., call a local
utility/find_free_port() or pytest-provided port fixture to get an available
port at test runtime) and change DYN_KVBM_LEADER_ZMQ_HOST from "0.0.0.0" to the
loopback "127.0.0.1"; update all occurrences of these keys (including the other
block around the 1574-1586 region) so tests use per-run ephemeral ports and
localhost binding.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py`:
- Line 1: The SPDX copyright header on the top of
tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py still shows "2022-2024";
update that header line (the string beginning "# SPDX-FileCopyrightText:
Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.")
to include the current modification year by replacing "2024" with "2026" so the
header reads "... Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All
rights reserved."
In `@tests/integration/defs/accuracy/test_disaggregated_serving.py`:
- Line 1: Add the required NVIDIA SPDX copyright/header comment block at the
very top of this Python source (above the existing "import concurrent" line)
containing the correct SPDX identifier, copyright notice and current
modification year; ensure the header matches the project's standard Python
header format used elsewhere and update the year if this is a modified file.
In `@tests/integration/defs/disaggregated/test_disaggregated.py`:
- Line 1: Update the SPDX copyright header range in the file by changing the
year suffix from "2024" to the current modification year (e.g., "2026") in the
SPDX line that starts with "# SPDX-FileCopyrightText:" so the header reads
"Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved."
🪄 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: 0dc2b02d-50bf-48bb-8cd5-38dc3b66a3ce
📒 Files selected for processing (10)
examples/disaggregated/slurm/benchmark/config.yamlexamples/disaggregated/slurm/benchmark/disaggr_torch.slurmexamples/disaggregated/slurm/benchmark/disaggr_torch_dwdp.slurmexamples/disaggregated/slurm/benchmark/start_worker.shexamples/disaggregated/slurm/benchmark/submit.pyexamples/disaggregated/slurm/benchmark/wait_server.shrequirements-dev.txttensorrt_llm/_torch/pyexecutor/kv_cache_connector.pytests/integration/defs/accuracy/test_disaggregated_serving.pytests/integration/defs/disaggregated/test_disaggregated.py
Comment on lines
+20
to
+21
| --kvbm-enabled) kvbm_enabled="$2"; shift 2 ;; | ||
| --kvbm-package) kvbm_package="$2"; shift 2 ;; |
Contributor
There was a problem hiding this comment.
--kvbm-enabled/--kvbm-package are parsed but never applied.
Right now these flags have no effect in this launcher, so KVBM enablement can be silently ignored.
Suggested fix
+# Install KVBM package when enabled
+if [ "${kvbm_enabled:-false}" = "true" ]; then
+ echo "Installing KVBM package: ${kvbm_package}..."
+ if ! srun --container-name="${container_name}" \
+ --container-mounts="${container_mount}" --no-container-mount-home \
+ --mpi=pmix --overlap -N "${SLURM_NNODES}" --ntasks-per-node=1 \
+ bash -c "pip install ${kvbm_package}" \
+ &> "${full_logdir}/2_install_kvbm.log"; then
+ cleanup_on_failure "KVBM installation failed. Check ${full_logdir}/2_install_kvbm.log for details"
+ fi
+fi🤖 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/benchmark/disaggr_torch_dwdp.slurm` around lines
20 - 21, The script parses --kvbm-enabled (kvbm_enabled) and --kvbm-package
(kvbm_package) but never uses them; wire these parsed values into the job
environment or launcher invocation so they take effect — e.g., if kvbm_enabled
is non-empty export KVBM_ENABLED="$kvbm_enabled" (or set an equivalent env var)
and export KVBM_PACKAGE="$kvbm_package" before the srun/sbatch/torchrun call, or
append appropriate flags/args to the command that launches the training process
(refer to the variables kvbm_enabled and kvbm_package where arguments are
constructed and ensure they are propagated).
Comment on lines
+20
to
+21
| --kvbm-enabled) kvbm_enabled="$2"; shift 2 ;; | ||
| --kvbm-package) kvbm_package="$2"; shift 2 ;; |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cat -n examples/disaggregated/slurm/benchmark/disaggr_torch.slurm | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 2191
🏁 Script executed:
sed -n '140,160p' examples/disaggregated/slurm/benchmark/disaggr_torch.slurmRepository: NVIDIA/TensorRT-LLM
Length of output: 1045
Initialize the new KVBM flags before first use.
With set -u, the script will abort at lines 44-45 if a caller omits the new --kvbm-enabled or --kvbm-package flags. Initialize kvbm_enabled and kvbm_package before the parse loop to maintain backward compatibility.
Suggested fix
+kvbm_enabled="false"
+kvbm_package="kvbm"
+
# Parse named arguments
while [[ $# -gt 0 ]]; do📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| --kvbm-enabled) kvbm_enabled="$2"; shift 2 ;; | |
| --kvbm-package) kvbm_package="$2"; shift 2 ;; | |
| kvbm_enabled="false" | |
| kvbm_package="kvbm" | |
| # Parse named arguments | |
| while [[ $# -gt 0 ]]; do | |
| case "$1" in | |
| --kvbm-enabled) kvbm_enabled="$2"; shift 2 ;; | |
| --kvbm-package) kvbm_package="$2"; shift 2 ;; |
🤖 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/benchmark/disaggr_torch.slurm` around lines 20 -
21, The script uses kvbm_enabled and kvbm_package in the argument parsing but
doesn't initialize them, causing failures under set -u; initialize kvbm_enabled
and kvbm_package to sensible defaults (e.g., empty string or "false"/"" as
appropriate) before the option-parsing loop so lines referencing kvbm_enabled
and kvbm_package later (the --kvbm-enabled/--kvbm-package handling and
subsequent uses) won't trigger unbound variable errors; update the top of the
script to set kvbm_enabled="" and kvbm_package="" (or an appropriate default)
prior to the parse loop.
Comment on lines
+147
to
+152
| echo "Installing kvbm package: ${kvbm_package}..." | ||
| if ! srun --container-name=${container_name} \ | ||
| --container-mounts=${container_mount} --no-container-mount-home \ | ||
| --mpi=pmix --overlap -N $SLURM_NNODES --ntasks-per-node=1 \ | ||
| bash -c "pip install ${kvbm_package}" \ | ||
| &> ${full_logdir}/2.1_kvbm_install.log; then |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cd examples/disaggregated/slurm/benchmark && wc -l disaggr_torch.slurmRepository: NVIDIA/TensorRT-LLM
Length of output: 87
🏁 Script executed:
cd examples/disaggregated/slurm/benchmark && sed -n '140,160p' disaggr_torch.slurmRepository: NVIDIA/TensorRT-LLM
Length of output: 1045
🏁 Script executed:
cd examples/disaggregated/slurm/benchmark && rg -A 50 "kvbm_package" disaggr_torch.slurm | head -80Repository: NVIDIA/TensorRT-LLM
Length of output: 2825
🏁 Script executed:
cd examples/disaggregated/slurm/benchmark && rg -B 20 "kvbm_package" disaggr_torch.slurm | head -100Repository: NVIDIA/TensorRT-LLM
Length of output: 2850
🏁 Script executed:
cd examples/disaggregated/slurm/benchmark && grep -n 'bash -c' disaggr_torch.slurmRepository: NVIDIA/TensorRT-LLM
Length of output: 392
Pass kvbm_package to pip without shell interpolation.
The kvbm_package variable is derived from user-supplied command-line arguments and is interpolated into a bash -c command with double quotes. This allows shell metacharacters or spaces in the package spec to execute arbitrary code. Pass the package as a positional argument instead using single quotes and the $1 parameter.
Additionally, similar shell injection vulnerabilities exist on lines 95, 115, 126, and 138 using the same pattern with other user-controlled variables (trtllm_wheel_path, trtllm_repo, build_command, get_env_file). These should be fixed consistently throughout the script.
Suggested fix
if ! srun --container-name=${container_name} \
--container-mounts=${container_mount} --no-container-mount-home \
--mpi=pmix --overlap -N $SLURM_NNODES --ntasks-per-node=1 \
- bash -c "pip install ${kvbm_package}" \
+ bash -c 'python3 -m pip install -- "$1"' _ "${kvbm_package}" \
&> ${full_logdir}/2.1_kvbm_install.log; then📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| echo "Installing kvbm package: ${kvbm_package}..." | |
| if ! srun --container-name=${container_name} \ | |
| --container-mounts=${container_mount} --no-container-mount-home \ | |
| --mpi=pmix --overlap -N $SLURM_NNODES --ntasks-per-node=1 \ | |
| bash -c "pip install ${kvbm_package}" \ | |
| &> ${full_logdir}/2.1_kvbm_install.log; then | |
| echo "Installing kvbm package: ${kvbm_package}..." | |
| if ! srun --container-name=${container_name} \ | |
| --container-mounts=${container_mount} --no-container-mount-home \ | |
| --mpi=pmix --overlap -N $SLURM_NNODES --ntasks-per-node=1 \ | |
| bash -c 'python3 -m pip install -- "$1"' _ "${kvbm_package}" \ | |
| &> ${full_logdir}/2.1_kvbm_install.log; then |
🤖 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/benchmark/disaggr_torch.slurm` around lines 147
- 152, The script dangerously interpolates user-controlled variables (e.g.,
kvbm_package) into bash -c double-quoted commands (used with srun and container
execution like the block using container_name and container_mount), enabling
shell injection; change each invocation to pass the user value as a positional
parameter instead (use single-quoted command string with $1 and supply the
variable as an argument to srun) for kvbm_package and the other variables
mentioned (trtllm_wheel_path, trtllm_repo, build_command, get_env_file) so the
container-run bash receives the package/path/command as $1 (or $2...) instead of
interpolating into the command string.
Comment on lines
+396
to
+398
| sd_config = config.get('service_discovery', {}) | ||
| sd_backend = sd_config.get('backend', 'http') | ||
|
|
Contributor
There was a problem hiding this comment.
Reject unsupported service-discovery backends explicitly.
Line 528 currently treats every non-"etcd" value as HTTP, so a typo in service_discovery.backend silently launches the wrong discovery mode. Fail fast here instead of defaulting on invalid input.
Suggested fix
sd_config = config.get('service_discovery', {})
sd_backend = sd_config.get('backend', 'http')
+ if sd_backend not in {'http', 'etcd'}:
+ raise ValueError(
+ f"Unsupported service_discovery.backend: {sd_backend!r}. "
+ "Expected 'http' or 'etcd'."
+ )Also applies to: 523-531
🤖 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/benchmark/submit.py` around lines 396 - 398, The
code currently sets sd_backend = sd_config.get('backend', 'http') and later
treats any non-"etcd" value as HTTP; instead validate the backend value
explicitly: read sd_backend from sd_config, check that it is one of the
supported values (e.g., "etcd" or "http"), and if not raise a clear
ValueError/RuntimeError with the invalid value and allowed options; update the
logic that branches on sd_backend (the places referencing sd_backend after the
sd_config assignment) to assume it's valid rather than silently defaulting to
HTTP.
Comment on lines
+567
to
+580
| # Start etcd on the disagg server node when using etcd discovery | ||
| if sd_backend == 'etcd': | ||
| etcd_port = sd_config.get('etcd_port', 2379) | ||
| cmd = [ | ||
| "srun -l", | ||
| f"--nodelist {disagg_server_hostname}", | ||
| f"--container-name={container_name}", | ||
| f"--container-mounts={container_mount_str}", | ||
| "--no-container-mount-home --mpi=pmix --overlap -N 1 -n 1", | ||
| f"bash -c 'etcd --listen-client-urls http://0.0.0.0:{etcd_port}" | ||
| f" --advertise-client-urls http://$(hostname):{etcd_port}'", | ||
| f"&> {log_dir}/2.5_etcd.log &", | ||
| ] | ||
| start_server_cmds.append(" ".join(cmd)) |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
fd submit.py | grep disaggregatedRepository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
wc -l examples/disaggregated/slurm/benchmark/submit.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 116
🏁 Script executed:
sed -n '560,600p' examples/disaggregated/slurm/benchmark/submit.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1609
🏁 Script executed:
sed -n '580,650p' examples/disaggregated/slurm/benchmark/submit.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 3033
🏁 Script executed:
grep -n "start_server_cmds" examples/disaggregated/slurm/benchmark/submit.py | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 436
🏁 Script executed:
grep -n "disagg_cluster_uri" examples/disaggregated/slurm/benchmark/submit.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 327
🏁 Script executed:
sed -n '685,750p' examples/disaggregated/slurm/benchmark/submit.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 3641
🏁 Script executed:
sed -n '750,850p' examples/disaggregated/slurm/benchmark/submit.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 3916
🏁 Script executed:
sed -n '650,690p' examples/disaggregated/slurm/benchmark/submit.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1814
🏁 Script executed:
grep -n "sleep\|wait\|time.sleep\|subprocess.run" examples/disaggregated/slurm/benchmark/submit.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 345
🏁 Script executed:
fd wait_server.sh | head -5Repository: NVIDIA/TensorRT-LLM
Length of output: 117
🏁 Script executed:
cat examples/disaggregated/slurm/benchmark/wait_server.sh 2>/dev/null | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 1149
Add readiness check for etcd before starting dependent services.
etcd is backgrounded at line 578 and immediately followed by server startup commands. Without waiting for etcd readiness, the server processes that depend on disagg_cluster_uri can fail or race against the registry coming up. Add a health check probe (e.g., curl to etcd's endpoint) after the etcd start command and before launching server processes that depend on it.
🤖 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/benchmark/submit.py` around lines 567 - 580, The
etcd background launch in the sd_backend == 'etcd' block (variables: etcd_port,
cmd, start_server_cmds) needs a readiness probe before server processes that use
disagg_cluster_uri are started; after building the etcd start command and before
appending subsequent server startup commands, add a looped health check (e.g.,
curl http://$(hostname):{etcd_port}/health or /version) with short sleeps and a
total timeout to verify etcd is responding, fail/exit on timeout, and only then
proceed to append the server startup commands to start_server_cmds so dependent
services won’t race the registry coming up.
Comment on lines
+507
to
+511
| logger.info( | ||
| f"bind_connector_meta: rank={mpi_rank()} " | ||
| f"type={type(metadata).__name__} " | ||
| f"len={len(metadata) if hasattr(metadata, '__len__') else 'N/A'} " | ||
| f"repr={repr(metadata)[:500]}") |
Contributor
There was a problem hiding this comment.
Avoid logging connector metadata payload at INFO level.
This can leak request metadata into logs and adds avoidable runtime overhead on every metadata bind.
Suggested fix
- logger.info(
- f"bind_connector_meta: rank={mpi_rank()} "
- f"type={type(metadata).__name__} "
- f"len={len(metadata) if hasattr(metadata, '__len__') else 'N/A'} "
- f"repr={repr(metadata)[:500]}")
+ logger.debug(
+ "bind_connector_meta: rank=%s type=%s len=%s",
+ mpi_rank(),
+ type(metadata).__name__,
+ len(metadata) if hasattr(metadata, "__len__") else "N/A",
+ )🤖 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/pyexecutor/kv_cache_connector.py` around lines 507 - 511,
The log statement in bind_connector_meta is emitting the full metadata payload
at INFO which can leak sensitive request data and add runtime cost; change the
logger call in bind_connector_meta (the site using mpi_rank() and metadata) to
avoid serializing or including repr(metadata) and instead log only safe summary
info (e.g., type(metadata).__name__ and length if available) or move the full
payload to a debug-level log behind logger.isEnabledFor(logging.DEBUG) so
payloads are never logged at INFO; ensure the modified log uses those symbols
(bind_connector_meta, mpi_rank, metadata) and does not call repr(metadata)
unconditionally.
Comment on lines
+1514
to
+1521
| extra_env = { | ||
| "DYN_KVBM_CPU_CACHE_GB": "32", | ||
| "DYN_KVBM_TRTLLM_ZMQ_PORT": "20081", | ||
| "DYN_KVBM_LEADER_ZMQ_HOST": "0.0.0.0", | ||
| "DYN_KVBM_NCCL_MLA_MODE": "true", | ||
| "DYN_KVBM_NCCL_MAX_CTAS": "1", | ||
| "TLLM_LOG_LEVEL": "DEBUG", | ||
| } |
Contributor
There was a problem hiding this comment.
Use dynamic ZMQ ports and loopback host for these tests.
Hardcoded ports (20081/20082) can collide in parallel runs, and 0.0.0.0 broadens exposure unnecessarily for local CI tests.
Suggested fix
+ kvbm_port = str(get_free_port())
extra_env = {
"DYN_KVBM_CPU_CACHE_GB": "32",
- "DYN_KVBM_TRTLLM_ZMQ_PORT": "20081",
- "DYN_KVBM_LEADER_ZMQ_HOST": "0.0.0.0",
+ "DYN_KVBM_TRTLLM_ZMQ_PORT": kvbm_port,
+ "DYN_KVBM_LEADER_ZMQ_HOST": "127.0.0.1",
"DYN_KVBM_NCCL_MLA_MODE": "true",
"DYN_KVBM_NCCL_MAX_CTAS": "1",
"TLLM_LOG_LEVEL": "DEBUG",
}+ kvbm_port = str(get_free_port())
extra_env = {
"DYN_KVBM_CPU_CACHE_GB": "8",
"DYN_KVBM_DISK_CACHE_GB": "16",
"DYN_KVBM_DISK_CACHE_DIR": disk_cache_dir,
@@
- "DYN_KVBM_TRTLLM_ZMQ_PORT": "20082",
- "DYN_KVBM_LEADER_ZMQ_HOST": "0.0.0.0",
+ "DYN_KVBM_TRTLLM_ZMQ_PORT": kvbm_port,
+ "DYN_KVBM_LEADER_ZMQ_HOST": "127.0.0.1",
"RUST_LOG": "warn,dynamo_llm::block_manager=warn",
}Also applies to: 1574-1586
🧰 Tools
🪛 Ruff (0.15.12)
[error] 1517-1517: Possible binding to all interfaces
(S104)
🤖 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/accuracy/test_disaggregated_serving.py` around lines
1514 - 1521, The extra_env dictionary is using hardcoded ZMQ ports and a broad
host; replace the fixed values for DYN_KVBM_TRTLLM_ZMQ_PORT (and any other
hardcoded ports like DYN_KVBM_TRTLLM_ZMQ_PORT_2 or similar found near lines
referenced) with dynamically allocated ports (e.g., call a local
utility/find_free_port() or pytest-provided port fixture to get an available
port at test runtime) and change DYN_KVBM_LEADER_ZMQ_HOST from "0.0.0.0" to the
loopback "127.0.0.1"; update all occurrences of these keys (including the other
block around the 1574-1586 region) so tests use per-run ephemeral ports and
localhost binding.
reasonsolo
force-pushed
the
kvbm
branch
2 times, most recently
from
May 12, 2026 11:32
bdd9842 to
b6099dd
Compare
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Expose disk cache size and path through the LLM API config and bridge them into the KV cache manager v2 disk tier. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…ics (NVIDIA#14506) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Collaborator
|
@reasonsolo should this PR be closed? |
4 tasks done
1 task done
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Improvements
Tests
Chores
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)
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.