Bindings: package separate wheels for CUDA 12 and 13 (#915)

* Package versioned wheel

Signed-off-by: Ovidiu Mara <[email protected]>

* Fix build for multiple python versions

Signed-off-by: Ovidiu Mara <[email protected]>

* Reformatting

Signed-off-by: Ovidiu Mara <[email protected]>

* Remove unused import

Signed-off-by: Ovidiu Mara <[email protected]>

* Add missing package

Signed-off-by: Ovidiu Mara <[email protected]>

* Install virtual environment the same way in all Dockerfiles for script compatibility

Signed-off-by: Ovidiu Mara <[email protected]>

* NIXL meta-package

Signed-off-by: Ovidiu Mara <[email protected]>

* Fix build in CI

Signed-off-by: Ovidiu Mara <[email protected]>

* Fix build

Signed-off-by: Ovidiu Mara <[email protected]>

* Revert changes to build.sh

Signed-off-by: Ovidiu Mara <[email protected]>

* Fix build issue

Signed-off-by: Ovidiu Mara <[email protected]>

* Simplify wheel build script

Signed-off-by: Ovidiu Mara <[email protected]>

* Fix build error due to wrong filename

Signed-off-by: Ovidiu Mara <[email protected]>

* Copy all wheels to dist/

Signed-off-by: Ovidiu Mara <[email protected]>

* Replace hardcoded version with template

Signed-off-by: Ovidiu Mara <[email protected]>

* Add missing file

Signed-off-by: Ovidiu Mara <[email protected]>

* Fix build

Signed-off-by: Ovidiu Mara <[email protected]>

---------

Signed-off-by: Ovidiu Mara <[email protected]>

Author: ovidiusm

benchmark/nixlbench/contrib/Dockerfile

@@ -183,7 +183,7 @@ 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
+RUN uv pip install --upgrade meson meson-python pybind11 patchelf pyYAML click tabulate auditwheel tomlkit
 # Install PyTorch
 RUN CUDA_SHORT_VERSION=cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d .) && \
     FLAGS="--index-url https://download.pytorch.org/whl/$CUDA_SHORT_VERSION" && \

contrib/Dockerfile

@@ -167,23 +167,41 @@ RUN wget --tries=3 --waitretry=5 --timeout=30 --read-timeout=60 \
     make install && \
     ldconfig
 
+# 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
+# Install PyTorch
+RUN export UV_INDEX="https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d .)" && \
+    uv pip install torch torchvision torchaudio
+
 WORKDIR /workspace/nixl
 COPY . /workspace/nixl
 
 ENV LD_LIBRARY_PATH=/usr/local/lib:$LIBFABRIC_INSTALL_PATH/lib:$LD_LIBRARY_PATH
 
-ENV VIRTUAL_ENV=/workspace/nixl/.venv
-RUN rm -rf $VIRTUAL_ENV && uv venv $VIRTUAL_ENV --python $DEFAULT_PYTHON_VERSION && \
-    # pybind11 pip install needed for ubuntu 22.04
-    uv pip install --upgrade "meson>=0.64.0" pybind11 patchelf
-
 # Install pybind11 via apt
 RUN apt-get update && apt-get install -y --no-install-recommends pybind11-dev
 
 ENV NIXL_PREFIX=$NIXL_PREFIX
 RUN rm -rf build && \
     mkdir build && \
-    uv run meson setup -Dlibfabric_path=$LIBFABRIC_INSTALL_PATH build/ --prefix=$NIXL_PREFIX && \
+    meson setup -Dlibfabric_path=$LIBFABRIC_INSTALL_PATH build/ --prefix=$NIXL_PREFIX && \
     cd build && \
     ninja && \
     ninja install
@@ -195,11 +213,13 @@ RUN echo "$NIXL_PREFIX/lib/$ARCH-linux-gnu" > /etc/ld.so.conf.d/nixl.conf && \
 RUN cd src/bindings/rust && cargo build --release --locked
 
 # Build wheel using the build-wheel.sh script for better UCX plugin bundling and library management
-RUN ./contrib/build-wheel.sh \
+RUN export PATH=$VIRTUAL_ENV/bin:$PATH && \
+    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
 
