diff --git a/docker/common/install_nixl.sh b/docker/common/install_nixl.sh index 905ff3ee87ce..f14398ced80c 100644 --- a/docker/common/install_nixl.sh +++ b/docker/common/install_nixl.sh @@ -4,7 +4,7 @@ set -ex GITHUB_URL="https://github.com" UCX_INSTALL_PATH="/usr/local/ucx/" CUDA_PATH="/usr/local/cuda" -NIXL_VERSION="0.9.0" +NIXL_VERSION="v1.0.1" NIXL_REPO="https://github.com/ai-dynamo/nixl.git" OLD_LD_LIBRARY_PATH=$LD_LIBRARY_PATH diff --git a/docker/common/install_ucx.sh b/docker/common/install_ucx.sh index 4a40679c804f..322298070cce 100644 --- a/docker/common/install_ucx.sh +++ b/docker/common/install_ucx.sh @@ -1,15 +1,20 @@ #!/bin/bash set -ex -UCX_VERSION="v1.20.x" -UCX_COMMIT="f656dbdf93e72e60b5d6ca78b9e3d9e744e789bd" +UCX_VERSION="v1.21.x" +UCX_COMMIT="167a4c6a311d9a42e30a37dcc01b8a3e73ea2826" UCX_INSTALL_PATH="/usr/local/ucx/" CUDA_PATH="/usr/local/cuda" UCX_REPO="https://github.com/openucx/ucx.git" mkdir -p /third-party-source -rm -rf ${UCX_INSTALL_PATH} +# Drop the trailing slash so this also removes a pre-existing *symlink* at +# ${UCX_INSTALL_PATH} (the NGC PyTorch base image ships /usr/local/ucx as a +# symlink to /opt/hpcx/ucx). Without this, `make install` writes through the +# symlink into /opt/hpcx/ucx, and the later /opt/hpcx/ucx -> /usr/local/ucx +# replacement creates an infinite symlink loop. +rm -rf "${UCX_INSTALL_PATH%/}" git clone -b ${UCX_VERSION} ${UCX_REPO} cd ucx git checkout ${UCX_COMMIT} @@ -30,6 +35,25 @@ cd ucx --with-dm \ --enable-mt make install -j$(nproc) +# Replace any pre-existing UCX (e.g. HPC-X's UCX shipped in the NGC PyTorch +# base image) with a symlink to our freshly-installed one, so every binary +# in the container — PyTorch, NCCL, UCC, plus the C++/Python NIXL stack — +# resolves the same UCX SONAMEs. +for d in /opt/hpcx/ucx /opt/hpcx-*/ucx; do + # `[ -e ]` catches real dirs/files; `[ -L ]` separately catches symlinks + # (including broken ones, where -e returns false). Together they replace + # any pre-existing UCX — real dir, stale symlink to an old UCX, or our + # own symlink on re-run (idempotent). + if [ -e "$d" ] || [ -L "$d" ]; then + echo "Replacing pre-existing UCX at $d with symlink to ${UCX_INSTALL_PATH%/}" + rm -rf "$d" + ln -s "${UCX_INSTALL_PATH%/}" "$d" + fi +done +# Make /usr/local/ucx/lib known to the system dynamic linker so binaries +# without an explicit RPATH (or with $ORIGIN-relative RPATH) still find it. +echo "${UCX_INSTALL_PATH%/}/lib" > /etc/ld.so.conf.d/ucx.conf +ldconfig cd .. rm -rf ucx # Remove UCX source to save space echo "export LD_LIBRARY_PATH=${UCX_INSTALL_PATH}/lib:\$LD_LIBRARY_PATH" >> "${ENV}" diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 87d1a31f9da6..0198f27f6583 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -4210,6 +4210,17 @@ def launchTestJobs(pipeline, testFilter) } echo "###### Run LLMAPI tests Start ######" + // Resolve the real tensorrt_llm install location after pip install, + // and expose UCX shared libraries shipped inside the wheel + // (tensorrt_llm/libs/ucx/*.so and libtensorrt_llm_ucx_wrapper.so) + // so dlopen can find them at test runtime. + def trtllmLibsDir = sh( + script: "python3 -c 'import os, tensorrt_llm; print(os.path.join(os.path.dirname(tensorrt_llm.__file__), \"libs\"))'", + returnStdout: true, + ).trim() + libEnv += ["LD_LIBRARY_PATH+trtllm_ucx=${trtllmLibsDir}/ucx"] + libEnv += ["LD_LIBRARY_PATH+trtllm_libs=${trtllmLibsDir}"] + def config = VANILLA_CONFIG if (cpu_arch == AARCH64_TRIPLE) { config = LINUX_AARCH64_CONFIG diff --git a/jenkins/current_image_tags.properties b/jenkins/current_image_tags.properties index 7eff9f5484eb..dc510f89c8e5 100644 --- a/jenkins/current_image_tags.properties +++ b/jenkins/current_image_tags.properties @@ -13,8 +13,8 @@ # images are adopted from PostMerge pipelines, the abbreviated commit hash is used instead. IMAGE_NAME=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm -LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.02-py3-x86_64-ubuntu24.04-trt10.15.1.29-skip-tritondevel-202605060827-13616 -LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.02-py3-sbsa-ubuntu24.04-trt10.15.1.29-skip-tritondevel-202605060827-13616 -LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-rocky8-x86_64-rocky8-py310-trt10.15.1.29-skip-tritondevel-202605060827-13616 -LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-rocky8-x86_64-rocky8-py312-trt10.15.1.29-skip-tritondevel-202605060827-13616 -LLM_SBSA_WHEEL_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-trt10.15.1.29-skip-tritondevel-202605060827-13616 +LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm-staging/tritondevel:x86_64-tritondevel-torch_skip-d697cde-github-pr-13841-38849 +LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm-staging/tritondevel:sbsa-tritondevel-torch_skip-d697cde-github-pr-13841-38849 +LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm-staging/tritondevel:x86_64-rockylinux8-torch_skip-py310-d697cde-github-pr-13841-38849 +LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm-staging/tritondevel:x86_64-rockylinux8-torch_skip-py312-d697cde-github-pr-13841-38849 +LLM_SBSA_WHEEL_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm-staging/tritondevel:sbsa-ubuntu24-torch_skip-py312-d697cde-github-pr-13841-38849 diff --git a/requirements-dev.txt b/requirements-dev.txt index 0b389f7844f3..402f2da82a5d 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -40,7 +40,7 @@ opentelemetry-semantic-conventions-ai>=0.4.1 fuzzywuzzy==0.18.0 aiperf==0.6.0 nanobind>=2.9.0 -nixl==0.9.0 +nixl==1.0.1 cupti-python>=13.0,<13.2 nvidia-cuda-cupti>=13.0,<13.2 cxxfilt diff --git a/tests/unittest/disaggregated/test_agent.py b/tests/unittest/disaggregated/test_agent.py index 134c1ea1a40f..496cd0e7c8ec 100644 --- a/tests/unittest/disaggregated/test_agent.py +++ b/tests/unittest/disaggregated/test_agent.py @@ -55,9 +55,12 @@ class MemoryManager: allocated_memory: list[torch.Tensor] = field(default_factory=list) def allocate_memory( - self, size: int, name: str, memory_type=MemoryType.VRAM, device_id: int = 0 + self, size: int, name: str, memory_type: str = "VRAM", device_id: int = 0 ) -> RegMemoryDescs: - device = torch.device(f"cuda:{device_id}" if memory_type == MemoryType.VRAM else "cpu") + # `memory_type` is the string used by RegMemoryDescs (see base/agent.py:81); + # do not compare against `MemoryType.VRAM`, which is overridden to a C++ enum + # when the C++ binding is available. + device = torch.device(f"cuda:{device_id}" if memory_type == "VRAM" else "cpu") # Allocate memory block using torch.Tensor and track it block = torch.zeros(size, dtype=torch.uint8, device=device) diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index cf9d544a9ce0..1a38fd8b8d51 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -261,6 +261,7 @@ def create_hybrid_cache_manager(mapping, mamba_layer_mask=mamba_layer_mask, mamba_cache_dtype=mamba_conv_dtype, mamba_ssm_cache_dtype=mamba_ssm_dtype, + is_disagg=True, kv_cache_config=KvCacheConfig( max_tokens=256, enable_block_reuse=False,