fix(container): prevent mixed UCX linkage in SGLang arm64 runtime

[Chokoyo]

Overview:

Fixes a UCX library conflict in recent SGLang runtime (nvcr.io/nvstaging/ai-dynamo/sglang-runtime:1.4.0-rc.0, upstream lmsysorg/sglang:v0.5.16-cu130-runtime) that causes multimodal encode-to-PD NIXL reads to fail with NIXL_ERR_BACKEND on GB200.

Verified on GB200: the upstream SGLang container resolves NVSHMEM’s generic UCX dependencies to Ubuntu UCX 1.16 while NIXL loads its bundled UCX 1.21, creating an in-process UCX conflict.

Details:

• Discover the installed nixl-cu* wheel and its auditwheel-renamed UCX libraries.
• Add generic libucp.so.* and libucs.so.* aliases beside NIXL’s private libraries.
• Expose that directory through a stable runtime path while preserving UCX’s $ORIGIN dependency and module lookup.
• Validate that NVSHMEM’s UCX transport and NIXL’s CUDA modules resolve without missing or distro UCX dependencies.
• Leave XPU images unchanged and avoid modifying existing wheel files or ELF metadata.

Validation

Validated on GB200 (aarch64) with Qwen/Qwen3-VL-2B-Instruct, multimodal E/P/D in --single-gpu mode, the default nixl-read embedding-transfer path, and the video from DYN-3695.

With this PR's compatibility layout:

• Two consecutive cache-enabled video requests returned HTTP 200.
• The second request logged Embedding cache hit for VIDEO URL index 0, covering both the initial population and cache-hit paths.
• A separate cache-disabled control also returned HTTP 200.
• The encode process still loaded nvshmem_transport_ucx.so.3, but mapped only NIXL's UCX 1.21 core libraries and CUDA modules.
• No system UCX mappings, UCX CUDA-support warnings, remote-key unpack failures, or NIXL_ERR_BACKEND errors were observed.

Where should the reviewer start?

• container/deps/sglang/install_nixl_ucx_compat.sh — installs the aliases and validates the resulting dependency graph.
• container/deps/sglang/discover_nixl_ucx_layout.py — discovers the NIXL and NVSHMEM wheel layout without hard-coding CUDA versions or library paths.
• container/templates/sglang_runtime.Dockerfile — wires the compatibility layout into non-XPU SGLang images.

Related Issues

• close DYN-3695

Summary by CodeRabbit

• New Features

  • Added automatic NIXL and UCX compatibility setup for CUDA-based runtime images.
  • Exposes required UCX libraries and transport plugins through a stable runtime path.
  • Supports compatible NVSHMEM transport discovery when available.

• Bug Fixes

  • Added validation for missing, ambiguous, inconsistent, or unresolved CUDA and UCX components.
  • Provides descriptive installation errors and warnings when required components are unavailable.
  • XPU runtime images remain unchanged.

[coderabbitai]

Walkthrough

The change adds NIXL and NVSHMEM UCX layout discovery, installs validated UCX compatibility aliases, checks dependency resolution, and enables the setup in non-XPU SGLang runtime images.

Changes

NIXL UCX compatibility

Layer / File(s)	Summary
CUDA and UCX layout discovery container/deps/sglang/discover_nixl_ucx_layout.py	The helper discovers CUDA distributions, NVSHMEM plugins, private UCX libraries, and SONAME aliases. It emits the layout as tab-separated records.
Compatibility links and dependency validation container/deps/sglang/install_nixl_ucx_compat.sh	The installer creates or verifies UCX aliases and a stable output link. It validates plugin and CUDA transport dependencies against private UCX libraries.
Runtime image integration container/templates/sglang_runtime.Dockerfile	Non-XPU images run the installer and prepend /opt/dynamo/nixl-ucx-compat to LD_LIBRARY_PATH.

Estimated code review effort: 3 (Moderate) | ~25 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Title check	✅ Passed	The title clearly identifies the container fix and its purpose of preventing mixed UCX linkage in the SGLang arm64 runtime.
Description check	✅ Passed	The description covers all required sections, explains the implementation, identifies review files, and documents validation results and the related issue.

---

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

[devin-ai-integration]

Devin Review found 3 potential issues.

[datadog-official]

🎯 Code Coverage (details)
• Patch Coverage: 100.00%
• Overall Coverage: 6.30% (-42.60%)

_{This comment will be updated automatically if new data arrives.
🔗 Commit SHA: 003d4b6 | Docs | Datadog PR Page | Give us feedback!}

[dmitry-tokarev-nv]

Reviewed the current head (b9f10fa) alongside DYN-3695 and the existing bot threads. The Devin findings on the not found scope and the dev/local-dev gate are already resolved, so the inline comments below are only the points not raised elsewhere.

The mechanism looks right to me: resolving the generic SONAME to the file NIXL already mapped lets glibc dedupe by (st_dev, st_ino) and reuse the existing mapping instead of creating a second UCX, and keeping the alias in the same directory to preserve $ORIGIN is the correct call — the comment explaining that is the comment a future reader needs.

Three PR-level asks beyond the inline notes:

1. Validation against the ticket's own acceptance criterion. DYN-3695 names test_sglang_serve_deployment[multimodal_epd_video_embedding_cache] (dynamo-qa) as "the only failing case in the 1.4.0-rc.0 sanity run without an owning bug." This PR verifies the linkage (NVSHMEM resolving to Ubuntu UCX 1.16 while NIXL loads its bundled 1.21), but doesn't state that the repro now returns 200 or that the QA case passes. Could you add a Validation section (also required by AGENTS.md) with that result?

2. Flag-specificity isn't explained by the stated mechanism. Mixed linkage is a static property of the image — present on every launch. DYN-3695 documents the identical launch without --multimodal-embedding-cache-capacity-gb passing on the same GB200 node. There is a plausible reconciliation — the cache path registers the domain that becomes md[1], and only that domain needs the CUDA support the mis-linked UCX lacks — but the ticket explicitly leaves root cause unestablished and prescribes the diagnostic: rerun both launches with UCX_LOG_LEVEL=debug / NIXL_LOG_LEVEL=DEBUG and diff the registered memory domains. That should settle it in one run.

Relatedly, the ucx_utils.cpp:444] UCX version is less than 1.19, CUDA support is limited line from the ticket is the strongest single piece of evidence for this fix — it is exactly what you would see if NIXL's runtime UCX were the distro 1.16 rather than its bundled 1.21. It currently lives only in Linear; it belongs in the PR body.

3. Nothing in CI will exercise this. The repro is already in-repo as tests/serve/test_sglang.py::video_e_pd_qwen (multimodal_epd.sh + --multimodal-embedding-cache-capacity-gb 0.1, pre_merge), and it runs against the runtime image this PR fixes (shared-test.yml:142 uses target_tag_plain). But pr.yaml:702,741,780 pin the GPU lanes to ["amd64"] — "No ARM GPUs available" — and the arm64 lanes are CPU-only with single-GPU tests skipped. So the failing configuration has no lane, and the build-time ldd check is the only standing regression guard. That is why the fail-open path flagged inline matters more than it first looks.

Also: DYN-3695 is tagged [release/1.4] against 1.4.0-rc.0 and this targets main — is a release/1.4.0 cherry-pick tracked anywhere?

[dmitry-tokarev-nv]

Follow-up after 69c5b315. All seven threads addressed, and the multi-arch build log settles the substantive ones. Recording the evidence in one place, since it is doing a lot of work:

#132 [linux/amd64 pre_runtime 16/29] ...
#132 0.733 nixl-ucx-compat: installed 2 aliases for nixl-cu13 UCX 1.21.0 in /usr/local/lib/python3.12/dist-packages/nixl_cu13.libs
#160 [linux/arm64 pre_runtime 16/29] ...
#160 0.805 nixl-ucx-compat: installed 2 aliases for nixl-cu13 UCX 1.21.0 in /usr/local/lib/python3.12/dist-packages/nixl_cu13.libs

That single line confirms three things at once: the new ctypes version probe works in a driver-less build container; the bundled UCX is 1.21.0, clear of nixl#1668's 1.20.0 hang; and — because the no-consumer gate now precedes alias installation — the NVSHMEM transport is discovered on both architectures.

I also had two things wrong, corrected in the threads above: the compat link targets nixl_cu13.libs (auditwheel's vendored closure), not .nixl_cu13.mesonpy.libs (meson-python's project libraries), so neither the LD_LIBRARY_PATH exposure nor the duplicate-prefix concern applies.

On the new code itself, nothing blocking. Two things I checked that are easy to get wrong and that you got right: while read … done < <(…) keeps the loop in the current shell, so die's exit 1 actually exits — a pipe there would have silently continued past a validation failure; and ((dependency_count += 1)) is safe under set -e because += 1 returns the new value. (((count++)) in that slot would abort on the first iteration — worth knowing for whoever edits it next.)

Three things still open, none of them code:

1. Validation section. AGENTS.md requires Summary and Validation in the description, and DYN-3695 names test_sglang_serve_deployment[multimodal_epd_video_embedding_cache] as its acceptance case. The description still verifies the linkage rather than the outcome — worth stating plainly that the repro now returns 200 and that the QA case passes.

2. Flag-specificity. Still unexplained by the stated mechanism: mixed linkage is a static property of the image, present on every launch, yet DYN-3695 records the identical launch without --multimodal-embedding-cache-capacity-gb passing on the same GB200 node. The ticket's own suggested diagnostic — rerun both with UCX_LOG_LEVEL=debug / NIXL_LOG_LEVEL=DEBUG and diff the registered memory domains — would settle it in one run, and would make the fix's scope defensible for the release rather than inferred.

3. release/1.4.0 cherry-pick. DYN-3695 is filed against 1.4.0-rc.0 and this targets main. Is that tracked anywhere?

Process note: 69c5b315 had a full green run including sglang-runtime / Build multi-arch cuda13.0 and sglang-runtime / Test cuda13.0, arm64, but the head is the subsequent main merge (003d4b6), whose checks are mostly skipped — it will need a fresh /ok to test 003d4b6 before merge.

[dmitry-tokarev-nv]

posted follow ups which IMO are non-blocking. approved

[Chokoyo]

Thanks @dmitry-tokarev-nv, added verification in PR description and opened a CP to release 1.4.0 at #12558