-RUN uv pip install dist/nixl-*cp${DEFAULT_PYTHON_VERSION//./}*.whl
+RUN uv pip install dist/nixl*cp${DEFAULT_PYTHON_VERSION//./}*.whl

contrib/Dockerfile.manylinux

@@ -166,10 +166,7 @@ ENV LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
 
 ENV CUDA_PATH=/usr/local/cuda
 
-WORKDIR /workspace/nixl
-ENV VIRTUAL_ENV=/workspace/nixl/.venv
-RUN uv venv $VIRTUAL_ENV --python $DEFAULT_PYTHON_VERSION && \
-    uv pip install --upgrade meson pybind11 patchelf
+WORKDIR /workspace
 
 RUN rm -rf /usr/lib/ucx
 RUN rm -rf /opt/hpcx/ucx
@@ -228,11 +225,30 @@ RUN wget --tries=3 --waitretry=5 --timeout=30 --read-timeout=60 \
     make install && \
     ldconfig
 
+# 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
+# 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 setuptools tabulate auditwheel tomlkit
+# Install PyTorch
+RUN export UV_INDEX="https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d .)" && \
+    uv pip install torch torchvision torchaudio
+
 COPY . /workspace/nixl
+WORKDIR /workspace/nixl
 
 RUN rm -rf build && \
     mkdir build && \
-    uv run meson setup build/ --prefix=/usr/local/nixl --buildtype=release \
+    meson setup build/ --prefix=/usr/local/nixl --buildtype=release \
     -Dcudapath_lib="/usr/local/cuda/lib64" \
     -Dcudapath_inc="/usr/local/cuda/include" && \
     cd build && \
@@ -252,7 +268,10 @@ RUN echo "/usr/local/nixl/lib/$ARCH-linux-gnu" > /etc/ld.so.conf.d/nixl.conf &&
 ARG WHL_PYTHON_VERSIONS="3.9,3.10,3.11,3.12"
 ARG WHL_PLATFORM="manylinux_2_28_$ARCH"
 RUN IFS=',' read -ra PYTHON_VERSIONS <<< "$WHL_PYTHON_VERSIONS" && \
+    export UV_INDEX="https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d .)" && \
+    mkdir -p dist && \
     for PYTHON_VERSION in "${PYTHON_VERSIONS[@]}"; do \
+        export PATH=$VIRTUAL_ENV/bin:$PATH && \
         ./contrib/build-wheel.sh \
             --python-version $PYTHON_VERSION \
             --platform $WHL_PLATFORM \
@@ -261,4 +280,7 @@ RUN IFS=',' read -ra PYTHON_VERSIONS <<< "$WHL_PYTHON_VERSIONS" && \
             --output-dir dist ; \
     done
 
-RUN uv pip install dist/nixl-*cp${DEFAULT_PYTHON_VERSION//./}*.whl
+# Copy the meta package wheel to the output directory
+RUN cp build/src/bindings/python/nixl-meta/nixl*.whl dist/
+
+RUN uv pip install dist/nixl*cp${DEFAULT_PYTHON_VERSION//./}*.whl

contrib/build-wheel.sh

@@ -73,23 +73,24 @@ done
 set -e
 set -x
 
-# Check for required dependencies
-if ! command -v uv &> /dev/null; then
-    echo "Required dependency: uv is not installed. Please install it from https://astral.sh/uv/install.sh"
-    exit 1
-fi
-
 # Build the wheel
 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
+fi
+PKG_NAME="nixl-cu${CUDA_MAJOR}"
+./contrib/tomlutil.py --set-name $PKG_NAME pyproject.toml
 uv build --wheel --out-dir $TMP_DIR --python $PYTHON_VERSION
 
 # Bundle libraries
-uv pip install auditwheel patchelf
-
-uv run auditwheel repair --exclude 'libcuda*' --exclude 'libcufile*' --exclude 'libssl*' --exclude 'libcrypto*' --exclude 'libefa*' --exclude 'libhwloc*' --exclude 'libfabric*' $TMP_DIR/nixl-*.whl --plat $WHL_PLATFORM --wheel-dir $OUTPUT_DIR
-
-uv run ./contrib/wheel_add_ucx_plugins.py --ucx-plugins-dir $UCX_PLUGINS_DIR --nixl-plugins-dir $NIXL_PLUGINS_DIR $OUTPUT_DIR/*.whl
+mkdir $TMP_DIR/dist
+auditwheel repair --exclude 'libcuda*' --exclude 'libcufile*' --exclude 'libssl*' --exclude 'libcrypto*' --exclude 'libefa*' --exclude 'libhwloc*' --exclude 'libfabric*' $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
 
 # Clean up
 rm -rf "$TMP_DIR"
-

contrib/tomlutil.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python3
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+
+import tomlkit
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--set-name", type=str, help="Set the project name")
+parser.add_argument("file", type=str, help="The toml file to modify")
+args = parser.parse_args()
+
+with open(args.file) as f:
+    doc = tomlkit.parse(f.read())
+
+if args.set_name:
+    doc["project"]["name"] = args.set_name
+
+with open(args.file, "w") as f:
+    f.write(tomlkit.dumps(doc))

contrib/wheel_add_ucx_plugins.py

@@ -189,9 +189,12 @@ def add_plugins(wheel_path, sys_plugins_dir, install_dirname):
     """
     temp_dir = extract_wheel(wheel_path)
 
-    pkg_libs_dir = os.path.join(temp_dir, "nixl.libs")
+    pkg_name = wheel_path.split("/")[-1].split("-")[0]
+    pkg_libs_dir = os.path.join(temp_dir, f"{pkg_name}.libs")
     if not os.path.exists(pkg_libs_dir):
-        raise FileNotFoundError(f"nixl.libs directory not found in wheel: {wheel_path}")
+        raise FileNotFoundError(
+            f"{pkg_name}.libs directory not found in wheel: {wheel_path}"
+        )
 
     logger.debug("Listing existing libs:")
     name_map = get_repaired_lib_name_map(pkg_libs_dir)

meson.build

@@ -101,10 +101,10 @@ if cuda_dep.found()
     nvcc_flags_link += ['-gencode=arch=compute_90,code=sm_90']
     add_project_link_arguments(nvcc_flags_link, language: 'cuda')
     message('nvcc version: ' + nvcc.version())
-    if nvcc.version().version_compare('>=12.8')
+    if nvcc.version().version_compare('>=12.8') and nvcc.version().version_compare('<13.0')
         doca_gpunetio_dep = dependency('doca-gpunetio', required : false)
     else
-        warning('CUDA version is less than 12.8, GPUNETIO plugin will be disabled')
+        warning('GPUNETIO plugin not supported in CUDA version: ' + nvcc.version())
         doca_gpunetio_dep = disabler()
     endif
 else
@@ -116,7 +116,7 @@ endif
 if cuda_dep.found()
   nvcc_cmd = find_program('nvcc', required: false)
   if nvcc_cmd.found()
-    if nvcc_cmd.version().version_compare('>=12.8')
+    if nvcc_cmd.version().version_compare('>=12.8') and nvcc_cmd.version().version_compare('<13.0')
       doca_gpunetio_dep = dependency('doca-gpunetio', required : false)
     else
       warning('CUDA version = ' + nvcc_cmd.version() + ', GPUNETIO plugin will be disabled')

src/bindings/python/meson.build

@@ -31,3 +31,5 @@ py.extension_module('_utils',
            subdir: ('nixl'),
            dependencies: [pybind_dep],
            install: true)
+
+subdir('nixl-meta')

src/bindings/python/nixl-meta/README.md

@@ -0,0 +1,9 @@
+# nixl
+
+This is a *meta package* that declares optional dependencies on CUDA variants.
+
+Install one of:
+```bash
+pip install "nixl[cu12]"   # for CUDA 12
+pip install "nixl[cu13]"   # for CUDA 13
+```

src/bindings/python/nixl-meta/meson.build

@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+source_dir = meson.current_source_dir()
+build_dir = meson.current_build_dir()
+
+configure_file(
+  input: 'pyproject.toml.in',
+  output: 'pyproject.toml',
+  configuration: { 'VERSION': meson.project_version() }
+)
+
+uv = find_program('uv', required: false)
+if uv.found()
+  wheel_name = 'nixl-@[email protected]'.format(meson.project_version())
+  meta_wheel = custom_target(
+    'build_nixl_meta',
+    output: [wheel_name],
+    command: [uv, 'build', '--wheel', '--out-dir', build_dir, build_dir],
+    install: false,
+    build_by_default: true
+  )
+else
+  warning('uv not found, skipping meta package build')
+endif