introduce libraft wheels (#2531)

Replaces #2306, contributes to rapidsai/build-planning#33. Proposes packaging `libraft` as a wheel, which is then re-used by: * `pylibraft-cu{11,12}` and `raft-cu{11,12}` (this PR) * `libcugraph-cu{11,12}`, `pylibcugraph-cu{11,12}`, and `cugraph-cu{11,12}` in rapidsai/cugraph#4804 * `libcuml-cu{11,12}` and `cuml-cu{11,12}` in rapidsai/cuml#6199 As part of this, also proposes: * introducing a new CMake option, `RAFT_COMPILE_DYNAMIC_ONLY`, to allow building/installing only the dynamic shared library (i.e. skipping the static library) * enforcing `rapids-cmake`'s preferred CMake style (#2531 (comment)) * making wheel-building CI jobs always depend on other wheel-building CI jobs, not tests or `*-publish` (to reduce end-to-end CI time) ## Notes for Reviewers ### Benefits of these changes * smaller wheels (see "Size Changes" below) * faster compile times (no more re-compiling RAFT in cuGraph and cuML CI) * other benefits mentioned in rapidsai/build-planning#33 ### Wheel contents `libraft`: * `libraft.so` (shared library) * RAFT headers * vendored dependencies (`fmt`, CCCL, `cuco`, `cute`, `cutlass`) `pylibraft`: * `pylibraft` Python / Cython code and compiled Cython extensions `raft-dask`: * `raft-dask` Python / Cython code and compiled Cython extension ### Dependency Flows In short.... `libraft` contains a `libraft.so` dynamic library and the headers to link against it. * Anything that needs to link against RAFT at build time pulls in `libraft` wheels as a build dependency. * Anything that needs RAFT's symbols at runtime pulls it in as a runtime dependency, and calls `libraft.load_library()`. For more details and some flowcharts, see rapidsai/build-planning#33 (comment) ### Size changes (CUDA 12, Python 3.12, x86_64) | wheel | num files (before) | num files (these PRs) | size (before) | size (these PRs) | |:---------------:|------------------:|-----------------:|--------------:|-------------:| | `libraft`. | --- | 3169 | --- | 19M | | `pylibraft` | 64 | 63 | 11M | 1M | | `raft-dask` | 29 | 28 | 188M | 188M | | `libcugraph` | --- | 1762 | --- | 903M | | `pylibcugraph` | 190 | 187 | 901M | 2M | | `cugraph` | 315 | 313 | 899M | 3.0M | | `libcuml` | --- | 1766 | --- | 289M | | `cuml` | 442 | --- | 517M | --- | |**TOTAL** | **1,040** | **7,268** | **2,516M** | **1,405M** | *NOTES: size = compressed, "before" = 2025-01-13 nightlies* <details><summary>how I calculated those (click me)</summary> * `cugraph`: nightly commit = rapidsai/cugraph@8507cbf, PR = rapidsai/cugraph#4804 * `cuml`: nightly commit = rapidsai/cuml@7c715c4, PR = rapidsai/cuml#6199 * `raft`: nightly commit = 1b62c41, PR = this PR ```shell docker run \ --rm \ --network host \ --env RAPIDS_NIGHTLY_DATE=2025-01-13 \ --env CUGRAPH_NIGHTLY_SHA=8507cbf63db2f349136b266d3e6e787b189f45a0 \ --env CUGRAPH_PR="pull-request/4804" \ --env CUGRAPH_PR_SHA="2ef32eaa006a84c0bd16220bb8e8af34198fbee8" \ --env CUML_NIGHTLY_SHA=7c715c494dff71274d0fdec774bdee12a7e78827 \ --env CUML_PR="pull-request/6199" \ --env CUML_PR_SHA="2ef32eaa006a84c0bd16220bb8e8af34198fbee8" \ --env RAFT_NIGHTLY_SHA=1b62c4117a35b11ce3c830daae248e32ebf75e3f \ --env RAFT_PR="pull-request/2531" \ --env RAFT_PR_SHA="0d6597b08919f2aae8ac268f1a68d6a8fe5beb4e" \ --env RAPIDS_PY_CUDA_SUFFIX=cu12 \ --env WHEEL_DIR_BEFORE=/tmp/wheels-before \ --env WHEEL_DIR_AFTER=/tmp/wheels-after \ -it rapidsai/ci-wheel:cuda12.5.1-rockylinux8-py3.12 \ bash # --- nightly wheels --- # mkdir -p ./wheels-before export RAPIDS_BUILD_TYPE=branch export RAPIDS_REF_NAME="branch-25.02" # pylibraft RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/raft \ RAPIDS_SHA=${RAFT_NIGHTLY_SHA} \ rapids-download-wheels-from-s3 python ./wheels-before # raft-dask RAPIDS_PY_WHEEL_NAME="raft_dask_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/raft \ RAPIDS_SHA=${RAFT_NIGHTLY_SHA} \ rapids-download-wheels-from-s3 python ./wheels-before # cugraph RAPIDS_PY_WHEEL_NAME="cugraph_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/cugraph \ RAPIDS_SHA=${CUGRAPH_NIGHTLY_SHA} \ rapids-download-wheels-from-s3 python ./wheels-before # pylibcugraph RAPIDS_PY_WHEEL_NAME="pylibcugraph_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/cugraph \ RAPIDS_SHA=${CUGRAPH_NIGHTLY_SHA} \ rapids-download-wheels-from-s3 python ./wheels-before # cuml RAPIDS_PY_WHEEL_NAME="cuml_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/cuml \ RAPIDS_SHA=${CUML_NIGHTLY_SHA} \ rapids-download-wheels-from-s3 python ./wheels-before # --- wheels from CI --- # mkdir -p ./wheels-after export RAPIDS_BUILD_TYPE="pull-request" # libraft RAPIDS_PY_WHEEL_NAME="libraft_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/raft \ RAPIDS_REF_NAME="${RAFT_PR}" \ RAPIDS_SHA="${RAFT_PR_SHA}" \ rapids-download-wheels-from-s3 cpp ./wheels-after # pylibraft RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/raft \ RAPIDS_REF_NAME="${RAFT_PR}" \ RAPIDS_SHA="${RAFT_PR_SHA}" \ rapids-download-wheels-from-s3 python ./wheels-after # raft-dask RAPIDS_PY_WHEEL_NAME="raft_dask_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/raft \ RAPIDS_REF_NAME="${RAFT_PR}" \ RAPIDS_SHA="${RAFT_PR_SHA}" \ rapids-download-wheels-from-s3 python ./wheels-after # libcugraph RAPIDS_PY_WHEEL_NAME="libcugraph_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/cugraph \ RAPIDS_REF_NAME="${CUGRAPH_PR}" \ RAPIDS_SHA="${CUGRAPH_PR_SHA}" \ rapids-download-wheels-from-s3 cpp ./wheels-after # pylibcugraph RAPIDS_PY_WHEEL_NAME="pylibcugraph_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/cugraph \ RAPIDS_REF_NAME="${CUGRAPH_PR}" \ RAPIDS_SHA="${CUGRAPH_PR_SHA}" \ rapids-download-wheels-from-s3 python ./wheels-after # cugraph RAPIDS_PY_WHEEL_NAME="cugraph_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/cugraph \ RAPIDS_REF_NAME="${CUGRAPH_PR}" \ RAPIDS_SHA="${CUGRAPH_PR_SHA}" \ rapids-download-wheels-from-s3 python ./wheels-after # libcuml RAPIDS_PY_WHEEL_NAME="libcuml_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/cuml \ RAPIDS_REF_NAME="${CUML_PR}" \ RAPIDS_SHA="${CUML_PR_SHA}" \ rapids-download-wheels-from-s3 cpp ./wheels-after # cuml RAPIDS_PY_WHEEL_NAME="cuml_${RAPIDS_PY_CUDA_SUFFIX}" \ RAPIDS_REPOSITORY=rapidsai/cuml \ RAPIDS_REF_NAME="${CUML_PR}" \ RAPIDS_SHA="${CUML_PR_SHA}" \ rapids-download-wheels-from-s3 python ./wheels-after pip install pydistcheck pydistcheck \ --inspect \ --select 'distro-too-large-compressed' \ ./wheels-before/*.whl \ | grep -E '^checking|files: | compressed' \ > ./before.txt # get more exact sizes du -sh ./wheels-before/* pydistcheck \ --inspect \ --select 'distro-too-large-compressed' \ ./wheels-after/*.whl \ | grep -E '^checking|files: | compressed' \ > ./after.txt # get more exact sizes du -sh ./wheels-after/* ``` </details> ### How I tested this These other PRs: * rapidsai/devcontainers#435 * rapidsai/cugraph-gnn#110 * rapidsai/cuml#6199 * rapidsai/cugraph#4804
rapidsai · Jan 16, 2025 · 8299f17 · 8299f17
1 parent fb6bfe6
commit 8299f17
Show file tree

Hide file tree

Showing 36 changed files with 643 additions and 176 deletions.
diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
@@ -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/[email protected]
+    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/[email protected]
+    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-build-libraft
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]
     with:
@@ -85,7 +108,9 @@ jobs:
       sha: ${{ inputs.sha }}
       date: ${{ inputs.date }}
       package-name: pylibraft
+      package-type: python
   wheel-build-raft-dask:
+    needs: wheel-build-libraft
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]
     with:
@@ -104,3 +129,4 @@ jobs:
       sha: ${{ inputs.sha }}
       date: ${{ inputs.date }}
       package-name: raft_dask
+      package-type: python
diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml
@@ -21,6 +21,7 @@ jobs:
       - conda-python-build
       - conda-python-tests
       - docs-build
+      - wheel-build-libraft
       - wheel-build-pylibraft
       - wheel-tests-pylibraft
       - wheel-build-raft-dask
@@ -116,10 +117,22 @@ jobs:
       arch: "amd64"
       container_image: "rapidsai/ci-conda:latest"
       run_script: "ci/build_docs.sh"
-  wheel-build-pylibraft:
+  wheel-build-libraft:
     needs: checks
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]
+    with:
+      build_type: pull-request
+      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, wheel-build-libraft]
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/[email protected]
     with:
       build_type: pull-request
       script: ci/build_wheel_pylibraft.sh
@@ -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/[email protected]
     with:

diff --git a/build.sh b/build.sh
@@ -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

diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
@@ -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.
@@ -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"
@@ -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
diff --git a/ci/build_wheel_libraft.sh b/ci/build_wheel_libraft.sh
@@ -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="${EXTRA_CMAKE_ARGS}"
+
+ci/build_wheel.sh libraft ${package_dir} cpp
+ci/validate_wheel.sh ${package_dir} final_dist libraft
diff --git a/ci/build_wheel_pylibraft.sh b/ci/build_wheel_pylibraft.sh
@@ -5,17 +5,16 @@ 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})"
 
