[https://nvbugs/6426834][fix] Deflake test_kv_transfer: cap NIXL progress threads, pin UCX env#16708
Conversation
a544a5b to
76d642a
Compare
|
/bot run --disable-fail-fast |
|
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 changes make UCX and NIXL environment settings deterministic across disaggregated, transceiver, and cache tests, and replace one integration-test waiver entry. ChangesTest environment configuration
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
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/disaggregated/test_kv_transfer.py`:
- Around line 18-25: Make the NIXL thread cap deterministic by replacing
os.environ.setdefault() with an explicit TRTLLM_NIXL_NUM_THREADS assignment to
"1" in tests/unittest/disaggregated/test_kv_transfer.py lines 18-25 and
tests/unittest/disaggregated/test_kv_transfer_mp.py lines 15-18. Both test
modules must override any externally supplied value.
🪄 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: dbc339ea-193c-47d9-992b-8b4eef1168ab
📒 Files selected for processing (3)
tests/integration/test_lists/waives.txttests/unittest/disaggregated/test_kv_transfer.pytests/unittest/disaggregated/test_kv_transfer_mp.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
|
PR_Github #60883 [ run ] triggered by Bot. Commit: |
|
PR_Github #60883 [ run ] completed with state |
c032f12 to
72e39ad
Compare
72e39ad to
624d927
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/disaggregated/test_kv_transfer.py`:
- Around line 8-20: Clear any inherited UCX_NET_DEVICES setting before
configuring transfer tests by adding os.environ.pop("UCX_NET_DEVICES", None) in
tests/unittest/disaggregated/test_kv_transfer.py lines 8-20,
tests/unittest/disaggregated/test_mamba_transfer.py lines 23-26, and
tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py lines 40-45. Apply
this alongside the existing UCX_TLS setup in each test module so UCX device
selection is not externally forced.
🪄 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: 1a32f858-98ae-4a0c-a751-f97ed626870f
📒 Files selected for processing (10)
tests/integration/test_lists/waives.txttests/unittest/disaggregated/test_agent.pytests/unittest/disaggregated/test_agent_multi_backends.pytests/unittest/disaggregated/test_cache_transceiver_single_process.pytests/unittest/disaggregated/test_deepseek_v4_kv_transfer.pytests/unittest/disaggregated/test_kv_transfer.pytests/unittest/disaggregated/test_kv_transfer_mp.pytests/unittest/disaggregated/test_mamba_transfer.pytests/unittest/disaggregated/test_py_cache_transceiver_mp.pytests/unittest/others/test_kv_cache_transceiver.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
🚧 Files skipped from review as they are similar to previous changes (6)
- tests/unittest/disaggregated/test_kv_transfer_mp.py
- tests/unittest/disaggregated/test_agent.py
- tests/unittest/disaggregated/test_agent_multi_backends.py
- tests/unittest/disaggregated/test_py_cache_transceiver_mp.py
- tests/unittest/others/test_kv_cache_transceiver.py
- tests/unittest/disaggregated/test_cache_transceiver_single_process.py
624d927 to
124d808
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/disaggregated/test_mamba_transfer.py`:
- Around line 23-27: Remove inherited UCX_NET_DEVICES before configuring UCX in
both NIXL test modules: add the environment-variable removal before the UCX
assignments in tests/unittest/disaggregated/test_mamba_transfer.py lines 23-27
and tests/unittest/disaggregated/test_agent.py lines 9-12, leaving the existing
UCX_TLS and TRTLLM_NIXL_NUM_THREADS settings 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: 32d429ce-1116-4d21-b14b-bbfdc1859042
📒 Files selected for processing (10)
tests/integration/test_lists/waives.txttests/unittest/disaggregated/test_agent.pytests/unittest/disaggregated/test_agent_multi_backends.pytests/unittest/disaggregated/test_cache_transceiver_single_process.pytests/unittest/disaggregated/test_deepseek_v4_kv_transfer.pytests/unittest/disaggregated/test_kv_transfer.pytests/unittest/disaggregated/test_kv_transfer_mp.pytests/unittest/disaggregated/test_mamba_transfer.pytests/unittest/disaggregated/test_py_cache_transceiver_mp.pytests/unittest/others/test_kv_cache_transceiver.py
🚧 Files skipped from review as they are similar to previous changes (7)
- tests/unittest/others/test_kv_cache_transceiver.py
- tests/unittest/disaggregated/test_agent_multi_backends.py
- tests/unittest/disaggregated/test_kv_transfer.py
- tests/unittest/disaggregated/test_kv_transfer_mp.py
- tests/unittest/disaggregated/test_cache_transceiver_single_process.py
- tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py
- tests/unittest/disaggregated/test_py_cache_transceiver_mp.py
…ress threads, pin UCX env test_transfer_worker_v2[tp4_pp1_to_tp2_pp2] (and other multi-worker cases) intermittently hit the 120s per-test timeout in CI. Root cause: each NIXL transfer agent spawns TRTLLM_NIXL_NUM_THREADS (default 8) busy-polling UCX progress threads, and a single case builds up to 8 TransferWorkers (64 spinning threads per pytest process). On CI nodes shared by up to 8 single-GPU Jenkins agents the resulting CPU oversubscription inflates agent construction from ~3s to ~27s each, so 8 agents alone exceed the 120s timeout depending on neighbor load. Measured on an H100 node with 4 concurrent pytest processes: 8 agents take ~178s with the default 8 progress threads vs ~23s with 1 thread. Fix, applied to test_kv_transfer.py and test_kv_transfer_mp.py: - setdefault TRTLLM_NIXL_NUM_THREADS=1: these tests verify transfer logic, not transfer-engine threading (sender-side KV_TRANSFER_NUM_THREADS worker threads are unchanged). - Force UCX_TLS=^ib,gdr_copy and drop UCX_NET_DEVICES instead of the previous setdefault, which was inert in CI (the CI agent bootstrap exports UCX_TLS before pytest starts) and broke local runs on clusters whose container runtime injects UCX_TLS/UCX_NET_DEVICES defaults. Un-waive the affected case. Validation: 8 concurrent whole-file runs x 5 iterations (40 runs) on one H100 node under CI-parity env all passed; before the fix the same setup timed out in the first iteration on every GPU. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
…t tests Address review: assign TRTLLM_NIXL_NUM_THREADS explicitly instead of setdefault so an externally supplied value cannot defeat the single-progress-thread cap, and apply the same deterministic env (UCX_TLS=^ib,gdr_copy, drop UCX_NET_DEVICES, TRTLLM_NIXL_NUM_THREADS=1) to every unit test that creates real UCX/NIXL transfer agents, not just test_kv_transfer*.py. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
…nsceiver unit tests Keep only UCX_TLS pinning and the NIXL progress-thread cap; removing UCX_NET_DEVICES from the environment is not needed. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
124d808 to
fd6764d
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/disaggregated/test_cache_transceiver_single_process.py (1)
924-925: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winDo not silently skip one-sided missing indexer data.
If either side returns
Nonewhile the other contains data, thiscontinuemakes the transfer test pass without verifying that the cache was received. Only skip when both sides areNone; otherwise raise an assertion identifying the request and layer.Proposed fix
if ctx_data is None or gen_data is None: + assert ctx_data is None and gen_data is None, ( + f"Missing indexer data at req={req_idx} layer={layer_idx}" + ) continue🤖 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_cache_transceiver_single_process.py` around lines 924 - 925, Update the conditional in the transfer test around ctx_data and gen_data so it continues only when both are None; when exactly one is missing, raise an assertion that identifies the request and layer, and preserve normal verification when both contain data.
🤖 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_agent_multi_backends.py`:
- Around line 6-10: Clear the inherited UCX_NET_DEVICES environment variable
alongside the forced UCX_TLS and TRTLLM_NIXL_NUM_THREADS settings, before the
subprocess environment is copied for launching test_agent.py. Use os.environ
removal so externally injected device selections cannot be propagated.
In `@tests/unittest/disaggregated/test_cache_transceiver_single_process.py`:
- Around line 31-34: Update the deterministic UCX/NIXL environment setup near
UCX_TLS and TRTLLM_NIXL_NUM_THREADS to unset UCX_NET_DEVICES as well, preventing
inherited cluster or CI values from overriding the test configuration.
---
Outside diff comments:
In `@tests/unittest/disaggregated/test_cache_transceiver_single_process.py`:
- Around line 924-925: Update the conditional in the transfer test around
ctx_data and gen_data so it continues only when both are None; when exactly one
is missing, raise an assertion that identifies the request and layer, and
preserve normal verification when both contain data.
🪄 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: db3fdcdc-bb35-4d2a-8dc5-a628ed4398b6
📒 Files selected for processing (9)
tests/integration/test_lists/waives.txttests/unittest/disaggregated/test_agent.pytests/unittest/disaggregated/test_agent_multi_backends.pytests/unittest/disaggregated/test_cache_transceiver_single_process.pytests/unittest/disaggregated/test_kv_transfer.pytests/unittest/disaggregated/test_kv_transfer_mp.pytests/unittest/disaggregated/test_mamba_transfer.pytests/unittest/disaggregated/test_py_cache_transceiver_mp.pytests/unittest/others/test_kv_cache_transceiver.py
🚧 Files skipped from review as they are similar to previous changes (7)
- tests/integration/test_lists/waives.txt
- tests/unittest/disaggregated/test_agent.py
- tests/unittest/disaggregated/test_py_cache_transceiver_mp.py
- tests/unittest/disaggregated/test_kv_transfer_mp.py
- tests/unittest/disaggregated/test_mamba_transfer.py
- tests/unittest/others/test_kv_cache_transceiver.py
- tests/unittest/disaggregated/test_kv_transfer.py
…e re-added during rebase (fixed by NVIDIA#16630) Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
|
/bot skip --comment "all tests have passed" |
|
PR_Github #61474 [ skip ] triggered by Bot. Commit: |
|
PR_Github #61474 [ skip ] completed with state |
test_transfer_worker_v2[tp4_pp1_to_tp2_pp2] (and other multi-worker cases) intermittently hit the 120s per-test timeout in CI. Root cause: each NIXL transfer agent spawns TRTLLM_NIXL_NUM_THREADS (default 8) busy-polling UCX progress threads, and a single case builds up to 8 TransferWorkers (64 spinning threads per pytest process). On CI nodes shared by up to 8 single-GPU Jenkins agents the resulting CPU oversubscription inflates agent construction from ~3s to ~27s each, so 8 agents alone exceed the 120s timeout depending on neighbor load.
Measured on an H100 node with 4 concurrent pytest processes: 8 agents take ~178s with the default 8 progress threads vs ~23s with 1 thread.
Fix, applied to all unit tests that create real UCX/NIXL transfer agents (test_kv_transfer.py, test_kv_transfer_mp.py, test_agent.py, test_agent_multi_backends.py, test_cache_transceiver_single_process.py, test_py_cache_transceiver_mp.py, test_deepseek_v4_kv_transfer.py, test_mamba_transfer.py, others/test_kv_cache_transceiver.py):
Un-waive the affected case.
Validation: 8 concurrent whole-file runs x 5 iterations (40 runs) on one H100 node under CI-parity env all passed; before the fix the same setup timed out in the first iteration on every GPU.
Dev Engineer Review
UCX_TLS=^ib,gdr_copy.TRTLLM_NIXL_NUM_THREADSto1to prevent CPU oversubscription.waives.txtwith valid formatting.QA Engineer Review
unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp4_pp1_to_tp2_pp2]is now un-waived.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-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.