NIXL EP wheels: Add support for multiple CUDA and PyTorch versions

contrib/Dockerfile.manylinux

@@ -354,11 +354,14 @@ RUN echo "/usr/local/nixl/lib/$ARCH-linux-gnu" > /etc/ld.so.conf.d/nixl.conf &&
 # Create the wheel
 # No need to specifically add path to libcuda.so here, meson finds the stubs and links them
 ARG WHL_PYTHON_VERSIONS="3.10,3.11,3.12,3.13,3.14"
+ARG WHL_TORCH_VERSIONS="2.11,2.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 .)" && \
-    export UV_INDEX_STRATEGY=unsafe-best-match && \
-    if [ "$BUILD_NIXL_EP" = "true" ]; then EP_BUILD_FLAG="--build-nixl-ep"; else EP_BUILD_FLAG=""; fi && \
+    if [ "$BUILD_NIXL_EP" = "true" ]; then \
+        EP_BUILD_FLAGS="--build-nixl-ep --torch-versions $WHL_TORCH_VERSIONS"; \
+    else \
+        EP_BUILD_FLAGS=""; \
+    fi && \
     rm -rf dist && mkdir -p dist && \
     for PYTHON_VERSION in "${PYTHON_VERSIONS[@]}"; do \
         export PATH=$VIRTUAL_ENV/bin:$PATH && \
@@ -368,7 +371,7 @@ RUN IFS=',' read -ra PYTHON_VERSIONS <<< "$WHL_PYTHON_VERSIONS" && \
             --ucx-plugins-dir /usr/lib64/ucx \
             --nixl-plugins-dir $NIXL_PLUGIN_DIR \
             --output-dir dist \
-            $EP_BUILD_FLAG ; \
+            $EP_BUILD_FLAGS ; \
     done
 
 # Copy the meta package wheel to the dist directory, which will be used to push to PyPI.

contrib/build-wheel.sh

@@ -23,6 +23,7 @@ UCX_PLUGINS_DIR="/usr/lib64/ucx"
 NIXL_PLUGINS_DIR="/usr/local/nixl/lib/$ARCH-linux-gnu/plugins"
 OUTPUT_DIR="dist"
 BUILD_NIXL_EP="false"
