[CI] [WheelNext] Build variant wheels (#11531)

* [CI] [WheelNext] Build variant wheels

* Fix script

* Create impl script

Author: hcho3

.github/workflows/python_wheels_variants.yml

@@ -0,0 +1,36 @@
+name: Build Python wheels using Wheel Variant prototype (WheelNext)
+
+on: [push, pull_request]
+
+permissions:
+  contents: read  # to fetch code (actions/checkout)
+
+defaults:
+  run:
+    shell: bash -l {0}
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+env:
+  BRANCH_NAME: >-
+    ${{ github.event.pull_request.number && 'PR-' }}${{ github.event.pull_request.number || github.ref_name }}
+
+jobs:
+  python-wheels-variants:
+    name: Build Python wheels using Wheel Variant prototype (WheelNext)
+    runs-on:
+      - runs-on=${{ github.run_id }}
+      - runner=linux-amd64-cpu
+      - tag=python-wheels-variants
+    steps:
+      # Restart Docker daemon so that it recognizes the ephemeral disks
+      - run: sudo systemctl restart docker
+      - uses: actions/checkout@v4
+        with:
+          submodules: "true"
+      - name: Log into Docker registry (AWS ECR)
+        run: bash ops/pipeline/login-docker-registry.sh
+      - run: |
+          bash ops/pipeline/build-variant-wheels.sh

ops/pipeline/build-variant-wheels-impl.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+## Build Python wheels using Wheel Variant prototype (WheelNext)
+## Companion script for ops/pipeline/build-variant-wheels.sh
+
+set -eo pipefail
+
+set -x
+gosu root chown -R $(id -u):$(id -g) /opt/miniforge/envs /opt/miniforge/pkgs/cache
+gosu root chown $(id -u):$(id -g) /opt/miniforge/pkgs
+set +x
+
+mamba create -y -n wheelnext python=3.13 python-build
+
+source activate wheelnext
+
+# Cannot set -u before Conda env activation
+set -xu
+
+python -m pip install -v \
+  git+https://github.com/wheelnext/pep_xxx_wheel_variants.git@f3b287090f8a6f510b0e1723896e1c7e638f6bff#subdirectory=pep_xxx_wheel_variants
+pip config set --site global.index-url https://variants-index.wheelnext.dev/
+variantlib make-variant -f python-package/dist/xgboost-*.whl \
+  -p "nvidia :: cuda :: 12" -o . --pyproject-toml python-package/pyproject.toml

ops/pipeline/build-variant-wheels.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+## Build Python wheels using Wheel Variant prototype (WheelNext)
+
+set -euo pipefail
+
+if [[ -z "${GITHUB_SHA:-}" ]]
+then
+  echo "Make sure to set environment variable GITHUB_SHA"
+  exit 1
+fi
+
+image_repo='xgb-ci.gpu_build_rockylinux8'
+
+source ops/pipeline/classify-git-branch.sh
+source ops/pipeline/get-docker-registry-details.sh
+source ops/pipeline/get-image-tag.sh
+
+WHEEL_TAG=manylinux_2_28_x86_64
+BUILD_IMAGE_URI="${DOCKER_REGISTRY_URL}/${image_repo}:${IMAGE_TAG}"
+MANYLINUX_IMAGE_URI="${DOCKER_REGISTRY_URL}/xgb-ci.${WHEEL_TAG}:${IMAGE_TAG}"
+
+echo "--- Build with CUDA"
+
+if [[ ($is_pull_request == 1) || ($is_release_branch == 0) ]]
+then
+  export BUILD_ONLY_SM75=1
+else
+  export BUILD_ONLY_SM75=0
+fi
+export USE_RMM=0
+
+set -x
+
+python3 ops/docker_run.py \
+  --image-uri ${BUILD_IMAGE_URI} \
+  --run-args='-e BUILD_ONLY_SM75 -e USE_RMM' \
+  -- ops/pipeline/build-cuda-impl.sh
+
+echo "--- Audit binary wheel to ensure it's compliant with ${WHEEL_TAG} standard"
+python3 ops/docker_run.py \
+  --image-uri ${MANYLINUX_IMAGE_URI} \
+  -- auditwheel repair --only-plat \
+  --plat ${WHEEL_TAG} python-package/dist/*.whl
+python3 -m wheel tags --python-tag py3 --abi-tag none --platform ${WHEEL_TAG} --remove \
+  wheelhouse/*.whl
+mv -v wheelhouse/*.whl python-package/dist/
+
+echo "--- Convert Python wheel to variant wheel"
+python3 ops/docker_run.py \
+  --image-uri ${BUILD_IMAGE_URI} \
+  -- ops/pipeline/build-variant-wheels-impl.sh

python-package/pyproject.toml

@@ -95,3 +95,10 @@ inspect = true
 ignore = ["compiled-objects-have-debug-symbols"]
 max_allowed_size_compressed = '300M'
 max_allowed_size_uncompressed = '500M'
+
+[variant.default-priorities]
+namespace = ["nvidia"]
+
+[variant.providers.nvidia]
+requires = ["nvidia-variant-provider>=0.0.1,<1.0.0"]
+plugin-api = "nvidia_variant_provider.plugin:NvidiaVariantPlugin"

python-package/pyproject.toml.in

@@ -94,3 +94,10 @@ inspect = true
 ignore = ["compiled-objects-have-debug-symbols"]
 max_allowed_size_compressed = '300M'
 max_allowed_size_uncompressed = '500M'
+
+[variant.default-priorities]
+namespace = ["nvidia"]
+
+[variant.providers.nvidia]
+requires = ["nvidia-variant-provider>=0.0.1,<1.0.0"]
+plugin-api = "nvidia_variant_provider.plugin:NvidiaVariantPlugin"