diff --git a/.github/workflows/rocm-build-check.yml b/.github/workflows/rocm-build-check.yml new file mode 100644 index 0000000000..e79504218c --- /dev/null +++ b/.github/workflows/rocm-build-check.yml @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# TODO: replace this compile-only host CI with end-to-end transfer validation +# on a self-hosted MI300X/MI355X runner once available. + +name: rocm-build-check + +on: + pull_request: + paths: + - 'src/**' + - 'benchmark/nixlbench/**' + - 'contrib/Dockerfile' + - 'contrib/build-container.sh' + - 'contrib/build-wheel.sh' + - 'meson.build' + - 'meson_options.txt' + - '.github/workflows/rocm-build-check.yml' + workflow_dispatch: + +jobs: + build: + name: ROCm container build + smoke + runs-on: ubuntu-latest + # Non-blocking for now: a from-scratch ROCm build exceeds the hosted-runner + # budget (the "-complete" image plus building UCX/gRPC/etc. from source runs + # well past the timeout). Flip to blocking once the build consumes the + # pre-built ghcr.io/ai-dynamo/nixl-rocm base image so only nixl is compiled. + continue-on-error: true + if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }} + timeout-minutes: 90 + steps: + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 + with: + fetch-depth: 0 + persist-credentials: false + + # The ROCm "-complete" base image is ~30GB extracted and does not fit on a + # stock ubuntu-latest runner (~14GB free); reclaim space before building. + - name: Free up disk space + run: | + sudo rm -rf /usr/share/dotnet /opt/ghc \ + /usr/local/lib/android /opt/hostedtoolcache \ + /usr/local/share/boost /usr/local/.ghcup + sudo docker image prune -af || true + sudo docker builder prune -af || true + df -h / + + # docker-container driver is required for --output=type=cacheonly below. + - name: Set up buildx (docker-container driver) + run: docker buildx create --use --driver docker-container --name nixl-ci + + # Build the full ROCm graph without exporting an image. The wheel-presence + # and `import nixl` -> nixl_rocm smoke checks run as RUN steps inside the + # rocm-build stage, so --output=cacheonly validates them without the + # multi-GB image export that OOMs the hosted runner during compression. + - name: Build ROCm container (validate, no image export) + run: ./contrib/build-container.sh --rocm --check-only --no-cache diff --git a/.github/workflows/wheel-rocm.yml b/.github/workflows/wheel-rocm.yml new file mode 100644 index 0000000000..732da3ed9a --- /dev/null +++ b/.github/workflows/wheel-rocm.yml @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Builds the nixl_rocm wheel on tagged releases and uploads it as a release asset. + +name: wheel-rocm + +on: + push: + tags: + - 'v*' + workflow_dispatch: + inputs: + ref: + description: 'Git ref to build' + required: false + type: string + +concurrency: + group: wheel-rocm-${{ github.ref }} + cancel-in-progress: false + +jobs: + build-wheel: + name: Build nixl_rocm wheel + runs-on: ubuntu-latest + timeout-minutes: 90 + steps: + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 + with: + ref: ${{ inputs.ref || github.ref }} + fetch-depth: 0 + persist-credentials: false + + - name: Build ROCm container + run: ./contrib/build-container.sh --rocm --tag nixl-rocm:wheel --no-cache + + - name: Extract wheels + run: | + mkdir -p dist + id=$(docker create nixl-rocm:wheel) + docker cp "$id":/workspace/nixl/dist/. dist/ + docker rm "$id" + ls -la dist/ + + - name: Upload wheel artifacts + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4 + with: + name: nixl-rocm-wheels + path: dist/*.whl + if-no-files-found: error + + - name: Attach to release + if: startsWith(github.ref, 'refs/tags/v') + uses: softprops/action-gh-release@3bb12739c298aeb8a4eeaf626c5b8d85266b0e65 # v2 + with: + files: dist/*.whl diff --git a/contrib/Dockerfile b/contrib/Dockerfile index 94861e6bfd..43f3c7c940 100644 --- a/contrib/Dockerfile +++ b/contrib/Dockerfile @@ -15,26 +15,20 @@ ARG BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base" ARG BASE_IMAGE_TAG="25.10-cuda13.0-devel-ubuntu24.04" -ARG OS -FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} +FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS common-base -# Set default OS if not provided -ARG OS=${OS:-ubuntu24} +ARG OS ARG ARCH="x86_64" ARG DEFAULT_PYTHON_VERSION="3.12" -ARG UCX_REF="v1.21.x" -ARG BUILD_NIXL_EP="true" ARG RDMA_CORE_PREFIX="/usr" +ARG UCX_REF="940c1c1d9" ARG UCX_PREFIX="/usr" ARG UCX_PLUGIN_DIR="$UCX_PREFIX/lib/ucx" -ARG DOCA_PREFIX="/opt/mellanox/doca" ARG NIXL_PREFIX="/usr/local/nixl" ARG NIXL_PLUGIN_DIR="$NIXL_PREFIX/lib/$ARCH-linux-gnu/plugins" ARG NPROC ARG WHL_DEFAULT_PYTHON_VERSIONS="3.12" -ARG LIBFABRIC_VERSION="v1.21.0" -ARG LIBFABRIC_INSTALL_PATH="/usr/local" ARG BUILD_TYPE="release" ARG ABSL_TAG="lts_2025_08_14" ARG GRPC_TAG="v1.73.0" @@ -44,6 +38,9 @@ RUN apt-get update -y && \ apt-get install -y ubuntu-keyring && \ apt-get update -y && \ DEBIAN_FRONTEND=noninteractive apt-get -y install \ + ca-certificates \ + curl \ + git \ ninja-build \ libclang-dev \ cmake \ @@ -61,31 +58,22 @@ RUN apt-get update -y && \ libgtest-dev \ build-essential \ python3.12-dev \ + python3.12-venv \ clang \ hwloc \ libhwloc-dev \ pkg-config \ + protobuf-compiler \ + wget \ # aws-sdk-cpp dependencies libcurl4-openssl-dev libssl-dev uuid-dev zlib1g-dev \ # azure-sdk-for-cpp dependencies libxml2-dev -# Add DOCA repo + upgrade always so the RDMA reinstall below resolves MLNX/DOCA versions. -# Skip only the DOCA SDK install when the base image already ships it. -RUN ARCH_SUFFIX=$(if [ "${ARCH}" = "aarch64" ]; then echo "arm64"; else echo "amd64"; fi) && \ - MELLANOX_OS="$(. /etc/lsb-release; echo ${DISTRIB_ID}${DISTRIB_RELEASE} | tr A-Z a-z | tr -d .)" && \ - wget --tries=3 --waitretry=5 --no-verbose https://www.mellanox.com/downloads/DOCA/DOCA_v3.3.0/host/doca-host_3.3.0-088000-26.01-${MELLANOX_OS}_${ARCH_SUFFIX}.deb -O doca-host.deb && \ - dpkg -i doca-host.deb && \ - apt-get update && \ - apt-get upgrade -y && \ - if dpkg -s doca-sdk-gpunetio >/dev/null 2>&1; then \ - echo "DOCA SDK already installed in base image, skipping SDK install"; \ - else \ - apt-get install -y --no-install-recommends doca-sdk-gpunetio libdoca-sdk-gpunetio-dev libdoca-sdk-verbs-dev libdoca-sdk-telemetry-exporter-dev collectx-clxapidev; \ - fi - -# Reinstall RDMA build deps (libfabric --enable-efa needs libibverbs-dev), required even when -# DOCA is from the base image. --reinstall because apt skips equal versions. +# common-base is shared by the ROCm and CUDA build stages and stays DOCA-free +# (ROCm does not need DOCA; the CUDA stage installs DOCA/GPUNetIO itself). +# Reinstall RDMA packages so Infiniband headers and providers are consistent. +# Upgrade is not sufficient if the version is the same since apt skips the installation. RUN DEBIAN_FRONTEND=noninteractive apt-get -y install \ --reinstall libibverbs-dev rdma-core ibverbs-utils libibumad-dev \ libnuma-dev librdmacm-dev ibverbs-providers @@ -181,6 +169,192 @@ RUN wget --tries=3 --waitretry=5 \ RUN rm -rf /usr/lib/ucx RUN rm -rf /opt/hpcx/ucx +RUN cd /tmp && \ + git clone --depth 1 https://github.com/google/gtest-parallel.git && \ + mkdir -p /usr/local/bin && \ + cp gtest-parallel/gtest-parallel gtest-parallel/gtest_parallel.py /usr/local/bin/ +ENV PATH=/usr/local/bin:$PATH + +# By default, uv downloads python packages to $HOME/.cache/uv and hard links them +# from the virtual environment. This means that the files reside in /root/.cache/uv, +# which is not what we want since some systems mount user home dir into /root, +# in which case the venv is broken when the container is started. +# Set a custom cache directory inside /workspace to avoid this. +ENV UV_CACHE_DIR=/workspace/.cache/uv +RUN mkdir -p $UV_CACHE_DIR +# Disable build isolation, i.e. uv should not create a new virtual environment for +# building wheels. This is faster as it skips installing dependencies twice. +ENV UV_NO_BUILD_ISOLATION=1 +# Disable syncing, i.e. uv will not download packages outside uv pip commands. +ENV UV_NO_SYNC=1 +# Create a new virtual environment +ENV VIRTUAL_ENV=/workspace/.venv +RUN rm -rf $VIRTUAL_ENV && uv venv $VIRTUAL_ENV --python $DEFAULT_PYTHON_VERSION +# Activate the virtual environment +ENV PATH="$VIRTUAL_ENV/bin:$PATH" +# Install python dependencies +RUN uv pip install --upgrade meson meson-python pybind11 patchelf pyYAML click tabulate auditwheel tomlkit +# Upgrade setuptools to latest version for compatibility with PEP 639 (license format) +RUN uv pip install --upgrade 'setuptools>=80.9.0' + +ENV LD_LIBRARY_PATH=/usr/local/lib:$RDMA_CORE_PREFIX/lib:$LD_LIBRARY_PATH + +# Install pybind11 via apt +RUN apt-get update && apt-get install -y --no-install-recommends pybind11-dev + +# Set PKG_CONFIG_PATH for shared RDMA/UCX dependencies. +ENV PKG_CONFIG_PATH=/usr/local/lib/pkgconfig:/usr/local/lib/$ARCH-linux-gnu/pkgconfig +ENV PKG_CONFIG_PATH=$RDMA_CORE_PREFIX/lib/$ARCH-linux-gnu/pkgconfig:$UCX_PREFIX/lib/pkgconfig:$PKG_CONFIG_PATH + +ENV NIXL_PREFIX=$NIXL_PREFIX +ENV NIXL_PLUGIN_DIR=$NIXL_PLUGIN_DIR + +FROM common-base AS rocm-build + +# ARGs from common-base must be re-declared in each child stage to be accessible. +ARG ARCH="x86_64" +ARG DEFAULT_PYTHON_VERSION="3.12" +ARG UCX_REF="940c1c1d9" +ARG UCX_PREFIX="/usr" +ARG UCX_PLUGIN_DIR="$UCX_PREFIX/lib/ucx" +ARG NIXL_PREFIX="/usr/local/nixl" +ARG NIXL_PLUGIN_DIR="$NIXL_PREFIX/lib/$ARCH-linux-gnu/plugins" +ARG NPROC +ARG BUILD_TYPE="release" +ARG ROCM_PATH="/opt/rocm" +ARG WHEEL_NAME="nixl_rocm" +ARG BACKEND_DIR="nixl_rocm" +ARG PYTORCH_INDEX_URL="https://download.pytorch.org/whl/rocm7.2" + +ENV ROCM_PATH=$ROCM_PATH +ENV PATH=$ROCM_PATH/bin:$PATH +ENV LD_LIBRARY_PATH=$ROCM_PATH/lib:$ROCM_PATH/lib64:$LD_LIBRARY_PATH +ENV PKG_CONFIG_PATH=$ROCM_PATH/lib/pkgconfig:$ROCM_PATH/lib64/pkgconfig:$PKG_CONFIG_PATH + +# rocm/dev-ubuntu-24.04 ships gcc 14; install gcc 13 to match the cuda-build +# stage and avoid -Werror promotion issues in UCX. +RUN apt-get update -y && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + gcc-13 g++-13 && \ + update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-13 100 && \ + update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-13 100 && \ + update-alternatives --install /usr/bin/cc cc /usr/bin/gcc 100 && \ + rm -rf /var/lib/apt/lists/* + +RUN cd /usr/local/src && \ + git clone https://github.com/openucx/ucx.git && \ + cd ucx && \ + echo "=== Using UCX_REF=$UCX_REF ===" && \ + git checkout $UCX_REF && \ + ./autogen.sh && \ + ./contrib/configure-release-mt \ + --prefix=$UCX_PREFIX \ + --enable-shared \ + --disable-static \ + --disable-doxygen-doc \ + --enable-optimizations \ + --enable-cma \ + --enable-devel-headers \ + --with-rocm=$ROCM_PATH \ + --with-verbs \ + --with-dm \ + --with-efa && \ + make -j${NPROC:-$(nproc)} && \ + make -j${NPROC:-$(nproc)} install-strip && \ + ldconfig + +RUN uv pip install --index-url "$PYTORCH_INDEX_URL" 'torch==2.11.*' 'torchvision==0.26.*' 'torchaudio==2.11.*' + +# PyTorch's dep pin can downgrade setuptools; re-upgrade to keep the PEP 639 +# license parser used by nixl-meta's pyproject.toml. +RUN uv pip install --upgrade 'setuptools>=80.9.0' + +WORKDIR /workspace/nixl +COPY . /workspace/nixl + +RUN rm -rf build && \ + mkdir build && \ + meson setup \ + -Ducx_path=$UCX_PREFIX \ + -Ddisable_gds_backend=true \ + -Ddisable_plugins=LIBFABRIC,GDS,GDS_MT,GPUNETIO \ + -Dbuild_nixl_ep=false \ + -Dpython_backend_dir=$BACKEND_DIR \ + -Dpython_backend_package=$WHEEL_NAME \ + build/ --prefix=$NIXL_PREFIX --buildtype=$BUILD_TYPE && \ + cd build && \ + ninja && \ + ninja install + +RUN echo "$NIXL_PREFIX/lib/$ARCH-linux-gnu" > /etc/ld.so.conf.d/nixl.conf && \ + echo "$NIXL_PLUGIN_DIR" >> /etc/ld.so.conf.d/nixl.conf && \ + ldconfig + +RUN cd src/bindings/rust && cargo build --release --locked + +RUN export PATH=$VIRTUAL_ENV/bin:$PATH && \ + export UV_INDEX="$PYTORCH_INDEX_URL" && \ + mkdir -p dist && \ + ./contrib/build-wheel.sh \ + --python-version $DEFAULT_PYTHON_VERSION \ + --platform manylinux_2_39_$ARCH \ + --ucx-plugins-dir $UCX_PLUGIN_DIR \ + --nixl-plugins-dir $NIXL_PLUGIN_DIR \ + --output-dir /workspace/nixl/dist \ + --wheel-name $WHEEL_NAME \ + --backend-dir $BACKEND_DIR + +RUN cp build/src/bindings/python/nixl-meta/nixl-*.whl dist/ +RUN PYTHON_ABI_TAG="$(echo "$DEFAULT_PYTHON_VERSION" | tr -d .)" && \ + uv pip install dist/*cp${PYTHON_ABI_TAG}*.whl dist/nixl-*-none-any.whl + +# In-build smoke checks: confirm both wheels exist and that `import nixl` +# dispatches to the ROCm backend. Running these here lets CI validate the build +# with --output=cacheonly, avoiding the multi-GB image export that OOMs the +# hosted runner during layer compression. +RUN ls /workspace/nixl/dist/nixl_rocm-*.whl /workspace/nixl/dist/nixl-*-none-any.whl && \ + export PATH=$VIRTUAL_ENV/bin:$PATH && \ + python -c 'import nixl; assert nixl._pkg.__name__ == "nixl_rocm", nixl._pkg.__name__' + +FROM common-base AS cuda-build + +# ARGs from common-base must be re-declared in each child stage to be accessible. +ARG ARCH="x86_64" +ARG DEFAULT_PYTHON_VERSION="3.12" +ARG UCX_REF="940c1c1d9" +ARG UCX_PREFIX="/usr" +ARG UCX_PLUGIN_DIR="$UCX_PREFIX/lib/ucx" +ARG NIXL_PREFIX="/usr/local/nixl" +ARG NIXL_PLUGIN_DIR="$NIXL_PREFIX/lib/$ARCH-linux-gnu/plugins" +ARG NPROC +ARG BUILD_TYPE="release" +ARG BUILD_NIXL_EP="true" +ARG DOCA_PREFIX="/opt/mellanox/doca" +ARG LIBFABRIC_VERSION="v1.21.0" +ARG LIBFABRIC_INSTALL_PATH="/usr/local" + +ENV LD_LIBRARY_PATH=/usr/local/lib:$LIBFABRIC_INSTALL_PATH/lib:$LD_LIBRARY_PATH +ENV PKG_CONFIG_PATH=$DOCA_PREFIX/lib/$ARCH-linux-gnu/pkgconfig:$PKG_CONFIG_PATH + +# CUDA-only: DOCA/GPUNetIO packages (v3.3.0 + telemetry exporter, forward-ported +# from main #1728). Skip only the DOCA SDK install when the base image ships it. +RUN ARCH_SUFFIX=$(if [ "${ARCH}" = "aarch64" ]; then echo "arm64"; else echo "amd64"; fi) && \ + MELLANOX_OS="$(. /etc/lsb-release; echo ${DISTRIB_ID}${DISTRIB_RELEASE} | tr A-Z a-z | tr -d .)" && \ + wget --tries=3 --waitretry=5 --no-verbose https://www.mellanox.com/downloads/DOCA/DOCA_v3.3.0/host/doca-host_3.3.0-088000-26.01-${MELLANOX_OS}_${ARCH_SUFFIX}.deb -O doca-host.deb && \ + dpkg -i doca-host.deb && \ + apt-get update && \ + apt-get upgrade -y && \ + if dpkg -s doca-sdk-gpunetio >/dev/null 2>&1; then \ + echo "DOCA SDK already installed in base image, skipping SDK install"; \ + else \ + apt-get install -y --no-install-recommends doca-sdk-gpunetio libdoca-sdk-gpunetio-dev libdoca-sdk-verbs-dev libdoca-sdk-telemetry-exporter-dev collectx-clxapidev; \ + fi + +# Reinstall RDMA packages from the DOCA repository for the CUDA image. +RUN DEBIAN_FRONTEND=noninteractive apt-get -y install \ + --reinstall libibverbs-dev rdma-core ibverbs-utils libibumad-dev \ + libnuma-dev librdmacm-dev ibverbs-providers + RUN cd /usr/local/src && \ git clone https://github.com/openucx/ucx.git && \ cd ucx && \ @@ -206,13 +380,7 @@ RUN cd /usr/local/src && \ make -j${NPROC:-$(nproc)} install-strip && \ ldconfig -RUN cd /tmp && \ - git clone --depth 1 https://github.com/google/gtest-parallel.git && \ - mkdir -p /usr/local/bin && \ - cp gtest-parallel/gtest-parallel gtest-parallel/gtest_parallel.py /usr/local/bin/ -ENV PATH=/usr/local/bin:$PATH - -# Build libfabric from source +# CUDA-only: build libfabric with CUDA HMEM/GDRCopy support. RUN wget --tries=3 --waitretry=5 --timeout=30 --read-timeout=60 \ "https://github.com/ofiwg/libfabric/releases/download/${LIBFABRIC_VERSION}/libfabric-${LIBFABRIC_VERSION#v}.tar.bz2" -O libfabric.tar.bz2 && \ tar xjf libfabric.tar.bz2 && rm libfabric.tar.bz2 && \ @@ -261,16 +429,6 @@ RUN uv pip install --system --upgrade 'setuptools>=80.9.0' WORKDIR /workspace/nixl COPY . /workspace/nixl -ENV LD_LIBRARY_PATH=/usr/local/lib:$LIBFABRIC_INSTALL_PATH/lib:$RDMA_CORE_PREFIX/lib:$LD_LIBRARY_PATH - -# Install pybind11 via apt -RUN apt-get update && apt-get install -y --no-install-recommends pybind11-dev - -# Set PKG_CONFIG_PATH for NIXL EP dependencies (rdma-core, UCX, DOCA) -ENV PKG_CONFIG_PATH=/usr/local/lib/pkgconfig:/usr/local/lib/$ARCH-linux-gnu/pkgconfig -ENV PKG_CONFIG_PATH=$RDMA_CORE_PREFIX/lib/$ARCH-linux-gnu/pkgconfig:$UCX_PREFIX/lib/pkgconfig:$DOCA_PREFIX/lib/$ARCH-linux-gnu/pkgconfig:$PKG_CONFIG_PATH - -ENV NIXL_PREFIX=$NIXL_PREFIX RUN rm -rf build && \ mkdir build && \ if [ "$BUILD_NIXL_EP" = "true" ]; then \ @@ -288,8 +446,7 @@ RUN echo "$NIXL_PREFIX/lib/$ARCH-linux-gnu" > /etc/ld.so.conf.d/nixl.conf && \ echo "$NIXL_PLUGIN_DIR" >> /etc/ld.so.conf.d/nixl.conf && \ ldconfig -# Set environment variables for NIXL EP -ENV NIXL_PLUGIN_DIR=$NIXL_PLUGIN_DIR +# CUDA-only: NIXL EP examples are built only when BUILD_NIXL_EP=true. ENV PYTHONPATH=/workspace/nixl/build/examples/device/ep RUN cd src/bindings/rust && cargo build --release --locked diff --git a/contrib/README.rocm.md b/contrib/README.rocm.md new file mode 100644 index 0000000000..26416c7ca6 --- /dev/null +++ b/contrib/README.rocm.md @@ -0,0 +1,62 @@ +# Building NIXL for AMD ROCm + +Supported on AMD Instinct accelerators (MI300X, MI350X, MI355X). + +## Build + +```bash +./contrib/build-container.sh --rocm +``` + +Produces a `nixl-rocm:` image containing the NIXL libraries, +the meta-loader `nixl` wheel, and the backend `nixl_rocm` wheel under +`/workspace/nixl/dist/`. + +## Run + +```bash +docker run --rm -it \ + --device /dev/kfd --device /dev/dri --device /dev/infiniband \ + --network=host --ipc=host --cap-add=IPC_LOCK \ + nixl-rocm: +``` + +## Runtime knobs for AMD Pensando AINIC (ionic) NICs + +The ionic kernel driver requires UCX's dmabuf-based VRAM registration +path. Without it, UCX falls back to `ibv_reg_mr` on raw VRAM pointers, +which the driver rejects with `EINVAL`. To enable the dmabuf path: + +```bash +docker run --rm -it \ + --device /dev/kfd --device /dev/dri --device /dev/infiniband \ + --network=host --ipc=host --cap-add=IPC_LOCK \ + -v /boot:/boot:ro \ + -e UCX_ROCM_COPY_DMABUF=yes \ + -e UCX_ROCM_IPC_MIN_ZCOPY=0 \ + nixl-rocm: +``` + +- `/boot:/boot:ro` lets UCX's `uct_rocm_base_is_dmabuf_supported()` + read `/boot/config-$(uname -r)` to verify `CONFIG_PCI_P2PDMA=y` and + `CONFIG_DMABUF_MOVE_NOTIFY=y`. +- `UCX_ROCM_COPY_DMABUF=yes` opts into the dmabuf path in the + `rocm_copy` memory domain (default is `no` upstream). +- `UCX_ROCM_IPC_MIN_ZCOPY=0` engages the rocm_ipc zero-copy path. + +These knobs are not needed for non-ionic RDMA NICs (e.g. Mellanox, +Broadcom Thor). + +## nixlbench + +The `nixlbench` benchmark binary is not built by the container by +default. To build it inside the running container (`$NIXL_PREFIX` and +`$ROCM_PATH` are set by the container; outside Docker, export them or +substitute `/usr/local/nixl` and `/opt/rocm` directly): + +```bash +cd benchmark/nixlbench +meson setup build -Duse_rocm=true -Drocm_path=$ROCM_PATH -Dnixl_path=$NIXL_PREFIX +ninja -C build +./build/nixlbench --help +``` diff --git a/contrib/build-container.sh b/contrib/build-container.sh index f99400e0a4..a04702ca80 100755 --- a/contrib/build-container.sh +++ b/contrib/build-container.sh @@ -35,12 +35,16 @@ ARCH=$(uname -m) WHL_BASE=manylinux_2_39 WHL_PLATFORM=${WHL_BASE}_${ARCH} WHL_PYTHON_VERSIONS="3.12" -UCX_REF=${UCX_REF:-v1.21.x} +UCX_REF=${UCX_REF:-940c1c1d9} BUILD_NIXL_EP="true" OS="ubuntu24" NPROC=${NPROC:-$(nproc)} GRPC_NPROC=${GRPC_NPROC:-$(nproc)} BUILD_TYPE="release" +DOCKER_TARGET="" +WHEEL_NAME="" +BACKEND_DIR="" +PYTORCH_INDEX_URL="" get_options() { while :; do @@ -84,6 +88,9 @@ get_options() { --no-cache) NO_CACHE=" --no-cache" ;; + --check-only) + CHECK_ONLY=1 + ;; --build-type) if [ "$2" ]; then BUILD_TYPE=$2 @@ -127,6 +134,17 @@ get_options() { --build-nixl-ep) BUILD_NIXL_EP=true ;; + --rocm) + DOCKER_FILE="${SOURCE_DIR}/Dockerfile" + BASE_IMAGE=rocm/dev-ubuntu-24.04 + BASE_IMAGE_TAG=7.2.4-complete + ARCH=x86_64 + DOCKER_TARGET="--target rocm-build" + WHEEL_NAME=nixl_rocm + BACKEND_DIR=nixl_rocm + PYTORCH_INDEX_URL=https://download.pytorch.org/whl/rocm7.2 + BUILD_NIXL_EP=false + ;; --arch) if [ "$2" ]; then ARCH=$2 @@ -160,7 +178,9 @@ get_options() { WHL_PLATFORM=${WHL_BASE}_${ARCH} - if [ -z "$TAG" ]; then + if [ -z "$TAG" ] && [ "$WHEEL_NAME" = "nixl_rocm" ]; then + TAG="--tag nixl-rocm:${VERSION}" + elif [ -z "$TAG" ]; then TAG="--tag nixl:${VERSION}" fi } @@ -180,6 +200,18 @@ show_build_options() { else echo "NIXL EP: Disabled" fi + if [ -n "$WHEEL_NAME" ]; then + echo "Wheel Name: ${WHEEL_NAME}" + fi + if [ -n "$BACKEND_DIR" ]; then + echo "Backend Dir: ${BACKEND_DIR}" + fi + if [ -n "$PYTORCH_INDEX_URL" ]; then + echo "PyTorch Index URL: ${PYTORCH_INDEX_URL}" + fi + if [ -n "$DOCKER_TARGET" ]; then + echo "Docker Target: ${DOCKER_TARGET#--target }" + fi echo "Build Type: ${BUILD_TYPE}" } @@ -189,12 +221,14 @@ show_help() { echo " [--base-image-tag base image tag]" echo " [--wheel-base base platform for wheel builds]" echo " [--no-cache disable docker build cache]" + echo " [--check-only build via buildx without exporting an image (CI validation)]" echo " [--os [ubuntu24|ubuntu22] to select Ubuntu version]" echo " [--build-type [debug|release] to select build type (default: release)]" echo " [--tag tag for image]" echo " [--python-versions python versions to build for, comma separated]" echo " [--ucx-ref ucx git reference (branch, tag, or sha)]" echo " [--build-nixl-ep build NIXL with NIXL EP support (requires UCX >= 1.21)]" + echo " [--rocm build the ROCm image and nixl_rocm wheel]" echo " [--arch [x86_64|aarch64] to select target architecture]" echo " [--dockerfile path to a dockerfile to use]" exit 0 @@ -226,7 +260,22 @@ BUILD_ARGS+=" --build-arg NPROC=$NPROC" BUILD_ARGS+=" --build-arg GRPC_NPROC=$GRPC_NPROC" BUILD_ARGS+=" --build-arg OS=$OS" BUILD_ARGS+=" --build-arg BUILD_TYPE=$BUILD_TYPE" +if [ -n "$WHEEL_NAME" ]; then + BUILD_ARGS+=" --build-arg WHEEL_NAME=$WHEEL_NAME" +fi +if [ -n "$BACKEND_DIR" ]; then + BUILD_ARGS+=" --build-arg BACKEND_DIR=$BACKEND_DIR" +fi +if [ -n "$PYTORCH_INDEX_URL" ]; then + BUILD_ARGS+=" --build-arg PYTORCH_INDEX_URL=$PYTORCH_INDEX_URL" +fi show_build_options -docker build --platform linux/$ARCH -f $DOCKER_FILE $BUILD_ARGS $TAG $NO_CACHE $BUILD_CONTEXT +if [ -n "$CHECK_ONLY" ]; then + # Build the full graph (incl. in-Dockerfile smoke RUNs) but do not export an + # image. Skipping the multi-GB image export avoids OOM-ing hosted CI runners. + docker buildx build --platform linux/$ARCH -f $DOCKER_FILE $DOCKER_TARGET $BUILD_ARGS --output type=cacheonly $NO_CACHE $BUILD_CONTEXT +else + docker build --platform linux/$ARCH -f $DOCKER_FILE $DOCKER_TARGET $BUILD_ARGS $TAG $NO_CACHE $BUILD_CONTEXT +fi diff --git a/contrib/build-wheel.sh b/contrib/build-wheel.sh index d765bf9ac3..ffc6bb8bda 100755 --- a/contrib/build-wheel.sh +++ b/contrib/build-wheel.sh @@ -23,6 +23,8 @@ NIXL_PLUGINS_DIR="/usr/local/nixl/lib/$ARCH-linux-gnu/plugins" OUTPUT_DIR="dist" BUILD_NIXL_EP="false" TORCH_VERSIONS="" +WHEEL_NAME="" +BACKEND_DIR="" while [[ $# -gt 0 ]]; do case $1 in @@ -51,13 +53,25 @@ while [[ $# -gt 0 ]]; do shift shift ;; + --wheel-name) + WHEEL_NAME=$2 + shift + shift + ;; + --backend-dir) + BACKEND_DIR=$2 + shift + shift + ;; --help) - echo "Usage: $0 [--python-version ] [--platform ] [--output-dir ] [--ucx-plugins-dir ] [--nixl-plugins-dir ]" + echo "Usage: $0 [--python-version ] [--platform ] [--output-dir ] [--ucx-plugins-dir ] [--nixl-plugins-dir ] [--wheel-name ] [--backend-dir ]" echo " --python-version: Python version to build the wheel for (default: $PYTHON_VERSION)" echo " --platform: Platform to build the wheel for (default: $WHL_PLATFORM)" echo " --output-dir: Directory to output the wheel to (default: $OUTPUT_DIR)" echo " --ucx-plugins-dir: Directory to find UCX plugins in (default: $UCX_PLUGINS_DIR)" echo " --nixl-plugins-dir: Directory to find NIXL plugins in (default: $NIXL_PLUGINS_DIR)" + echo " --wheel-name: Python wheel distribution name (default: auto-detected CUDA wheel)" + echo " --backend-dir: Python import package directory (default: wheel name with '-' replaced by '_')" echo " --build-nixl-ep: Build wheel with nixl_ep package included (requires a CUDA sm_90 or newer target environment)" echo " --torch-versions: Comma-separated list of torch versions to build the wheel for (default: $TORCH_VERSIONS)" echo " --help: Show this help message" @@ -91,12 +105,56 @@ set -x TMP_DIR=$(mktemp -d) -CUDA_MAJOR=$(nvcc --version | grep -Eo 'release [0-9]+\.[0-9]+' | cut -d' ' -f2 | cut -d'.' -f1) -# Must be 12 or 13 -if [ "$CUDA_MAJOR" -ne 12 ] && [ "$CUDA_MAJOR" -ne 13 ]; then - echo "Invalid CUDA_MAJOR: '$CUDA_MAJOR'" - exit 1 +if [ -z "$WHEEL_NAME" ]; then + CUDA_MAJOR=$(nvcc --version | grep -Eo 'release [0-9]+\.[0-9]+' | cut -d' ' -f2 | cut -d'.' -f1) + # Must be 12 or 13 + if [ "$CUDA_MAJOR" -ne 12 ] && [ "$CUDA_MAJOR" -ne 13 ]; then + echo "Invalid CUDA_MAJOR: '$CUDA_MAJOR'" + exit 1 + fi + WHEEL_NAME="nixl-cu${CUDA_MAJOR}" + BACKEND_DIR="nixl_cu${CUDA_MAJOR}" +elif [ -z "$BACKEND_DIR" ]; then + BACKEND_DIR="${WHEEL_NAME//-/_}" +fi + +# ROCm (or any explicit --wheel-name) build: nixl and torch are already built +# and installed in the build stage, so do one self-contained wheel build and +# skip the CUDA-only nvcc / multi-torch machinery below. +if [ -z "$CUDA_MAJOR" ]; then + ./contrib/tomlutil.py --wheel-name "$WHEEL_NAME" pyproject.toml + + UV_BUILD_ARGS=( + --wheel + --out-dir "$TMP_DIR" + --python "$PYTHON_VERSION" + -Csetup-args="-Dpython_backend_dir=$BACKEND_DIR" + -Csetup-args="-Dpython_backend_package=$WHEEL_NAME" + ) + if [ "$BUILD_NIXL_EP" = "true" ]; then + UV_BUILD_ARGS+=( + -Csetup-args=-Dbuild_nixl_ep=true + -Csetup-args=-Dbuild_examples=true + ) + fi + uv build "${UV_BUILD_ARGS[@]}" + + mkdir "$TMP_DIR/dist" + auditwheel repair \ + --exclude 'libcuda*' --exclude 'libcufile*' \ + --exclude 'libssl*' --exclude 'libcrypto*' \ + --exclude 'libefa*' --exclude 'libhwloc*' --exclude 'libfabric*' \ + --exclude 'libtorch*' --exclude 'libc10*' --exclude 'libdoca*' \ + --exclude 'libred_client*' --exclude 'libred_async*' --exclude 'liblz4*' \ + --exclude 'libamdhip64*' --exclude 'libhsa-runtime64*' \ + --exclude 'libroc*' --exclude 'librocm*' --exclude 'libhip*' --exclude 'libamd_comgr*' \ + "$TMP_DIR"/nixl*.whl --plat "$WHL_PLATFORM" --wheel-dir "$TMP_DIR/dist" + ./contrib/wheel_add_ucx_plugins.py --ucx-plugins-dir "$UCX_PLUGINS_DIR" --nixl-plugins-dir "$NIXL_PLUGINS_DIR" "$TMP_DIR"/dist/*.whl + cp "$TMP_DIR"/dist/*.whl "$OUTPUT_DIR" + rm -rf "$TMP_DIR" + exit 0 fi + AUDITWHEEL_EXCLUDES="--exclude libcuda* --exclude libcufile* --exclude libssl* --exclude libcrypto* --exclude libefa* --exclude libhwloc* --exclude libfabric* --exclude libtorch* --exclude libc10* --exclude libdoca* --exclude libred_client* --exclude libred_async* --exclude liblz4*" PKG_NAME="nixl-cu${CUDA_MAJOR}" diff --git a/meson.build b/meson.build index 0c23b3537f..454f7e911e 100644 --- a/meson.build +++ b/meson.build @@ -309,14 +309,21 @@ else cuda_wheel_dir = 'nixl_cu12' endif -# Allow callers (e.g. ROCm CI/Dockerfile) to override the wheel variant suffix -# without otherwise affecting the build. Non-empty value wins over the -# autodetected CUDA-major value above: e.g. -Dwheel_variant=rocm -> nixl_rocm. wheel_variant = get_option('wheel_variant').strip() if wheel_variant != '' cuda_wheel_dir = 'nixl_' + wheel_variant endif +python_backend_dir_opt = get_option('python_backend_dir') +if python_backend_dir_opt != '' + cuda_wheel_dir = python_backend_dir_opt +endif + +python_backend_package = get_option('python_backend_package') +if python_backend_package == '' + python_backend_package = cuda_wheel_dir.replace('_', '-') +endif + # Check for etcd-cpp-api - use multiple methods for discovery etcd_dep = dependency('etcd-cpp-api', required : false) etcd_inc_path = get_option('etcd_inc_path') diff --git a/meson_options.txt b/meson_options.txt index eb875fc36c..8f2b0d2aa6 100644 --- a/meson_options.txt +++ b/meson_options.txt @@ -33,6 +33,8 @@ option('enable_plugins', type: 'string', value: '', description: 'Comma-separate option('disable_plugins', type: 'string', value: '', description: 'Comma-separated list of plugins to exclude from build. Cannot be used with enable_plugins.') option('build_docs', type: 'boolean', value: false, description: 'Build Doxygen documentation') option('release_wheel', type: 'boolean', value: false, description: 'Add nixl-cu12 and nixl-cu13 dependencies to the meta wheel') +option('python_backend_dir', type: 'string', value: '', description: 'Override Python backend import package directory') +option('python_backend_package', type: 'string', value: '', description: 'Override Python backend distribution package name') option('log_level', type: 'combo', choices: ['trace', 'debug', 'info', 'warning', 'error', 'fatal', 'auto'], value: 'auto', description: 'Log Level (auto: auto-detect based on build type: trace for debug builds, info for release builds)') option('rust', type: 'boolean', value: false, description: 'Build Rust bindings') option('sanitizer', type: 'combo', diff --git a/src/bindings/python/nixl-meta/meson.build b/src/bindings/python/nixl-meta/meson.build index 579bbdf1a9..4821ee3415 100644 --- a/src/bindings/python/nixl-meta/meson.build +++ b/src/bindings/python/nixl-meta/meson.build @@ -24,7 +24,7 @@ pyproject_toml = configure_file( 'VERSION': meson.project_version(), 'WHEEL_DEPS': get_option('release_wheel') \ ? '"nixl-cu12==' + meson.project_version() + '", "nixl-cu13==' + meson.project_version() + '"' \ - : '"' + cuda_wheel_dir.replace('_', '-') + '==' + meson.project_version() + '"', + : '"' + python_backend_package + '==' + meson.project_version() + '"', } ) diff --git a/src/bindings/python/nixl-meta/nixl/__init__.py b/src/bindings/python/nixl-meta/nixl/__init__.py index 283a472448..496b8d1aee 100644 --- a/src/bindings/python/nixl-meta/nixl/__init__.py +++ b/src/bindings/python/nixl-meta/nixl/__init__.py @@ -25,7 +25,25 @@ def _get_torch_cuda_major() -> int | None: return int(_torch_cuda_ver.split(".")[0]) if _torch_cuda_ver else None -def _load_cuda_backend() -> str: +def _get_torch_rocm_version() -> str | None: + """Return the ROCm/HIP version that torch was built for, or None.""" + from torch import version as _torch_version + + return getattr(_torch_version, "hip", None) + + +def _load_backend() -> str: + rocm_version = _get_torch_rocm_version() + if rocm_version is not None: + try: + return importlib.import_module("nixl_rocm").__name__ + except ModuleNotFoundError as e: + if e.name != "nixl_rocm": + raise + raise ImportError( + f"torch reports ROCm {rocm_version} but nixl_rocm is not installed" + ) from e + cuda_major = _get_torch_cuda_major() if cuda_major is not None: pip_name = f"nixl-cu{cuda_major}" @@ -39,7 +57,7 @@ def _load_cuda_backend() -> str: f"torch reports CUDA {cuda_major} but {pip_name} is not installed" ) from e # CPU-only torch — use whatever backend is installed - for mod_name in ("nixl_cu13", "nixl_cu12"): + for mod_name in ("nixl_rocm", "nixl_cu13", "nixl_cu12"): try: return importlib.import_module(mod_name).__name__ except ModuleNotFoundError as e: @@ -47,10 +65,10 @@ def _load_cuda_backend() -> str: # Re-raise if the error is not about the module we're trying to import raise continue - raise ImportError("No NIXL CUDA backend found") + raise ImportError("No NIXL backend found") -_pkg = sys.modules[_load_cuda_backend()] +_pkg = sys.modules[_load_backend()] submodules = ["_api", "_bindings", "_utils", "logging"] for sub_name in submodules: