Skip to content

[https://nvbugs/6426834][fix] Deflake test_kv_transfer: cap NIXL progress threads, pin UCX env#16708

Merged
chuangz0 merged 4 commits into
NVIDIA:mainfrom
chuangz0:fix/kv-transfer-flaky-nixl-threads
Jul 24, 2026
Merged

[https://nvbugs/6426834][fix] Deflake test_kv_transfer: cap NIXL progress threads, pin UCX env#16708
chuangz0 merged 4 commits into
NVIDIA:mainfrom
chuangz0:fix/kv-transfer-flaky-nixl-threads

Conversation

@chuangz0

@chuangz0 chuangz0 commented Jul 22, 2026

Copy link
Copy Markdown
Collaborator

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):

  • Set TRTLLM_NIXL_NUM_THREADS=1 with an explicit assignment (not setdefault, so an externally supplied value cannot defeat the cap): 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 with an explicit assignment instead of the previous setdefault, which was inert in CI (the CI agent bootstrap exports UCX_TLS before pytest starts).

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

  • Standardizes UCX/NIXL test configuration across KV-transfer and transceiver tests:
    • Forces UCX_TLS=^ib,gdr_copy.
    • Limits TRTLLM_NIXL_NUM_THREADS to 1 to prevent CPU oversubscription.
  • Configuration changes are consistent and scoped to tests using real UCX/NIXL agents.
  • The affected KV-transfer test waiver was removed; the unrelated proxy shutdown test was added to waives.txt with valid formatting.
  • No public APIs or test logic were changed. CI completed successfully, including 40 whole-file runs on an H100 node.

QA Engineer Review

  • Modified test modules configure runtime environments; no test functions were added, modified, or removed.
  • The previously waived unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp4_pp1_to_tp2_pp2] is now un-waived.
  • Test coverage is validated by the successful CI pipeline and 40 repeated whole-file runs.
  • Verdict: sufficient

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.

@chuangz0

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@coderabbitai

coderabbitai Bot commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

Walkthrough

The changes make UCX and NIXL environment settings deterministic across disaggregated, transceiver, and cache tests, and replace one integration-test waiver entry.

Changes

Test environment configuration

Layer / File(s) Summary
Disaggregated transport setup
tests/unittest/disaggregated/test_kv_transfer.py, tests/unittest/disaggregated/test_kv_transfer_mp.py, tests/unittest/disaggregated/test_mamba_transfer.py, tests/integration/test_lists/waives.txt
Disaggregated tests force UCX_TLS to ^ib,gdr_copy and TRTLLM_NIXL_NUM_THREADS to 1; the waiver changes from a KV-transfer case to a proxy-session shutdown test.
Transceiver environment setup
tests/unittest/disaggregated/test_agent*.py, tests/unittest/disaggregated/test_cache_transceiver_single_process.py, tests/unittest/disaggregated/test_py_cache_transceiver_mp.py, tests/unittest/others/test_kv_cache_transceiver.py
Transceiver and cache tests override UCX_TLS and set TRTLLM_NIXL_NUM_THREADS to 1, while removing the previous 0 fallback in the single-process cache test.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title matches the main fix by capping NIXL threads and pinning UCX env to deflake tests.
Description check ✅ Passed The description explains the issue, fix, and validation, but the template sections are not fully populated in their designated places.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 9714902 and 76d642a.

📒 Files selected for processing (3)
  • tests/integration/test_lists/waives.txt
  • tests/unittest/disaggregated/test_kv_transfer.py
  • tests/unittest/disaggregated/test_kv_transfer_mp.py
💤 Files with no reviewable changes (1)
  • tests/integration/test_lists/waives.txt

Comment thread tests/unittest/disaggregated/test_kv_transfer.py Outdated
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #60883 [ run ] triggered by Bot. Commit: 76d642a Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #60883 [ run ] completed with state SUCCESS. Commit: 76d642a
/LLM/main/L0_MergeRequest_PR pipeline #49152 completed with status: 'SUCCESS'

CI Report

Link to invocation

