Skip to content
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

WIP: [DO NOT MERGE] introduce libraft wheels #2531

Draft
wants to merge 32 commits into
base: branch-25.02
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
32 commits
Select commit Hold shift + click to select a range
24a4db8
start adding libraft wheels
jameslamb Dec 16, 2024
c0f059c
libraft minimally working
jameslamb Dec 16, 2024
d116417
pylibraft builds mostly working
jameslamb Dec 16, 2024
eb17735
more changes
jameslamb Dec 17, 2024
d317235
comment out more CI
jameslamb Dec 17, 2024
679d289
more changes
jameslamb Dec 17, 2024
51e7041
more dependencies.yaml fiddling
jameslamb Dec 17, 2024
6d38c19
actually run builds
jameslamb Dec 17, 2024
717df6f
pre-commit
jameslamb Dec 17, 2024
d815242
merge branch-25.02
jameslamb Dec 17, 2024
11168a1
install libraft wheels
jameslamb Dec 17, 2024
46bfcfa
fewer builds, fix libraft downloads
jameslamb Dec 17, 2024
97f5d31
fix branches
jameslamb Dec 17, 2024
f492d59
more fixes to CI scripts
jameslamb Dec 17, 2024
ad13d08
download libraft in wheel test jobs, change file-size limits
jameslamb Dec 17, 2024
6bf5eba
pylibraft no longer needs CUDA-wheels runtime dependency
jameslamb Dec 17, 2024
9783244
update thresholds
jameslamb Dec 18, 2024
d0b6385
Merge branch 'branch-25.02' into libraft-wheels
jameslamb Dec 18, 2024
ee99702
try re-enabling conda CI
jameslamb Dec 24, 2024
0ac23f4
Merge branch 'branch-25.02' of github.com:rapidsai/raft into libraft-…
jameslamb Dec 30, 2024
48c98b8
more changes
jameslamb Jan 6, 2025
2070b11
Merge branch 'branch-25.02' into libraft-wheels
jameslamb Jan 6, 2025
345f0e5
Merge branch 'branch-25.02' into libraft-wheels
jameslamb Jan 6, 2025
4c657b2
comments
jameslamb Jan 7, 2025
dbbfbb5
merge branch-25.02
jameslamb Jan 7, 2025
6deeb51
Merge branch 'libraft-wheels' of github.com:jameslamb/raft into libra…
jameslamb Jan 7, 2025
c6579e1
exclude static raft from wheels
jameslamb Jan 9, 2025
1fd7603
add a RAFT_BUILD_STATIC_LIB option
jameslamb Jan 9, 2025
a1a905d
be more thorough in skipping the static library
jameslamb Jan 10, 2025
c224e27
more fixes
jameslamb Jan 10, 2025
76d788b
clean up dependencies
jameslamb Jan 10, 2025
d275c99
Merge branch 'branch-25.02' into libraft-wheels
jameslamb Jan 13, 2025
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
25 changes: 25 additions & 0 deletions .github/workflows/build.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,30 @@ jobs:
node_type: "gpu-v100-latest-1"
run_script: "ci/build_docs.sh"
sha: ${{ inputs.sha }}
wheel-build-libraft:
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.02
with:
build_type: ${{ inputs.build_type || 'branch' }}
branch: ${{ inputs.branch }}
sha: ${{ inputs.sha }}
date: ${{ inputs.date }}
script: ci/build_wheel_libraft.sh
# build for every combination of arch and CUDA version, but only for the latest Python
matrix_filter: group_by([.ARCH, (.CUDA_VER|split(".")|map(tonumber)|.[0])]) | map(max_by(.PY_VER|split(".")|map(tonumber)))
wheel-publish-libraft:
needs: wheel-build-libraft
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-25.02
with:
build_type: ${{ inputs.build_type || 'branch' }}
branch: ${{ inputs.branch }}
sha: ${{ inputs.sha }}
date: ${{ inputs.date }}
package-name: libraft
package-type: cpp
wheel-build-pylibraft:
needs: wheel-publish-libraft
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.02
with:
Expand All @@ -85,6 +108,7 @@ jobs:
sha: ${{ inputs.sha }}
date: ${{ inputs.date }}
package-name: pylibraft
package-type: python
wheel-build-raft-dask:
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.02
Expand All @@ -104,3 +128,4 @@ jobs:
sha: ${{ inputs.sha }}
date: ${{ inputs.date }}
package-name: raft_dask
package-type: python
71 changes: 42 additions & 29 deletions .github/workflows/pr.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,16 @@ jobs:
- checks
- conda-cpp-build
- conda-cpp-tests
- conda-cpp-checks
# - conda-cpp-checks
- conda-python-build
- conda-python-tests
- docs-build
# - docs-build
- wheel-build-libraft
- wheel-build-pylibraft
- wheel-tests-pylibraft
- wheel-build-raft-dask
- wheel-tests-raft-dask
- devcontainer
# - devcontainer
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-25.02
if: always()
Expand Down Expand Up @@ -86,13 +87,13 @@ jobs:
if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp
with:
build_type: pull-request
conda-cpp-checks:
needs: conda-cpp-build
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-25.02
with:
build_type: pull-request
enable_check_symbols: true
# conda-cpp-checks:
# needs: conda-cpp-build
# secrets: inherit
# uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-25.02
# with:
# build_type: pull-request
# enable_check_symbols: true
conda-python-build:
needs: conda-cpp-build
secrets: inherit
Expand All @@ -106,18 +107,30 @@ jobs:
if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python
with:
build_type: pull-request
docs-build:
needs: conda-python-build
# docs-build:
# needs: conda-python-build
# secrets: inherit
# uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.02
# with:
# build_type: pull-request
# node_type: "gpu-v100-latest-1"
# arch: "amd64"
# container_image: "rapidsai/ci-conda:latest"
# run_script: "ci/build_docs.sh"
wheel-build-libraft:
needs: checks
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.02
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.02
with:
build_type: pull-request
node_type: "gpu-v100-latest-1"
arch: "amd64"
container_image: "rapidsai/ci-conda:latest"
run_script: "ci/build_docs.sh"
branch: ${{ inputs.branch }}
sha: ${{ inputs.sha }}
date: ${{ inputs.date }}
script: ci/build_wheel_libraft.sh
# build for every combination of arch and CUDA version, but only for the latest Python
matrix_filter: group_by([.ARCH, (.CUDA_VER|split(".")|map(tonumber)|.[0])]) | map(max_by(.PY_VER|split(".")|map(tonumber)))
wheel-build-pylibraft:
needs: checks
needs: [checks, wheel-build-libraft]
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.02
with:
Expand All @@ -132,7 +145,7 @@ jobs:
build_type: pull-request
script: ci/test_wheel_pylibraft.sh
wheel-build-raft-dask:
needs: wheel-tests-pylibraft
needs: [checks, wheel-build-libraft]
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.02
with:
Expand All @@ -146,13 +159,13 @@ jobs:
with:
build_type: pull-request
script: ci/test_wheel_raft_dask.sh
devcontainer:
secrets: inherit
uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@branch-25.02
with:
arch: '["amd64"]'
cuda: '["12.5"]'
build_command: |
sccache -z;
build-all -DBUILD_PRIMS_BENCH=ON --verbose;
sccache -s;
# devcontainer:
# secrets: inherit
# uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@branch-25.02
# with:
# arch: '["amd64"]'
# cuda: '["12.5"]'
# build_command: |
# sccache -z;
# build-all -DBUILD_PRIMS_BENCH=ON --verbose;
# sccache -s;
7 changes: 1 addition & 6 deletions build.sh
Original file line number Diff line number Diff line change
Expand Up @@ -347,13 +347,8 @@ if [[ ${CMAKE_TARGET} == "" ]]; then
CMAKE_TARGET="all"
fi

# Append `-DFIND_RAFT_CPP=ON` to EXTRA_CMAKE_ARGS unless a user specified the option.
SKBUILD_EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS}"
if [[ "${EXTRA_CMAKE_ARGS}" != *"DFIND_RAFT_CPP"* ]]; then
SKBUILD_EXTRA_CMAKE_ARGS="${SKBUILD_EXTRA_CMAKE_ARGS} -DFIND_RAFT_CPP=ON"
fi
# Replace spaces with semicolons in SKBUILD_EXTRA_CMAKE_ARGS
SKBUILD_EXTRA_CMAKE_ARGS=$(echo ${SKBUILD_EXTRA_CMAKE_ARGS} | sed 's/ /;/g')
SKBUILD_EXTRA_CMAKE_ARGS=$(echo ${EXTRA_CMAKE_ARGS} | sed 's/ /;/g')

# If clean given, run it prior to any other steps
if (( ${CLEAN} == 1 )); then
Expand Down
9 changes: 8 additions & 1 deletion ci/build_wheel.sh
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ set -euo pipefail

package_name=$1
package_dir=$2
package_type=$3
underscore_package_name=$(echo "${package_name}" | tr "-" "_")

# Clear out system ucx files to ensure that we're getting ucx from the wheel.
Expand Down Expand Up @@ -39,6 +40,12 @@ case "${RAPIDS_CUDA_VERSION}" in
;;
esac

