CI: enhance NVLS tests (#1269)

ikryukov · web-flow · commit b7d4f7675f1b · 2026-03-10T13:47:59.000+02:00
## What

1. Added correctness tests for NVLS collectives
2. Updated CUDA to 13.1.1 for NVLS path
3. Use official CUDA base image
`nvcr.io/nvidia/cuda:${CUDA_VER}-devel-ubuntu24.04` (more lightweight)
4. Use official HPCX instead of building from source
5. Added smoke testing of nvlink on allocated nodes to separate node
issues from UCC

---------

Signed-off-by: Ilya Kryukov &lt;ikryukov@nvidia.com&gt;
diff --git a/.ci/Dockerfile.nvls b/.ci/Dockerfile.nvls
@@ -0,0 +1,74 @@
+ARG CUDA_VER='13.1.1'
+FROM nvcr.io/nvidia/cuda:${CUDA_VER}-devel-ubuntu24.04
+
+ARG _UID=6213
+ARG _GID=11429
+ARG _LOGIN=swx-jenkins
+ARG _GROUP=swx-jenkins
+ARG _HOME=/labhome
+ARG UCC_ENABLE_NVLS=yes
+ARG UCC_ENABLE_GTEST=no
+ARG UCC_BUILD_TLS=cuda,ucp
+
+#==============================================================================
+# Build tools
+#==============================================================================
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    sudo \
+    build-essential \
+    autoconf \
+    automake \
+    libtool \
+    numactl \
+    libnuma-dev \
+    wget \
+    ca-certificates \
+    && rm -rf /var/lib/apt/lists/*
+
+#==============================================================================
+# Install HPC-X (provides UCX + OpenMPI)
+#==============================================================================
+ARG HPCX_VERSION=v2.26
+ARG HPCX_CUDA=cuda13
+ARG HPCX_OS=ubuntu24.04
+RUN cd /tmp && \
+    HPCX_ARCH=$(uname -m) && \
+    HPCX_FILENAME="hpcx-${HPCX_VERSION}-gcc-inbox-${HPCX_OS}-${HPCX_CUDA}-${HPCX_ARCH}" && \
+    wget -q "https://content.mellanox.com/hpc/hpc-x/${HPCX_VERSION}_${HPCX_CUDA}/${HPCX_FILENAME}.tbz" && \
+    tar xf "${HPCX_FILENAME}.tbz" && \
+    mv "${HPCX_FILENAME}" /opt/hpcx && \
+    rm -f "${HPCX_FILENAME}.tbz"
+
+#==============================================================================
+# Environment
+#==============================================================================
+ENV CUDA_HOME=/usr/local/cuda
+ENV SRC_DIR=/opt/nvidia/src
+ENV UCX_INSTALL_DIR=/opt/hpcx/ucx
+ENV UCC_INSTALL_DIR=/opt/nvidia/bin/ucc/build
+ENV PATH=/opt/hpcx/ompi/bin:${PATH}
+ENV LD_LIBRARY_PATH=/opt/hpcx/ompi/lib:/opt/hpcx/ucx/lib:${LD_LIBRARY_PATH}
+ENV OPAL_PREFIX=/opt/hpcx/ompi
+
+#==============================================================================
+# Build UCC
+#==============================================================================
+RUN rm -rf ${SRC_DIR}/ucc
+COPY . ${SRC_DIR}/ucc
+
+ENV UCC_ENABLE_NVLS=${UCC_ENABLE_NVLS}
+ENV UCC_ENABLE_GTEST=${UCC_ENABLE_GTEST}
+ENV UCC_BUILD_TLS=${UCC_BUILD_TLS}
+RUN ${SRC_DIR}/ucc/.ci/scripts/build_ucc.sh
+
+#==============================================================================
+# User setup
+#==============================================================================
+RUN echo "${_LOGIN} ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers
+RUN chown -R ${_UID}:${_GID} /opt/nvidia
+RUN if ! getent group "${_GID}" > /dev/null 2>&1; then \
+        groupadd -g "${_GID}" "${_GROUP}"; \
+    fi && \
+    useradd --no-create-home --uid ${_UID} --gid ${_GID} --home ${_HOME}/${_LOGIN} ${_LOGIN}
+
+USER ${_LOGIN}
diff --git a/.ci/pipeline/test_nvls_matrix.yaml b/.ci/pipeline/test_nvls_matrix.yaml
@@ -18,7 +18,7 @@ volumes:
   - { mountPath: "/home/svcnbu-swx-hpcx", hostPath: "/labhome/svcnbu-swx-hpcx" }
 
 env:
-  CUDA_VER: 13.0
+  CUDA_VER: 13.1.1
   UCC_URI_SUFFIX: "ucc/${UCC_VERSION}/aarch64/ubuntu24.04/cuda${CUDA_VER}"
   DOCKER_IMAGE_TAG: "${BUILD_NUMBER}"
   SLURM_NODES: 2
@@ -28,6 +28,7 @@ env:
   SLURM_JOB_NAME: 'ucc-ci-test-nvls-${BUILD_NUMBER}'
   SLURM_JOB_TIMEOUT: '00:40:00'
   TEST_TIMEOUT_MINUTES: 35
+  NVLS_MPI_PPN: 4
 
 kubernetes:
   cloud: il-ipp-blossom-prod
@@ -38,12 +39,12 @@ kubernetes:
 # cloud pod to build the shared docker image
 runs_on_dockers:
   - {
-      file: ".ci/Dockerfile.ngc_pytorch",
-      name: "ngc_pytorch",
+      file: ".ci/Dockerfile.nvls",
+      name: "nvls",
       tag: "${DOCKER_IMAGE_TAG}",
       arch: "aarch64",
       uri: "${UCC_URI_SUFFIX}",
-      build_args: "--no-cache --build-arg ARCH=aarch64 --build-arg OS=ubuntu24.04 --build-arg CUDA_VER=${CUDA_VER} --build-arg _UID=149917 --build-arg _GID=30 --build-arg _LOGIN=svcnbu-swx-hpcx --build-arg _GROUP=svcnbu-swx-hpcx --build-arg UCC_ENABLE_NVLS=yes --build-arg UCC_ENABLE_GTEST=no",
+      build_args: "--no-cache --build-arg CUDA_VER=${CUDA_VER} --build-arg _UID=149917 --build-arg _GID=30 --build-arg _LOGIN=svcnbu-swx-hpcx --build-arg _GROUP=svcnbu-swx-hpcx",
     }
   - {
       file: ".ci/dockerfiles/Dockerfile.build_helper",
@@ -68,13 +69,24 @@ steps:
       touch job_id.txt && chown svcnbu-swx-hpcx job_id.txt
       sudo -E -u svcnbu-swx-hpcx ${WORKSPACE}/.ci/scripts/run_slurm_allocation.sh
 
-  - name: Run UCC NVLS tests
+  - name: Run UCC NVLS perftest
     containerSelector: "{name: 'build_helper'}"
     timeout: "${TEST_TIMEOUT_MINUTES}"
     run: |
       set -x
       export DOCKER_IMAGE_NAME="${registry_host}#torch-ucc/${UCC_URI_SUFFIX}:${DOCKER_IMAGE_TAG}"
       export SLURM_JOB_ID=$(cat ${WORKSPACE}/job_id.txt)
-      sudo -E -u svcnbu-swx-hpcx ${WORKSPACE}/.ci/scripts/run_tests_ucc_nvls_slurm.sh
+      sudo -E -u svcnbu-swx-hpcx ${WORKSPACE}/.ci/scripts/run_nvls_slurm.sh '/opt/nvidia/src/ucc/.ci/scripts/run_tests_ucc_nvls_all.sh' ${NVLS_MPI_PPN:-4}
+    onfail: |
+      sudo -E -u svcnbu-swx-hpcx ${WORKSPACE}/.ci/scripts/stop_slurm_allocation.sh
+
+  - name: Run UCC NVLS MPI tests
+    containerSelector: "{name: 'build_helper'}"
+    timeout: "${TEST_TIMEOUT_MINUTES}"
+    run: |
+      set -x
+      export DOCKER_IMAGE_NAME="${registry_host}#torch-ucc/${UCC_URI_SUFFIX}:${DOCKER_IMAGE_TAG}"
+      export SLURM_JOB_ID=$(cat ${WORKSPACE}/job_id.txt)
+      sudo -E -u svcnbu-swx-hpcx ${WORKSPACE}/.ci/scripts/run_nvls_slurm.sh '/opt/nvidia/src/ucc/.ci/scripts/run_tests_ucc_nvls_mpi.sh' ${NVLS_MPI_PPN:-4}
     always: |
       sudo -E -u svcnbu-swx-hpcx ${WORKSPACE}/.ci/scripts/stop_slurm_allocation.sh
diff --git a/.ci/scripts/build_ucc.sh b/.ci/scripts/build_ucc.sh
@@ -5,6 +5,7 @@ export CXXFLAGS="-Wno-error=maybe-uninitialized"
 
 export UCC_ENABLE_GTEST=${UCC_ENABLE_GTEST:-yes}
 export UCC_ENABLE_NVLS=${UCC_ENABLE_NVLS:-no}
+export UCC_BUILD_TLS=${UCC_BUILD_TLS:-cuda,nccl,self,sharp,shm,ucp,mlx5}
 
 # In containers, calculate based on memory limits to avoid OOM
 # Determine number of parallel build jobs based on available system memory if running inside a container/Kubernetes
@@ -39,7 +40,7 @@ cd "${UCC_SRC_DIR}/build"
 # Build base configure flags
 CONFIGURE_FLAGS="--with-ucx=${UCX_INSTALL_DIR} --with-cuda=${CUDA_HOME} \
     --prefix=${UCC_INSTALL_DIR} --with-mpi \
-    --with-tls=cuda,nccl,self,sharp,shm,ucp,mlx5"
+    --with-tls=${UCC_BUILD_TLS}"
 
 # Add NVLS support if enabled
 if [ "${UCC_ENABLE_NVLS}" = "yes" ] || [ "${UCC_ENABLE_NVLS}" = "true" ] || [ "${UCC_ENABLE_NVLS}" = "1" ]; then
diff --git a/.ci/scripts/check_nvls_fabric.sh b/.ci/scripts/check_nvls_fabric.sh
@@ -0,0 +1,39 @@
+#!/bin/bash -xe
+
+echo "===== NVLS Fabric Smoke Test ($(hostname)) ====="
+
+echo "INFO: Checking GPU driver ..."
+nvidia-smi --query-gpu=index,name,uuid --format=csv,noheader
+NGPUS=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)
+if [ "$NGPUS" -eq 0 ]; then
+    echo "ERROR: No GPUs found"
+    exit 1
+fi
+echo "INFO: Found $NGPUS GPUs"
+
+echo "INFO: Checking NVLink fabric registration ..."
+FABRIC_OUTPUT=$(nvidia-smi -q | grep 'Fabric' -A 4)
+echo "$FABRIC_OUTPUT"
+
+COMPLETED_COUNT=$(echo "$FABRIC_OUTPUT" | grep -c 'State.*:.*Completed' || true)
+if [ "$COMPLETED_COUNT" -ne "$NGPUS" ]; then
+    echo "ERROR: Expected $NGPUS GPUs with Fabric State 'Completed', found $COMPLETED_COUNT"
+    exit 1
+fi
+
+FAILURES=$(echo "$FABRIC_OUTPUT" | grep 'Status' | grep -cv 'Success' || true)
+if [ "$FAILURES" -ne 0 ]; then
+    echo "ERROR: Some GPUs have Fabric Status != 'Success'"
+    exit 1
+fi
+echo "INFO: All $NGPUS GPUs registered to NVLink fabric successfully"
+
+echo "INFO: Checking NVLink link status ..."
+nvidia-smi nvlink --status
+echo "INFO: NVLink link status ... DONE"
+
+echo "INFO: Checking GPU P2P topology ..."
+nvidia-smi topo -p2p n
+echo "INFO: GPU P2P topology ... DONE"
+
+echo "===== NVLS Fabric Smoke Test PASSED ($(hostname)) ====="
diff --git a/.ci/scripts/run_nvls_slurm.sh b/.ci/scripts/run_nvls_slurm.sh
@@ -1,35 +1,33 @@
 #!/bin/bash -xe
 
+# Generic NVLS Slurm test runner.
+# Usage: run_nvls_slurm.sh <container_script> [ntasks_per_node]
 
 SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 source "${SCRIPT_DIR}/env.sh"
 
-readonly SLURM_COMMAND="srun --jobid=${SLURM_JOB_ID} --nodes=${SLURM_NODES} --mpi=pmi2 --ntasks-per-node=1 --container-image=${DOCKER_IMAGE_NAME} '/opt/nvidia/src/ucc/.ci/scripts/run_tests_ucc_nvls.sh'"
+CONTAINER_SCRIPT=${1:?"Usage: run_nvls_slurm.sh <container_script> [ntasks_per_node]"}
+NTASKS_PER_NODE=${2:-1}
+
+readonly SLURM_COMMAND="srun --jobid=${SLURM_JOB_ID} --nodes=${SLURM_NODES} --mpi=pmix --ntasks-per-node=${NTASKS_PER_NODE} --container-image=${DOCKER_IMAGE_NAME} '${CONTAINER_SCRIPT}'"
 
-# Validate SLURM_HEAD_NODE is set
 if [ -z "${SLURM_HEAD_NODE}" ]; then
     echo "ERROR: SLURM_HEAD_NODE is not set or empty"
     exit 1
 fi
 
-# Execute based on head node type
 case "${SLURM_HEAD_NODE}" in
     scctl)
         echo "Using scctl client to connect and execute slurm command"
         scctl client connect -- "${SLURM_COMMAND}"
         ;;
     dlcluster*)
         echo "Connecting to SLURM head node: ${SLURM_HEAD_NODE}"
-        # Escape the # character in the command for SSH transmission
         SLURM_COMMAND_ESCAPED="${SLURM_COMMAND//\#/\\#}"
         eval "${SSH_CMD} ${SLURM_HEAD_NODE} \"${SLURM_COMMAND_ESCAPED}\""
         ;;
-    "")
-        echo "ERROR: Invalid SLURM_HEAD_NODE value: ${SLURM_HEAD_NODE}"
-        exit 1
-        ;;
     *)
         echo "Connecting to SLURM head node: ${SLURM_HEAD_NODE}"
-        eval "${SSH_CMD} ${SLURM_HEAD_NODE} ${SLURM_COMMAND}"
+        eval "${SSH_CMD} ${SLURM_HEAD_NODE} \"${SLURM_COMMAND}\""
         ;;
 esac
diff --git a/.ci/scripts/run_tests_ucc_nvls_all.sh b/.ci/scripts/run_tests_ucc_nvls_all.sh
@@ -0,0 +1,25 @@
+#!/bin/bash -xe
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+source "${SCRIPT_DIR}/env.sh"
+
+if [ "${SLURM_LOCALID:-0}" = "0" ]; then
+    "${SCRIPT_DIR}/check_nvls_fabric.sh"
+fi
+
+export OMPI_MCA_coll=^hcoll
+export OMPI_MCA_coll_ucc_enable=0
+export UCC_LOG_LEVEL=info
+export UCC_TL_CUDA_NVLS_SM_COUNT=4
+export UCC_TLS=cuda,ucp
+
+PERFTEST=/opt/nvidia/bin/ucc/build/bin/ucc_perftest
+
+echo "INFO: NVLS perftest (allreduce) ..."
+UCC_TL_CUDA_TUNE=allreduce:cuda:@0 $PERFTEST -c allreduce -F -m cuda -b 1k -e 32M -d bfloat16 -o sum
+echo "INFO: NVLS perftest (allreduce) ... DONE"
+
+# Disabled: reduce_scatter NVLS is tested via MPI tests instead.
+#echo "INFO: NVLS perftest (reduce_scatter) ..."
+#UCC_TL_CUDA_TUNE=reduce_scatter:cuda:@3 $PERFTEST -c reduce_scatter -F -m cuda -b 1k -e 32M -d bfloat16 -o sum
+#echo "INFO: NVLS perftest (reduce_scatter) ... DONE"
diff --git a/.ci/scripts/run_tests_ucc_nvls_mpi.sh b/.ci/scripts/run_tests_ucc_nvls_mpi.sh
@@ -0,0 +1,25 @@
+#!/bin/bash -xe
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+source "${SCRIPT_DIR}/env.sh"
+
+if [ "${SLURM_LOCALID:-0}" = "0" ]; then
+    "${SCRIPT_DIR}/check_nvls_fabric.sh"
+fi
+
+export OMPI_MCA_coll=^hcoll
+export OMPI_MCA_coll_ucc_enable=0
+export UCC_TLS=cuda,ucp
+export UCC_LOG_LEVEL=info
+export UCC_TL_CUDA_NVLS_SM_COUNT=4
+
+EXE="/opt/nvidia/src/ucc/build/test/mpi/ucc_test_mpi"
+EXE+=" --set_device 2 --mtypes cuda"
+
+echo "INFO: NVLS MPI tests (allreduce) ..."
+UCC_TL_CUDA_TUNE="allreduce:cuda:@0" $EXE -c allreduce -d float32 -o sum -m 1024:33554432
+echo "INFO: NVLS MPI tests (allreduce) ... DONE"
+
+# echo "INFO: NVLS MPI tests (reduce_scatter) ..."
+# UCC_TL_CUDA_TUNE="reduce_scatter:cuda:@3" $EXE -c reduce_scatter -d float32 -o sum -m 1024:33554432
+# echo "INFO: NVLS MPI tests (reduce_scatter) ... DONE"
