Skip to content

WIP: build wheels with CUDA 13.0.x, test wheels against mix of CTK versions #604

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.

Sign up for GitHub

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

Draft
jameslamb wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

WIP: build wheels with CUDA 13.0.x, test wheels against mix of CTK versions #604

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b1e5230
build wheels with CUDA 13.0.x, test wheels against mix of CTK versions
jameslamb Mar 5, 2026
bf3b694
loosen 12.2 constraint slightly
jameslamb Mar 5, 2026
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
6 changes: 3 additions & 3 deletions .github/workflows/build.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ jobs:
sha: ${{ inputs.sha }}
wheel-build-libucxx:
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@python-3.14
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@build-wheels-old-ctk
with:
build_type: ${{ inputs.build_type || 'branch' }}
branch: ${{ inputs.branch }}
Expand All @@ -106,7 +106,7 @@ jobs:
wheel-build-ucxx:
needs: wheel-build-libucxx
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@python-3.14
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@build-wheels-old-ctk
with:
build_type: ${{ inputs.build_type || 'branch' }}
branch: ${{ inputs.branch }}
Expand All @@ -132,7 +132,7 @@ jobs:
publish-wheel-search-key: ucxx_wheel_python_abi3
wheel-build-distributed-ucxx:
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@python-3.14
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@build-wheels-old-ctk
with:
build_type: ${{ inputs.build_type || 'branch' }}
branch: ${{ inputs.branch }}
Expand Down
6 changes: 3 additions & 3 deletions .github/workflows/pr.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,7 +211,7 @@ jobs:
wheel-build-libucxx:
needs: checks
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@python-3.14
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@build-wheels-old-ctk
with:
build_type: pull-request
node_type: cpu8
Expand All @@ -223,7 +223,7 @@ jobs:
wheel-build-ucxx:
needs: wheel-build-libucxx
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@python-3.14
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@build-wheels-old-ctk
with:
build_type: pull-request
node_type: cpu8
Expand All @@ -244,7 +244,7 @@ jobs:
wheel-build-distributed-ucxx:
needs: checks
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@python-3.14
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@build-wheels-old-ctk
with:
build_type: pull-request
node_type: cpu8
Expand Down
3 changes: 3 additions & 0 deletions ci/build_wheel_distributed_ucxx.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@ set -euo pipefail

source rapids-init-pip

# TODO(jameslamb): revert before merging
source ci/use_wheels_from_prs.sh

package_dir="python/distributed-ucxx"

./ci/build_wheel.sh distributed-ucxx "${package_dir}"
Expand Down
3 changes: 3 additions & 0 deletions ci/build_wheel_libucxx.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@ set -euo pipefail

source rapids-init-pip

# TODO(jameslamb): revert before merging
source ci/use_wheels_from_prs.sh

package_name="libucxx"
package_dir="python/libucxx"

Expand Down
3 changes: 3 additions & 0 deletions ci/build_wheel_ucxx.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@ set -euo pipefail

source rapids-init-pip

# TODO(jameslamb): revert before merging
source ci/use_wheels_from_prs.sh

package_name="ucxx"
package_dir="python/ucxx"

Expand Down
21 changes: 21 additions & 0 deletions ci/test_wheel_distributed_ucxx.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,18 @@

set -euo pipefail

# TODO(jameslamb): revert before merging
git clone --branch generate-pip-constraints \
https://github.com/rapidsai/gha-tools.git \
/tmp/gha-tools

export PATH="/tmp/gha-tools/tools:${PATH}"

source rapids-init-pip

# TODO(jameslamb): revert before merging
source ci/use_wheels_from_prs.sh

package_name="distributed_ucxx"

source "$(dirname "$0")/test_common.sh"
Expand All @@ -16,8 +26,19 @@ distributed_ucxx_wheelhouse=$(rapids-download-from-github "$(rapids-package-name
ucxx_wheelhouse=$(rapids-download-from-github "$(rapids-package-name "wheel_python" ucxx --stable --cuda "$RAPIDS_CUDA_VERSION")")
libucxx_wheelhouse=$(RAPIDS_PY_WHEEL_NAME="libucxx_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp)

# generate constraints (possibly pinning to oldest support versions of dependencies)
rapids-generate-pip-constraints test_python "${PIP_CONSTRAINT}"

# notes:
#
# * echo to expand wildcard before adding `[test]` requires for pip
# * just providing --constraint="${PIP_CONSTRAINT}" to be explicit, and because
# that environment variable is ignored if any other --constraint are passed via the CLI
#
rapids-pip-retry install \
-v \
--prefer-binary \
--constraint "${PIP_CONSTRAINT}" \
"${libucxx_wheelhouse}/libucxx_${RAPIDS_PY_CUDA_SUFFIX}"*.whl \
"${ucxx_wheelhouse}/ucxx_${RAPIDS_PY_CUDA_SUFFIX}"*.whl \
"$(echo "${distributed_ucxx_wheelhouse}"/"${package_name}_${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test]"
Expand Down
21 changes: 21 additions & 0 deletions ci/test_wheel_ucxx.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,18 @@

set -euo pipefail

# TODO(jameslamb): revert before merging
git clone --branch generate-pip-constraints \
https://github.com/rapidsai/gha-tools.git \
/tmp/gha-tools

export PATH="/tmp/gha-tools/tools:${PATH}"

source rapids-init-pip

# TODO(jameslamb): revert before merging
source ci/use_wheels_from_prs.sh

package_name="ucxx"

source "$(dirname "$0")/test_common.sh"
Expand All @@ -17,8 +27,19 @@ RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")"
libucxx_wheelhouse=$(RAPIDS_PY_WHEEL_NAME="libucxx_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp)
ucxx_wheelhouse=$(rapids-download-from-github "$(rapids-package-name "wheel_python" ucxx --stable --cuda "$RAPIDS_CUDA_VERSION")")

# generate constraints (possibly pinning to oldest support versions of dependencies)
rapids-generate-pip-constraints test_python "${PIP_CONSTRAINT}"

# notes:
#
# * echo to expand wildcard before adding `[test]` requires for pip
# * just providing --constraint="${PIP_CONSTRAINT}" to be explicit, and because
# that environment variable is ignored if any other --constraint are passed via the CLI
#
rapids-pip-retry install \
-v \
--prefer-binary \
--constraint "${PIP_CONSTRAINT}" \
"${libucxx_wheelhouse}/libucxx_${RAPIDS_PY_CUDA_SUFFIX}"*.whl \
"$(echo "${ucxx_wheelhouse}"/"${package_name}_${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test]"

Expand Down
23 changes: 23 additions & 0 deletions ci/use_wheels_from_prs.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
# SPDX-License-Identifier: BSD-3-Clause

# initialize PIP_CONSTRAINT
source rapids-init-pip

RAPIDS_PY_CUDA_SUFFIX=$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")

# download wheels, store the directories holding them in variables
RMM_COMMIT=af96977c404565bbb0657c490d407a561cedc3fc
LIBRMM_WHEELHOUSE=$(
RAPIDS_PY_WHEEL_NAME="librmm_${RAPIDS_PY_CUDA_SUFFIX}" rapids-get-pr-artifact rmm 2270 cpp wheel "${RMM_COMMIT}"
)
RMM_WHEELHOUSE=$(
rapids-get-pr-artifact rmm 2270 python wheel --pkg_name rmm --stable "${RMM_COMMIT}"
)

# write a pip constraints file saying e.g. "whenever you encounter a requirement for 'librmm-cu12', use this wheel"
cat >> "${PIP_CONSTRAINT}" <<EOF
librmm-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${LIBRMM_WHEELHOUSE}"/librmm*.whl)
rmm-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${RMM_WHEELHOUSE}"/rmm_*.whl)
EOF
47 changes: 47 additions & 0 deletions dependencies.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,6 +205,8 @@ dependencies:
- output_types: [conda, requirements]
packages:
- pre-commit
# 'cuda_version' intentionally does not contain fallback entries... we want
# a loud error if an unsupported 'cuda' value is passed
cuda_version:
specific:
- output_types: conda
Expand Down Expand Up @@ -233,6 +235,51 @@ dependencies:
cuda: "13.1"
packages:
- cuda-version=13.1
- output_types: requirements
matrices:
# if use_cuda_wheels=false is provided, do not add dependencies on any CUDA wheels
# (e.g. for DLFW and pip devcontainers)
- matrix:
use_cuda_wheels: "false"
packages:
- matrix:
arch: aarch64
cuda: "12.2"
use_cuda_wheels: "true"
packages:
# some components (like nvidia-cublas-cu12 and nvidia-cuda-nvcc-cu12) didn't have
# aarch64 wheels until CTK 12.3, so allow a slightly looser bound here
- cuda-toolkit>=12.2,<12.4
Copy link
Member Author

@jameslamb jameslamb Mar 5, 2026

Choose a reason for hiding this comment

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

For this repo, it's nvidia-cuda-nvcc-cu12 that ends up being required but can't be installed on aarch64 + CUDA 12.2:

docker run --rm -it rapidsai/citestwheel:26.04-cuda12.2.2-ubuntu22.04-py3.11 bash

pip install 'cuda-toolkit[nvcc,nvrtc]==12.2.*'
# ERROR: ResolutionImpossible

pip install 'cuda-toolkit[nvcc,nvrtc]==12.3.0'
# ERROR: No matching distribution found for nvidia-cuda-nvcc-cu12

pip install 'cuda-toolkit[nvcc,nvrtc]==12.3.1'
# ERROR: No matching distribution found for nvidia-cuda-nvcc-cu12

pip install 'cuda-toolkit[nvcc,nvrtc]==12.3.2'
# Successfully installed cuda-toolkit-12.3.2 nvidia-cuda-nvcc-cu12-12.3.107 nvidia-cuda-nvrtc-cu12-12.3.107

But trying to keep this cuda_version: list consistent across all of RAPIDs, so including cuBLAS in the comment too.

- matrix:
cuda: "12.2"
use_cuda_wheels: "true"
packages:
- cuda-toolkit==12.2.*
- matrix:
cuda: "12.5"
use_cuda_wheels: "true"
packages:
- cuda-toolkit==12.5.*
- matrix:
cuda: "12.8"
use_cuda_wheels: "true"
packages:
- cuda-toolkit==12.8.*
- matrix:
cuda: "12.9"
use_cuda_wheels: "true"
packages:
- cuda-toolkit==12.9.*
- matrix:
cuda: "13.0"
use_cuda_wheels: "true"
packages:
- cuda-toolkit==13.0.*
- matrix:
cuda: "13.1"
use_cuda_wheels: "true"
packages:
- cuda-toolkit==13.1.*
cuda:
common:
- output_types: [conda]
Expand Down