[https://nvbugs/6541356][fix] align cache transceiver harness setup - #17121
[https://nvbugs/6541356][fix] align cache transceiver harness setup#17121chuangz0 wants to merge 5 commits into
Conversation
chuangz0
changed the title
458ac02 to
ce1f8d9
Compare
|
/bot run --stage-list "GB300-56_GPUs-14_Nodes-PyTorch-Disagg-PerfSanity-CTX12-NODE1-GPU4-GEN1-NODE2-GPU8-Post-Merge*" |
|
PR_Github #63016 [ run ] triggered by Bot. Commit: |
|
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:
WalkthroughThe change centralizes fabric-memory setup for Python transceivers, applies it before KV-cache allocation, propagates shell-quoted ChangesFabric-memory setup and precheck integration
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant PerfSubmit
participant PrecheckConfig
participant run_precheck
participant FabricMemoryHelper
participant KVCacheManager
PerfSubmit->>PrecheckConfig: pass shell-quoted LLM_MODELS_ROOT
PrecheckConfig->>run_precheck: generate context and generation commands
run_precheck->>FabricMemoryHelper: configure resolved transceiver and cache manager
FabricMemoryHelper->>KVCacheManager: set fabric-memory default before construction
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py (1)
38-46: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the public helper parameters.
Add a Google-style
Args:section forcache_transceiver_configandkv_cache_manager_cls. This helper is imported outside its defining module.Proposed documentation update
"""Default the C++ V1 KV pool to fabric memory for the Python transceiver. This must run before any KV pool allocation because the C++ environment getter caches the value on first read. Explicit user settings are always respected. + + Args: + cache_transceiver_config: Configuration that selects the transceiver. + kv_cache_manager_cls: KV-cache manager class to evaluate. """🤖 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_transceiver.py` around lines 38 - 46, Update the docstring for maybe_enable_fabric_memory_for_python_transceiver with a Google-style Args: section documenting cache_transceiver_config and kv_cache_manager_cls, including their roles and expected types. Preserve the existing behavior and explanatory text.Source: Coding guidelines
tests/unittest/others/test_kv_cache_transceiver.py (1)
47-67: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the new test to
tests/integration/test_lists/test-db/l0_sanity_check.yml.The test covers four parameterized cases, but only other tests from
test_kv_cache_transceiver.pyare listed. Coverage verdict: insufficient.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/others/test_kv_cache_transceiver.py` around lines 47 - 67, Add test_maybe_enable_fabric_memory_for_python_transceiver from test_kv_cache_transceiver.py to tests/integration/test_lists/test-db/l0_sanity_check.yml so all four parameterized cases are included in the integration test list.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 `@tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py`:
- Line 167: The shlex.quote() call on line 167 produces single-quote escaping
that breaks when the quoted value is embedded in the double-quoted export
statements on lines 177-180. Replace the quoting strategy for model_root_env to
survive embedding in double-quoted strings, either by passing LLM_MODELS_ROOT as
a runtime environment variable without shell string interpolation or by quoting
the entire export value instead of the path alone. Add a regression test that
validates the fix with a path containing an apostrophe.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py`:
- Around line 38-46: Update the docstring for
maybe_enable_fabric_memory_for_python_transceiver with a Google-style Args:
section documenting cache_transceiver_config and kv_cache_manager_cls, including
their roles and expected types. Preserve the existing behavior and explanatory
text.
In `@tests/unittest/others/test_kv_cache_transceiver.py`:
- Around line 47-67: Add test_maybe_enable_fabric_memory_for_python_transceiver
from test_kv_cache_transceiver.py to
tests/integration/test_lists/test-db/l0_sanity_check.yml so all four
parameterized cases are included in the integration test list.
🪄 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: 49468c21-5a6c-412d-820a-cb887e0184ef
📒 Files selected for processing (7)
examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.pytests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.pytests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.pytests/unittest/others/test_cache_transceiver_precheck_config.pytests/unittest/others/test_kv_cache_transceiver.py
|
PR_Github #63016 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #63342 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/others/test_cache_transceiver_precheck_config.py`:
- Around line 368-376: Update the test setup before the precheck_prefix_lines
call to also remove the TRTLLM_DISAGG_CT_PRECHECK environment variable with
monkeypatch, while preserving the existing LLM_MODELS_ROOT cleanup and test
arguments.
In `@tests/unittest/scripts/test_perf_submit.py`:
- Around line 103-123: Add a test near test_ci_missing_model_root_is_detectable
that passes an unterminated-quote pytestCommand prefix to
ci_submit_module._get_pytest_command_env_var and asserts it raises ValueError.
Keep the existing valid-command and missing-variable coverage 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: 0e901ae8-3a8a-490a-a39c-a64512c2bbcc
📒 Files selected for processing (8)
jenkins/scripts/perf/cluster_env.pyjenkins/scripts/perf/local/submit.pyjenkins/scripts/perf/submit.pytests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.pytests/unittest/others/test_cache_transceiver_precheck_config.pytests/unittest/others/test_cache_transceiver_precheck_run.pytests/unittest/scripts/test_cluster_env.pytests/unittest/scripts/test_perf_submit.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py
d26e663 to
f6e7500
Compare
|
PR_Github #63342 [ run ] completed with state
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/others/test_cache_transceiver_precheck_config.py (1)
365-365: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winApply the function-annotation rule consistently to both new tests.
tests/unittest/others/test_cache_transceiver_precheck_config.py#L365-L365: annotate the fixture and model-root parameters and returnNone, or verify an explicit permissive typing exception.tests/unittest/others/test_cache_transceiver_precheck_run.py#L400-L400: annotate the fixture parameters and returnNone, or verify an explicit permissive typing exception.As per coding guidelines, Python functions must be annotated. Based on learnings, permissive untyped test methods require an explicit submodule typing configuration.
🤖 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/others/test_cache_transceiver_precheck_config.py` at line 365, Annotate the new test function test_precheck_commands_propagate_model_root in tests/unittest/others/test_cache_transceiver_precheck_config.py at lines 365-365 with types for its fixture and model-root parameters and a None return type, or add the required explicit permissive typing configuration. Apply the same fixture-parameter and None-return annotations, or equivalent typing exception, to the new test at tests/unittest/others/test_cache_transceiver_precheck_run.py lines 400-400.Sources: Coding guidelines, Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unittest/others/test_cache_transceiver_precheck_config.py`:
- Line 365: Annotate the new test function
test_precheck_commands_propagate_model_root in
tests/unittest/others/test_cache_transceiver_precheck_config.py at lines 365-365
with types for its fixture and model-root parameters and a None return type, or
add the required explicit permissive typing configuration. Apply the same
fixture-parameter and None-return annotations, or equivalent typing exception,
to the new test at tests/unittest/others/test_cache_transceiver_precheck_run.py
lines 400-400.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 082e16ca-d7db-4e84-b872-c679fca70dd2
📒 Files selected for processing (13)
examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.pyjenkins/scripts/perf/cluster_env.pyjenkins/scripts/perf/local/submit.pyjenkins/scripts/perf/submit.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.pytests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.pytests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.pytests/unittest/others/test_cache_transceiver_precheck_config.pytests/unittest/others/test_cache_transceiver_precheck_run.pytests/unittest/others/test_kv_cache_transceiver.pytests/unittest/scripts/test_cluster_env.pytests/unittest/scripts/test_perf_submit.py
🚧 Files skipped from review as they are similar to previous changes (11)
- tests/unittest/others/test_kv_cache_transceiver.py
- jenkins/scripts/perf/local/submit.py
- tests/unittest/scripts/test_cluster_env.py
- tensorrt_llm/_torch/pyexecutor/_util.py
- jenkins/scripts/perf/cluster_env.py
- tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
- tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py
- tests/unittest/scripts/test_perf_submit.py
- jenkins/scripts/perf/submit.py
- examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py
- tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py
|
/bot run --stage-list "GB300-56_GPUs-14_Nodes-PyTorch-Disagg-PerfSanity-CTX12-NODE1-GPU4-GEN1-NODE2-GPU8-Post-Merge*" |
|
PR_Github #63388 [ run ] triggered by Bot. Commit: |
|
PR_Github #63388 [ run ] completed with state
|
| "UCX_NET_DEVICES=" | ||
| "rdma_vf_rail0:1,rdma_vf_rail1:1,rdma_vf_rail2:1,rdma_vf_rail3:1 " | ||
| "UCX_IB_TRAFFIC_CLASS=96 TRTLLM_NIXL_NUM_THREADS=1", | ||
| "export UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp UCX_TCP_AF_PRIO=inet", |
There was a problem hiding this comment.
Could you clarify what motivated replacing the oci-aga RDMA/VF configuration with this TCP-based transport allowlist?
There was a problem hiding this comment.
On oci-aga, our previous environment variables worked in the standalone cache transceiver test, but fail in the real benchmark . After testing, the following set is confirmed to work on this cluster:
export UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp
export UCX_TCP_AF_PRIO=inet
in our test, MPI (pml_ucx) and the KV transceiver fail to create UCX workers inside the serve worker processes on this cluster.
the rail VFs only carry IPv6 addresses that TCP cannot bind to, so UCX must be forced to use IPv4 via UCX_TCP_AF_PRIO=inet.
ce7c69f to
b9b831b
Compare
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 `@tests/unittest/others/test_cache_transceiver_precheck_config.py`:
- Line 399: Update the assertion in the command-building test to verify that the
LLM_MODELS_ROOT assignment token appears before the executable token, rather
than only checking token presence. Preserve the existing expected value and use
the command tokens produced by shlex.split.
🪄 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: bcaa1360-d8b7-4194-a596-66c91adfbf3d
📒 Files selected for processing (14)
examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.pyjenkins/scripts/perf/cluster_env.pyjenkins/scripts/perf/local/submit.pyjenkins/scripts/perf/submit.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.pytests/integration/test_lists/test-db/l0_sanity_check.ymltests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.pytests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.pytests/unittest/others/test_cache_transceiver_precheck_config.pytests/unittest/others/test_cache_transceiver_precheck_run.pytests/unittest/others/test_kv_cache_transceiver.pytests/unittest/scripts/test_cluster_env.pytests/unittest/scripts/test_perf_submit.py
🚧 Files skipped from review as they are similar to previous changes (12)
- jenkins/scripts/perf/cluster_env.py
- tests/integration/test_lists/test-db/l0_sanity_check.yml
- jenkins/scripts/perf/local/submit.py
- jenkins/scripts/perf/submit.py
- tests/unittest/others/test_kv_cache_transceiver.py
- tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py
- tests/unittest/scripts/test_perf_submit.py
- tests/unittest/scripts/test_cluster_env.py
- tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
- tensorrt_llm/_torch/pyexecutor/_util.py
- tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py
- examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py
|
|
||
| assert len(commands) == 2 | ||
| for command in commands: | ||
| assert f"LLM_MODELS_ROOT={model_root}" in shlex.split(command) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Assert that the assignment precedes the executable.
Line 399 checks only token presence. A regression that moves LLM_MODELS_ROOT=... after python3 would still pass, but the shell would treat it as an argument instead of an environment assignment.
Suggested assertion
- assert f"LLM_MODELS_ROOT={model_root}" in shlex.split(command)
+ tokens = shlex.split(command)
+ assignment = f"LLM_MODELS_ROOT={model_root}"
+ assert assignment in tokens
+ assert tokens.index(assignment) < tokens.index("python3")📝 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.
| assert f"LLM_MODELS_ROOT={model_root}" in shlex.split(command) | |
| tokens = shlex.split(command) | |
| assignment = f"LLM_MODELS_ROOT={model_root}" | |
| assert assignment in tokens | |
| assert tokens.index(assignment) < tokens.index("python3") |
🤖 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/others/test_cache_transceiver_precheck_config.py` at line 399,
Update the assertion in the command-building test to verify that the
LLM_MODELS_ROOT assignment token appears before the executable token, rather
than only checking token presence. Preserve the existing expected value and use
the command tokens produced by shlex.split.
|
/bot run --stage-list "GB300-56_GPUs-14_Nodes-PyTorch-Disagg-PerfSanity-CTX12-NODE1-GPU4-GEN1-NODE2-GPU8-Post-Merge*, DGX_B200-16_GPUs-2_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU4-GEN1-NODE1-GPU8-Post-Merge-1" |
|
PR_Github #63628 [ run ] triggered by Bot. Commit: |
|
PR_Github #63628 [ 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>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
4652eda to
69b17af
Compare
|
/bot run --stage-list "GB300-56_GPUs-14_Nodes-PyTorch-Disagg-PerfSanity-CTX12-NODE1-GPU4-GEN1-NODE2-GPU8-Post-Merge*, DGX_B200-16_GPUs-2_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU4-GEN1-NODE1-GPU8-Post-Merge-1" |
|
PR_Github #63741 [ run ] triggered by Bot. Commit: |
|
PR_Github #63741 [ run ] completed with state
|
5 participants
Description
Fixes NVBug 6541356.
LLM_MODELS_ROOTandPYTEST_COMMON_VARSto the cache-transceiver precheck command so automatic KV-cache-manager selection can load the model configuration instead of silently falling back to V1.examples/disaggregated/slurm/cache_transceiver_testbefore KV-pool allocation.TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORYvalues and leave C++ transceiver and KV cache manager V2 behavior unchanged.Root cause
The generated precheck command did not forward the model-root/common environment. Consequently, auto manager resolution could not read the model configuration and selected cache manager V1. In addition, the standalone precheck and SLURM transceiver harness instantiate KV-cache managers directly, bypassing the serving path where #16832 enables fabric memory for the Python transceiver with V1.
Validation
pre-commit run --files <changed files>pytest -q tests/unittest/others/test_cache_transceiver_precheck_config.py tests/unittest/others/test_cache_transceiver_precheck_run.py tests/unittest/scripts/test_perf_submit.py tests/unittest/disaggregated/test_cache_transceiver_harness_report.py -k 'not TestInternalApiContract'(104 passed, 11 deselected)python3 -m py_compilefor the changed Python modules and testsThe full internal-API/GPU E2E was not run in this clean worktree because the available local TensorRT-LLM bindings predate the current
mainAPI.Dev Engineer Review
KVCacheManagerV2behavior.LLM_MODELS_ROOTandPYTEST_COMMON_VARSwith shell-safe quoting.oci-agaUCX configuration consistently.QA Engineer Review
Test code changed in:
tests/unittest/others/test_cache_transceiver_precheck_config.pytests/unittest/others/test_kv_cache_transceiver.pytests/unittest/others/test_cache_transceiver_precheck_run.pyKVCacheManagerV2selection for DeepSeek V4.tests/unittest/scripts/test_cluster_env.pyoci-aga-cs-001UCX environment.tests/unittest/scripts/test_perf_submit.pyLLM_MODELS_ROOTfrom quoted and unquoted commands, including special characters and missing values.The fabric-memory test is listed in
tests/integration/test_lists/test-db/l0_sanity_check.yml. The other test functions are not listed intests/integration/test_lists/test-db or QA files.Targeted validation passed.
Verdict: sufficient