Skip to content

[CI] [WheelNext] Build variant wheels #11531

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jul 4, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions .github/workflows/python_wheels_variants.yml
Original file line number Diff line number Diff line change
@@ -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
23 changes: 23 additions & 0 deletions ops/pipeline/build-variant-wheels-impl.sh
Original file line number Diff line number Diff line change
@@ -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 \
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you please provide some comments on what this line does? Skimming the pipeline, this follows a CUDA build, I don't see any potential variant here since there's no CPU build and NCCL is still installed.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is an initial stage, where we build a single variant for the CUDA build. In the future, I plan to build the CPU variant and other variants.

-p "nvidia :: cuda :: 12" -o . --pyproject-toml python-package/pyproject.toml
51 changes: 51 additions & 0 deletions ops/pipeline/build-variant-wheels.sh
Original file line number Diff line number Diff line change
@@ -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
7 changes: 7 additions & 0 deletions python-package/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
7 changes: 7 additions & 0 deletions python-package/pyproject.toml.in
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Loading