@chuangz0
chuangz0 requested a review from Shixiaowei02 July 22, 2026 08:31
@chuangz0
chuangz0 force-pushed the fix/kv-transfer-flaky-nixl-threads branch from c032f12 to 72e39ad Compare July 22, 2026 08:31
Comment thread tests/integration/test_lists/waives.txt
@chuangz0 chuangz0 changed the title [https://nvbugs/6426834][fix] Deflake test_kv_transfer: cap NIXL progress threads, pin UCX env [https://nvbugs/6426834][fix] fix ucx env in test_kv_transfer.py Jul 22, 2026
@chuangz0 chuangz0 changed the title [https://nvbugs/6426834][fix] fix ucx env in test_kv_transfer.py [https://nvbugs/6426834][fix] Deflake test_kv_transfer: cap NIXL progress threads, pin UCX env Jul 22, 2026
@chuangz0
chuangz0 force-pushed the fix/kv-transfer-flaky-nixl-threads branch from 72e39ad to 624d927 Compare July 22, 2026 08:52

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 72e39ad and 624d927.

📒 Files selected for processing (10)
  • tests/integration/test_lists/waives.txt
  • tests/unittest/disaggregated/test_agent.py
  • tests/unittest/disaggregated/test_agent_multi_backends.py
  • tests/unittest/disaggregated/test_cache_transceiver_single_process.py
  • tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py
  • tests/unittest/disaggregated/test_kv_transfer.py
  • tests/unittest/disaggregated/test_kv_transfer_mp.py
  • tests/unittest/disaggregated/test_mamba_transfer.py
  • tests/unittest/disaggregated/test_py_cache_transceiver_mp.py
  • tests/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

Comment thread tests/unittest/disaggregated/test_kv_transfer.py
@chuangz0
chuangz0 force-pushed the fix/kv-transfer-flaky-nixl-threads branch from 624d927 to 124d808 Compare July 23, 2026 02:48

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 624d927 and 124d808.

📒 Files selected for processing (10)
  • tests/integration/test_lists/waives.txt
  • tests/unittest/disaggregated/test_agent.py
  • tests/unittest/disaggregated/test_agent_multi_backends.py
  • tests/unittest/disaggregated/test_cache_transceiver_single_process.py
  • tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py
  • tests/unittest/disaggregated/test_kv_transfer.py
  • tests/unittest/disaggregated/test_kv_transfer_mp.py
  • tests/unittest/disaggregated/test_mamba_transfer.py
  • tests/unittest/disaggregated/test_py_cache_transceiver_mp.py
  • tests/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

Comment thread tests/unittest/disaggregated/test_mamba_transfer.py
chuangz0 added 2 commits July 23, 2026 09:38
…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>
@chuangz0
chuangz0 force-pushed the fix/kv-transfer-flaky-nixl-threads branch from 124d808 to fd6764d Compare July 23, 2026 09:47

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 win

Do not silently skip one-sided missing indexer data.

If either side returns None while the other contains data, this continue makes the transfer test pass without verifying that the cache was received. Only skip when both sides are None; 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

📥 Commits

Reviewing files that changed from the base of the PR and between 124d808 and fd6764d.

📒 Files selected for processing (9)
  • tests/integration/test_lists/waives.txt
  • tests/unittest/disaggregated/test_agent.py
  • tests/unittest/disaggregated/test_agent_multi_backends.py
  • tests/unittest/disaggregated/test_cache_transceiver_single_process.py
  • tests/unittest/disaggregated/test_kv_transfer.py
  • tests/unittest/disaggregated/test_kv_transfer_mp.py
  • tests/unittest/disaggregated/test_mamba_transfer.py
  • tests/unittest/disaggregated/test_py_cache_transceiver_mp.py
  • tests/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

Comment thread tests/unittest/disaggregated/test_agent_multi_backends.py
Comment thread tests/integration/test_lists/waives.txt Outdated
Comment thread tests/unittest/disaggregated/test_agent.py
…e re-added during rebase (fixed by NVIDIA#16630)

Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
@chuangz0
chuangz0 enabled auto-merge (squash) July 24, 2026 02:47
@chuangz0

Copy link
Copy Markdown
Collaborator Author

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61474 [ skip ] triggered by Bot. Commit: 127e7f4 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61474 [ skip ] completed with state SUCCESS. Commit: 127e7f4
Skipping testing for commit 127e7f4

Link to invocation

@chuangz0
chuangz0 merged commit 7d3a4e9 into NVIDIA:main Jul 24, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants