[None][feat] Upgrade NIXL to v1.0.1 and UCX to v1.21.x - #14056
Open
nv-anants wants to merge 3 commits into
Open
Conversation
Contributor
📝 WalkthroughWalkthroughThis PR upgrades NIXL (0.9.0 → v1.0.1) and UCX (v1.20.x → v1.21.x) build tools across Docker scripts and development dependencies. UCX installation now includes symlink setup and dynamic linker registration. Test memory allocation helpers transition to string-based memory types, and hybrid cache managers are configured for disaggregated mode. ChangesBuild tooling and disaggregated memory test support
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/unittest/disaggregated/test_agent.py (1)
1-1:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd/update NVIDIA copyright header for this modified file.
Line 1 indicates this modified Python file is missing the required NVIDIA header/year update.
As per coding guidelines: “Include NVIDIA copyright header on all new files; update year on modified files.”
🤖 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_agent.py` at line 1, Add or update the required NVIDIA copyright header at the top of this modified file (tests/unittest/disaggregated/test_agent.py) so the file includes the standard header block with the current year and NVIDIA ownership; ensure the header appears before any code (before the existing "from dataclasses import dataclass, field" import) and matches the project's canonical header format used in other files.tests/unittest/others/test_kv_cache_transceiver.py (1)
1-1:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd/update NVIDIA copyright header for this modified file.
Line 1 shows this modified Python source file without the required NVIDIA header/year update.
As per coding guidelines: “All TensorRT-LLM source files (.cpp, .h, .cu, .py) must contain an NVIDIA copyright header with the year of latest meaningful modification.”
🤖 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` at line 1, The file tests/unittest/others/test_kv_cache_transceiver.py is missing the required NVIDIA copyright header/updated year; add the standard NVIDIA copyright comment block at the top of test_kv_cache_transceiver.py (above the existing imports like the "import time" line), updating the year to the latest meaningful modification and matching the project's header format and wording used in other .py files so the file complies with the repository's copyright/header guidelines.
🧹 Nitpick comments (1)
tests/unittest/others/test_kv_cache_transceiver.py (1)
333-537: QA integration test-list updates are unnecessary for this PR scope.This file change is limited to
tests/unittest/**; notests/integration/defs/**coverage matrix update is needed.As per coding guidelines: “If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional.”
🤖 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 333 - 537, The PR only modifies unit tests under tests/unittest (not integration matrix), so QA integration test-list updates are unnecessary; update the PR description and commit message to explicitly state that QA list updates are not required for this change and reference the modified tests (test_kybrid... — e.g., test_hybrid_cache_transceiver_single_process and test_hybrid_cache_transceiver_cancel_request in test_kv_cache_transceiver.py) to make the intent clear per the coding guideline. Ensure the PR description contains the phrase that QA integration defs updates are unnecessary for this unit-scope change.
🤖 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 `@docker/common/install_ucx.sh`:
- Around line 37-42: The loop that iterates "for d in /opt/hpcx/ucx
/opt/hpcx-*/ucx; do" skips symlinks due to the "[ ! -L "$d" ]" check, leaving
stale symlinks pointing to old UCX; change the conditional to act on any
existing path (e.g., use "[ -e "$d" ]" only) or explicitly treat symlinks as
removable (e.g., "[ -e "$d" ] || [ -L "$d" ]" or just "[ -e "$d" ]" and then rm
-rf "$d"), then rm -rf "$d" and recreate the symlink to "${UCX_INSTALL_PATH%/}"
so both real dirs and symlinks are replaced; update the block containing the for
loop and UCX_INSTALL_PATH usage accordingly.
---
Outside diff comments:
In `@tests/unittest/disaggregated/test_agent.py`:
- Line 1: Add or update the required NVIDIA copyright header at the top of this
modified file (tests/unittest/disaggregated/test_agent.py) so the file includes
the standard header block with the current year and NVIDIA ownership; ensure the
header appears before any code (before the existing "from dataclasses import
dataclass, field" import) and matches the project's canonical header format used
in other files.
In `@tests/unittest/others/test_kv_cache_transceiver.py`:
- Line 1: The file tests/unittest/others/test_kv_cache_transceiver.py is missing
the required NVIDIA copyright header/updated year; add the standard NVIDIA
copyright comment block at the top of test_kv_cache_transceiver.py (above the
existing imports like the "import time" line), updating the year to the latest
meaningful modification and matching the project's header format and wording
used in other .py files so the file complies with the repository's
copyright/header guidelines.
---
Nitpick comments:
In `@tests/unittest/others/test_kv_cache_transceiver.py`:
- Around line 333-537: The PR only modifies unit tests under tests/unittest (not
integration matrix), so QA integration test-list updates are unnecessary; update
the PR description and commit message to explicitly state that QA list updates
are not required for this change and reference the modified tests
(test_kybrid... — e.g., test_hybrid_cache_transceiver_single_process and
test_hybrid_cache_transceiver_cancel_request in test_kv_cache_transceiver.py) to
make the intent clear per the coding guideline. Ensure the PR description
contains the phrase that QA integration defs updates are unnecessary for this
unit-scope change.
🪄 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: 82fbdb43-18f7-47d8-b049-1b699c2fdc60
📒 Files selected for processing (5)
docker/common/install_nixl.shdocker/common/install_ucx.shrequirements-dev.txttests/unittest/disaggregated/test_agent.pytests/unittest/others/test_kv_cache_transceiver.py
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
/bot run --add-multi-gpu-test --disable-fail-fast |
Collaborator
|
PR_Github #48038 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #48039 [ run ] triggered by Bot. Commit: |
svc-trtllm-gh-bot
added
the
Community want to contribute
Collaborator
|
/bot run --add-multi-gpu-test --disable-fail-fast |
Collaborator
|
PR_Github #48049 [ run ] triggered by Bot. Commit: |
Collaborator
|
@nv-anants This PR does not update the Docker image links, so the test results may not be valid. |
Author
Thanks for letting me know @Shixiaowei02. What should i update the docker image links to? |
Collaborator
Steps
|
Collaborator
|
PR_Github #48049 [ run ] completed with state
|
Bump NIXL from 0.9.0 to v1.0.1 (requires UCX 1.21.x) and UCX from v1.20.x to v1.21.x in the dev container. After installing UCX 1.21 to /usr/local/ucx, replace any pre-existing HPC-X UCX directory shipped in the NGC PyTorch base image (/opt/hpcx*/ucx) with a symlink to /usr/local/ucx and refresh the system linker cache. This ensures PyTorch/NCCL/UCC and the NIXL stack resolve the same single UCX 1.21 binary, fixing the dual-UCX-load failure (mix of UCX 1.20 and 1.21 in the same process) that prevented the NIXL UCX backend from initializing in disagg tests. Update the matching test fixtures: switch test_agent.py to the string "VRAM" memory_type (NIXL 1.0.1 rebinds MemoryType to a C++ enum at runtime, so direct enum comparison no longer matches the strings flowing through RegMemoryDescs), and pass is_disagg=True to create_hybrid_cache_manager in test_kv_cache_transceiver.py. Signed-off-by: Anant Sharma <anants@nvidia.com>
Change the conditional in the /opt/hpcx*/ucx replacement loop from `[ -e "$d" ] && [ ! -L "$d" ]` to `[ -e "$d" ] || [ -L "$d" ]` so the script also overwrites stale or broken symlinks at $d, not just real directories. The previous condition skipped any symlink, which would leave a pre-existing symlink to an old UCX in place. The new check is idempotent: re-running the script over our own symlink simply recreates an identical symlink. Signed-off-by: Anant Sharma <anants@nvidia.com>
…its symlink
Change `rm -rf ${UCX_INSTALL_PATH}` (with trailing slash) to
`rm -rf "${UCX_INSTALL_PATH%/}"`. The NGC PyTorch base image ships
/usr/local/ucx as a symlink to /opt/hpcx/ucx; with the trailing
slash, rm followed the symlink and only emptied its target, leaving
the symlink intact. `make install --prefix=/usr/local/ucx/` then
wrote UCX 1.21 *through* the symlink into /opt/hpcx/ucx, and the
subsequent /opt/hpcx/ucx -> /usr/local/ucx replacement created an
infinite symlink loop, breaking the NIXL UCX-plugin compile with
"Too many levels of symbolic links".
Without the trailing slash, rm removes the symlink itself, `make
install` creates a fresh real directory at /usr/local/ucx, and the
HPC-X path is symlinked one-way to it. Verified in the rebuilt
dev container: only one physical UCX install, torch and the NIXL
UCX backend share /usr/local/ucx in the same process.
Signed-off-by: Anant Sharma <anants@nvidia.com>
Shixiaowei02
force-pushed
the
upgrade/nixl-v1.0.1-ucx-v1.21
branch
from
May 13, 2026 06:40
58a3f20 to
0d8502a
Compare
Collaborator
|
/bot run --stage-list "Build-Docker-Images" |
Collaborator
|
PR_Github #48121 [ run ] triggered by Bot. Commit: |
zackyoray
added a commit
to zackyoray/TensorRT-LLM
that referenced
this pull request
May 13, 2026
Cherry-picked from PR NVIDIA#14056. Two related fixes to docker/common/install_ucx.sh: - Strip the trailing slash before `rm -rf "${UCX_INSTALL_PATH%/}"` so that a pre-existing symlink at /usr/local/ucx (NGC PyTorch ships it as a symlink to /opt/hpcx/ucx) is removed instead of being followed. - After `make install`, replace any pre-existing UCX under /opt/hpcx/ucx and /opt/hpcx-*/ucx with a symlink to the freshly installed /usr/local/ucx, and register /usr/local/ucx/lib with the dynamic linker via /etc/ld.so.conf.d/ucx.conf + ldconfig, so all binaries in the container resolve the same UCX SONAMEs. Signed-off-by: Yoray Zack <62789610+zackyoray@users.noreply.github.com>
Collaborator
|
PR_Github #48121 [ run ] completed with state
|
Collaborator
|
/bot run --stage-list "Build-Docker-Images" |
Collaborator
|
PR_Github #48296 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #48296 [ run ] completed with state
|
zackyoray
added a commit
to zackyoray/TensorRT-LLM
that referenced
this pull request
May 19, 2026
Cherry-picked from PR NVIDIA#14056. Two related fixes to docker/common/install_ucx.sh: - Strip the trailing slash before `rm -rf "${UCX_INSTALL_PATH%/}"` so that a pre-existing symlink at /usr/local/ucx (NGC PyTorch ships it as a symlink to /opt/hpcx/ucx) is removed instead of being followed. - After `make install`, replace any pre-existing UCX under /opt/hpcx/ucx and /opt/hpcx-*/ucx with a symlink to the freshly installed /usr/local/ucx, and register /usr/local/ucx/lib with the dynamic linker via /etc/ld.so.conf.d/ucx.conf + ldconfig, so all binaries in the container resolve the same UCX SONAMEs. Signed-off-by: Yoray Zack <62789610+zackyoray@users.noreply.github.com>
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
6 participants
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.
Bump NIXL from 0.9.0 to v1.0.1 (requires UCX 1.21.x) and UCX from v1.20.x to v1.21.x in the dev container.
After installing UCX 1.21 to /usr/local/ucx, replace any pre-existing HPC-X UCX directory shipped in the NGC PyTorch base image (/opt/hpcx*/ucx) with a symlink to /usr/local/ucx and refresh the system linker cache. This ensures PyTorch/NCCL/UCC and the NIXL stack resolve the same single UCX 1.21 binary, fixing the dual-UCX-load failure (mix of UCX 1.20 and 1.21 in the same process) that prevented the NIXL UCX backend from initializing in disagg tests.
Update the matching test fixtures: switch test_agent.py to the string "VRAM" memory_type (NIXL 1.0.1 rebinds MemoryType to a C++ enum at runtime, so direct enum comparison no longer matches the strings flowing through RegMemoryDescs), and pass is_disagg=True to create_hybrid_cache_manager in test_kv_cache_transceiver.py.
Summary by CodeRabbit
Chores
Tests