-# 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}"
+# 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"
 
-ci/build_wheel.sh pylibraft ${package_dir}
+ci/build_wheel.sh pylibraft ${package_dir} python
 ci/validate_wheel.sh ${package_dir} final_dist pylibraft
diff --git a/ci/build_wheel_raft_dask.sh b/ci/build_wheel_raft_dask.sh
@@ -5,8 +5,16 @@ set -euo pipefail
 
 package_dir="python/raft-dask"
 
-# Set up skbuild options. Enable sccache in skbuild config options
-export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF"
+RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
 
-ci/build_wheel.sh raft-dask ${package_dir}
+# 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"
+
+ci/build_wheel.sh raft-dask ${package_dir} python
 ci/validate_wheel.sh ${package_dir} final_dist raft-dask
diff --git a/ci/check_style.sh b/ci/check_style.sh
@@ -14,5 +14,12 @@ rapids-dependency-file-generator \
 rapids-mamba-retry env create --yes -f env.yaml -n checks
 conda activate checks
 
+# get config for cmake-format checks
+RAPIDS_VERSION_MAJOR_MINOR="$(rapids-version-major-minor)"
+FORMAT_FILE_URL="https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-${RAPIDS_VERSION_MAJOR_MINOR}/cmake-format-rapids-cmake.json"
+export RAPIDS_CMAKE_FORMAT_FILE=/tmp/rapids_cmake_ci/cmake-formats-rapids-cmake.json
+mkdir -p $(dirname ${RAPIDS_CMAKE_FORMAT_FILE})
+wget -O ${RAPIDS_CMAKE_FORMAT_FILE} ${FORMAT_FILE_URL}
+
 # Run pre-commit checks
 pre-commit run --all-files --show-diff-on-failure
diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh
@@ -43,6 +43,8 @@ echo "${NEXT_FULL_TAG}" > VERSION
 
 DEPENDENCIES=(
   dask-cuda
+  libraft
+  librmm
   pylibraft
   rmm
   rapids-dask-dependency

diff --git a/ci/test_wheel_pylibraft.sh b/ci/test_wheel_pylibraft.sh
@@ -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
diff --git a/ci/test_wheel_raft_dask.sh b/ci/test_wheel_raft_dask.sh
@@ -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]"
 

diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh
@@ -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}"