+TORCH_VERSIONS=""
 
 # Parse arguments
 while [[ $# -gt 0 ]]; do
@@ -69,6 +70,11 @@ while [[ $# -gt 0 ]]; do
             BUILD_NIXL_EP="true"
             shift
             ;;
+        --torch-versions)
+            TORCH_VERSIONS=$2
+            shift
+            shift
+            ;;
         *)
             echo "Unknown argument: $1"
             exit 1
@@ -88,21 +94,248 @@ if [ "$CUDA_MAJOR" -ne 12 ] && [ "$CUDA_MAJOR" -ne 13 ]; then
     echo "Invalid CUDA_MAJOR: '$CUDA_MAJOR'"
     exit 1
 fi
+AUDITWHEEL_EXCLUDES="--exclude libcuda* --exclude libcufile* --exclude libssl* --exclude libcrypto* --exclude libefa* --exclude libhwloc* --exclude libfabric* --exclude libtorch* --exclude libc10* --exclude libdoca*"
+
 PKG_NAME="nixl-cu${CUDA_MAJOR}"
+CU_TAG="cu$(nvcc --version | grep -Eo 'release [0-9]+\.[0-9]+' | cut -d' ' -f2 | tr -d .)"
 ./contrib/tomlutil.py --wheel-name $PKG_NAME pyproject.toml
-if [ "$BUILD_NIXL_EP" = "true" ]; then
-    uv build --wheel --out-dir $TMP_DIR --python $PYTHON_VERSION \
-        -Csetup-args=-Dbuild_nixl_ep=true \
-        -Csetup-args=-Dbuild_examples=true
+
+# Index URLs for the resolved CUDA tag. Stable cu index hosts released
+# wheels; nightly index hosts dev/pre-release wheels.
+TORCH_STABLE_INDEX="https://download.pytorch.org/whl/${CU_TAG}"
+TORCH_NIGHTLY_INDEX="https://download.pytorch.org/whl/nightly/${CU_TAG}"
+
+# Build deps installed into the per-iteration venv (torch is added
+# separately with channel-appropriate constraints).
+BUILD_DEPS=(
+    "meson"
+    "meson-python"
+    "pybind11"
+    "patchelf"
+    "pyyaml"
+    "types-PyYAML"
+    "setuptools>=80.9.0"
+)
+
+# Classification cache for repeated lookups within a single script run.
+# Values: "stable" | "nightly" | "unavailable".
+declare -A TORCH_CLASS_CACHE
+
+# Slugify a dotted version (e.g. "2.13" -> "213", "3.10" -> "310") so it can
+# be used unambiguously as a path component.
+slug() { echo "${1//./}"; }
+
+# Path for a per-iteration build venv. One venv per (python, torch) tuple
+# so torch's transitive footprint (nvidia-*, triton, sympy, …) never bleeds
+# across torch versions. Lives in /workspace, not /tmp, so it inherits the
+# image's UV_CACHE_DIR layout and is visible to debugging.
+venv_path() {
+    local VER=${1:-}
+    if [ -n "$VER" ]; then
+        echo "/workspace/venv-torch$(slug "$VER")-py$(slug "$PYTHON_VERSION")"
+    else
+        echo "/workspace/venv-py$(slug "$PYTHON_VERSION")"
+    fi
+}
+
+# Classify a requested torch version for the target Python: does
+# `torch==${VER}.*` resolve from the stable cu index alone, only from the
+# nightly index (with --pre), or from neither? Echoes one of:
+#   stable       — resolvable from $TORCH_STABLE_INDEX without --pre
+#   nightly      — resolvable from $TORCH_NIGHTLY_INDEX with --pre
+#   unavailable  — not in either index for this Python
+# Result is cached so each (PYTHON_VERSION, VER) probe runs at most once
+# across the script's filter and build phases.
+torch_classify() {
+    local VER=$1
+    if [ -n "${TORCH_CLASS_CACHE[$VER]:-}" ]; then
+        echo "${TORCH_CLASS_CACHE[$VER]}"
+        return
+    fi
+    local CLASS="unavailable"
+    local PROBE="/workspace/venv-probe-py$(slug "$PYTHON_VERSION")"
+    rm -rf "$PROBE"
+    if uv venv "$PROBE" --python "$PYTHON_VERSION" >/dev/null 2>&1; then
+        if uv pip install --dry-run \
+            --python "$PROBE/bin/python" \
+            --index-url "$TORCH_STABLE_INDEX" \
+            "torch==${VER}.*" >/dev/null 2>&1; then
+            CLASS="stable"
+        elif uv pip install --dry-run --pre \
+            --python "$PROBE/bin/python" \
+            --extra-index-url "$TORCH_STABLE_INDEX" \
+            --extra-index-url "$TORCH_NIGHTLY_INDEX" \
+            --index-strategy unsafe-best-match \
+            "torch==${VER}.*" >/dev/null 2>&1; then
+            CLASS="nightly"
+        fi
+    fi
+    rm -rf "$PROBE"
+    TORCH_CLASS_CACHE[$VER]="$CLASS"
+    echo "$CLASS"
+}
+
+# Echo the torch requirement spec for the given (version, channel). Stable
+# stays as `torch==X.Y.*`; nightly uses an explicit `.dev0` lower bound so
+# pre-release/dev wheels are admissible without --prerelease=allow.
+torch_spec() {
+    local VER=$1
+    local CHANNEL=$2
+    if [ "$CHANNEL" = "nightly" ]; then
+        local MAJOR="${VER%%.*}"
+        local MINOR="${VER##*.}"
+        echo "torch>=${MAJOR}.${MINOR}.0.dev0,<${MAJOR}.$((MINOR + 1))"
+    else
+        echo "torch==${VER}.*"
+    fi
+}
+
+# Echo the uv index/resolution flags appropriate for the given channel,
+# one flag per line. Stable builds use only the stable cu index and no
+# pre-release allowance; nightly builds add the nightly index and
+# --prerelease=allow.
+torch_uv_flags() {
+    local CHANNEL=$1
+    if [ "$CHANNEL" = "nightly" ]; then
+        printf '%s\n' \
+            --extra-index-url "$TORCH_STABLE_INDEX" \
+            --extra-index-url "$TORCH_NIGHTLY_INDEX" \
+            --index-strategy unsafe-best-match \
+            --prerelease allow
+    else
+        printf '%s\n' \
+            --extra-index-url "$TORCH_STABLE_INDEX" \
+            --index-strategy unsafe-best-match
+    fi
+}
+
+# Build the wheel for the current PYTHON_VERSION (and optional torch VER).
+# Creates a fresh venv at venv_path, installs build deps + torch with the
+# channel-appropriate flags, runs `uv build --no-build-isolation`, and
+# tears the venv down so the next iteration starts from a clean slate.
+# Doing it this way instead of `pip install --reinstall torch` avoids
+# orphan packages from the previous torch's transitive footprint
+# (nvidia-* wheels, triton, sympy, …) bleeding across iterations.
+build_wheel() {
+    local OUT_DIR=$1
+    local VER=${2:-}
+
+    local VENV_PATH
+    VENV_PATH=$(venv_path "$VER")
+    local CHANNEL="stable"
+    [ -n "$VER" ] && CHANNEL=$(torch_classify "$VER")
+
+    local UV_FLAGS=()
+    while IFS= read -r f; do UV_FLAGS+=("$f"); done < <(torch_uv_flags "$CHANNEL")
+
+    local TORCH_PKG=()
+    [ -n "$VER" ] && TORCH_PKG+=("$(torch_spec "$VER" "$CHANNEL")")
+
+    echo "=== Provisioning ${VENV_PATH} (python ${PYTHON_VERSION}${VER:+, torch ${VER} [${CHANNEL}]}) ==="
+    rm -rf "$VENV_PATH"
+    uv venv "$VENV_PATH" --python "$PYTHON_VERSION"
+    uv pip install \
+        --python "$VENV_PATH/bin/python" \
+        "${UV_FLAGS[@]}" \
+        "${BUILD_DEPS[@]}" \
+        "${TORCH_PKG[@]}"
+
+    # Activate so meson's `find_installation('python3')` resolves to this
+    # venv's interpreter (which has the right torch). Deactivate before
+    # returning so the caller's auditwheel keeps using the orchestration
+    # venv on PATH.
+    # shellcheck disable=SC1091
+    source "$VENV_PATH/bin/activate"
+
+    local BUILD_ARGS=(
+        --wheel
+        --no-build-isolation
+        --out-dir "$OUT_DIR"
+        --python "$VENV_PATH/bin/python"
+    )
+    if [ "$BUILD_NIXL_EP" = "true" ]; then
+        BUILD_ARGS+=(
+            -Csetup-args=-Dbuild_nixl_ep=true
+            -Csetup-args=-Dbuild_examples=true
+        )
+    fi
+    uv build "${BUILD_ARGS[@]}"
+
+    deactivate
+    # Free disk: torch + nvidia-* wheels in a venv add up to several GB;
+    # 3 torches × 5 pythons would otherwise blow the docker layer budget.
+    rm -rf "$VENV_PATH"
+}
+
+repair_wheel() {
+    local IN_DIR=$1
+    local OUT_DIR=$2
+    mkdir -p "$OUT_DIR"
+    auditwheel repair $AUDITWHEEL_EXCLUDES "$IN_DIR"/nixl*.whl --plat $WHL_PLATFORM --wheel-dir "$OUT_DIR"
+    ./contrib/wheel_add_ucx_plugins.py --ucx-plugins-dir $UCX_PLUGINS_DIR --nixl-plugins-dir $NIXL_PLUGINS_DIR "$OUT_DIR"/*.whl
+}
+
+if [ "$BUILD_NIXL_EP" = "true" ] && [ -n "$TORCH_VERSIONS" ]; then
+    # Multi-torch: build full wheel with first torch, then merge extra .so from others.
+    IFS=',' read -ra TORCH_REQUESTED <<< "$TORCH_VERSIONS"
+
+    # Filter to torch versions actually resolvable for this (Python, CUDA) combo.
+    TORCH_ARRAY=()
+    SKIPPED=()
+    for TORCH in "${TORCH_REQUESTED[@]}"; do
+        if [ "$(torch_classify "$TORCH")" = "unavailable" ]; then
+            SKIPPED+=("$TORCH")
+        else
+            TORCH_ARRAY+=("$TORCH")
+        fi
+    done
+
+    if [ ${#SKIPPED[@]} -gt 0 ]; then
+        echo "=== Skipping torch versions (no wheel on index for Python ${PYTHON_VERSION} + ${CU_TAG}): ${SKIPPED[*]} ==="
+    fi
+    if [ ${#TORCH_ARRAY[@]} -eq 0 ]; then
+        echo "ERROR: none of the requested torch versions (${TORCH_REQUESTED[*]}) are available for Python ${PYTHON_VERSION} + ${CU_TAG}"
+        exit 1
+    fi
+    echo "=== Building for torch versions: ${TORCH_ARRAY[*]} ==="
+
+    FIRST_TORCH="${TORCH_ARRAY[0]}"
+    echo "=== Building wheel with torch ${FIRST_TORCH} ==="
+    build_wheel "$TMP_DIR" "$FIRST_TORCH"
+    repair_wheel "$TMP_DIR" "$TMP_DIR/dist"
+    BASE_WHL=$(ls "$TMP_DIR"/dist/*.whl)
+
+    for ((i=1; i<${#TORCH_ARRAY[@]}; i++)); do
+        TORCH="${TORCH_ARRAY[$i]}"
+        echo "=== Building nixl_ep .so for torch ${TORCH} ==="
+
+        EP_TMP=$(mktemp -d)
+        build_wheel "$EP_TMP" "$TORCH"
+        # Repair so the .so passes auditwheel for the manylinux tag before we extract it
+        repair_wheel "$EP_TMP" "$EP_TMP/dist"
+
+        # Merge the torch-versioned .so from the freshly-built wheel into
+        # the base wheel and regenerate RECORD. Both wheels were built
+        # against the same outer C++ build, so the .so's DT_NEEDED entries
+        # (libucp-<hash>.so etc.) match the libs already bundled in
+        # $BASE_WHL by auditwheel.
+        TORCH_MM=$(echo "$TORCH" | tr -d '.')
+        EP_WHL=$(ls "$EP_TMP"/dist/*.whl)
+        ./contrib/wheel_merge.py \
+            --base-wheel "$BASE_WHL" \
+            --source-wheel "$EP_WHL" \
+            --pattern "nixl_ep_cpp_torch${TORCH_MM}.*" \
+            --target-dir "nixl_ep_cu${CUDA_MAJOR}"
+
+        rm -rf "$EP_TMP"
+    done
+
+    cp "$BASE_WHL" "$OUTPUT_DIR"
 else
-    uv build --wheel --out-dir $TMP_DIR --python $PYTHON_VERSION
+    build_wheel "$TMP_DIR"
+    repair_wheel "$TMP_DIR" "$TMP_DIR/dist"
+    cp "$TMP_DIR"/dist/*.whl "$OUTPUT_DIR"
 fi
 
-# Bundle libraries
-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*' $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"