if [[ ${package_name} != "libraft" ]]; then
EXCLUDE_ARGS+=(
--exclude "libraft.so"
)
fi

sccache --zero-stats

rapids-logger "Building '${package_name}' wheel"
Expand All @@ -55,4 +62,4 @@ sccache --show-adv-stats
mkdir -p final_dist
python -m auditwheel repair -w final_dist "${EXCLUDE_ARGS[@]}" dist/*

RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 python final_dist
RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 ${package_type} final_dist
43 changes: 43 additions & 0 deletions ci/build_wheel_libraft.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
#!/bin/bash
# Copyright (c) 2024, NVIDIA CORPORATION.

set -euo pipefail

package_name="libraft"
package_dir="python/libraft"

rapids-logger "Generating build requirements"
matrix_selectors="cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION};cuda_suffixed=true"

rapids-dependency-file-generator \
--output requirements \
--file-key "py_build_${package_name}" \
--file-key "py_rapids_build_${package_name}" \
--matrix "${matrix_selectors}" \
| tee /tmp/requirements-build.txt

rapids-logger "Installing build requirements"
python -m pip install \
-v \
--prefer-binary \
-r /tmp/requirements-build.txt

# build with '--no-build-isolation', for better sccache hit rate
# 0 really means "add --no-build-isolation" (ref: https://github.com/pypa/pip/issues/5735)
export PIP_NO_BUILD_ISOLATION=0

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

case "${RAPIDS_CUDA_VERSION}" in
12.*)
EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=ON"
;;
11.*)
EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=OFF"
;;
esac

export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF${EXTRA_CMAKE_ARGS}"

ci/build_wheel.sh libraft ${package_dir} cpp
ci/validate_wheel.sh ${package_dir} final_dist libraft
22 changes: 12 additions & 10 deletions ci/build_wheel_pylibraft.sh
Original file line number Diff line number Diff line change
Expand Up @@ -5,17 +5,19 @@ set -euo pipefail

package_dir="python/pylibraft"

case "${RAPIDS_CUDA_VERSION}" in
12.*)
EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=ON"
;;
11.*)
EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=OFF"
;;
esac
RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"

# Downloads libraft wheels from this current build,
# then ensures 'pylibraft' wheel builds always use the 'libraft' just built in the same CI run.
#
# Using env variable PIP_CONSTRAINT is necessary to ensure the constraints
# are used when creating the isolated build environment.
RAPIDS_PY_WHEEL_NAME="libraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 cpp /tmp/libraft_dist
echo "libraft-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo /tmp/libraft_dist/libraft_*.whl)" > /tmp/constraints.txt
export PIP_CONSTRAINT="/tmp/constraints.txt"

# Set up skbuild options. Enable sccache in skbuild config options
export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF${EXTRA_CMAKE_ARGS}"
export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;"

ci/build_wheel.sh pylibraft ${package_dir}
ci/build_wheel.sh pylibraft ${package_dir} python
ci/validate_wheel.sh ${package_dir} final_dist pylibraft
15 changes: 13 additions & 2 deletions ci/build_wheel_raft_dask.sh
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,19 @@ set -euo pipefail

package_dir="python/raft-dask"

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

# Downloads libraft wheels from this current build,
# then ensures 'raft-dask' wheel builds always use the 'libraft' just built in the same CI run.
#
# Using env variable PIP_CONSTRAINT is necessary to ensure the constraints
# are used when creating the isolated build environment.
RAPIDS_PY_WHEEL_NAME="libraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 cpp /tmp/libraft_dist
echo "libraft-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo /tmp/libraft_dist/libraft_*.whl)" > /tmp/constraints.txt
export PIP_CONSTRAINT="/tmp/constraints.txt"

# Set up skbuild options. Enable sccache in skbuild config options
export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF"
export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;"

ci/build_wheel.sh raft-dask ${package_dir}
ci/build_wheel.sh raft-dask ${package_dir} python
ci/validate_wheel.sh ${package_dir} final_dist raft-dask
8 changes: 6 additions & 2 deletions ci/test_wheel_pylibraft.sh
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,13 @@ set -euo pipefail

mkdir -p ./dist
RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist
RAPIDS_PY_WHEEL_NAME="libraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 cpp ./local-libraft-dep
RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 python ./dist


# echo to expand wildcard before adding `[extra]` requires for pip
python -m pip install $(echo ./dist/pylibraft*.whl)[test]
python -m pip install \
./local-libraft-dep/libraft*.whl \
"$(echo ./dist/pylibraft*.whl)[test]"

python -m pytest ./python/pylibraft/pylibraft/test
8 changes: 4 additions & 4 deletions ci/test_wheel_raft_dask.sh
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,13 @@ set -euo pipefail

mkdir -p ./dist
RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
RAPIDS_PY_WHEEL_NAME="raft_dask_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist

# Download the pylibraft built in the previous step
RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./local-pylibraft-dep
RAPIDS_PY_WHEEL_NAME="libraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 cpp ./local-libraft-dep
RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 python ./local-pylibraft-dep
RAPIDS_PY_WHEEL_NAME="raft_dask_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 python ./dist

# echo to expand wildcard before adding `[extra]` requires for pip
python -m pip install -v \
./local-libraft-dep/libraft*.whl \
./local-pylibraft-dep/pylibraft*.whl \
"$(echo ./dist/raft_dask_${RAPIDS_PY_CUDA_SUFFIX}*.whl)[test]"

Expand Down
16 changes: 5 additions & 11 deletions ci/validate_wheel.sh
Original file line number Diff line number Diff line change
Expand Up @@ -10,23 +10,17 @@ package_name=$3
RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}"

# some packages are much larger on CUDA 11 than on CUDA 12
if [[ "${package_name}" == "raft-dask" ]]; then
PYDISTCHECK_ARGS=(
--max-allowed-size-compressed '200M'
)
elif [[ "${package_name}" == "pylibraft" ]]; then
PYDISTCHECK_ARGS=()
if [[ "${package_name}" == "libraft" ]]; then
if [[ "${RAPIDS_CUDA_MAJOR}" == "11" ]]; then
PYDISTCHECK_ARGS=(
--max-allowed-size-compressed '600M'
PYDISTCHECK_ARGS+=(
--max-allowed-size-compressed '750M'
)
else
PYDISTCHECK_ARGS=(
PYDISTCHECK_ARGS+=(
--max-allowed-size-compressed '100M'
)
fi
else
echo "Unsupported package name: ${package_name}"
exit 1
fi

cd "${package_dir}"
Expand Down
Loading
Loading