Build C++ wheel (#2264)

This PR changes wheel building in raft to create a separate C++ wheel that is then found from the Python wheel. The C++ wheel is now a hard dependency of the Python wheel. This allows Python packaging to more closely mirror the structure of our conda packaging, and the way we would normally wish to package these in any other manager. It also allows us to reduce package sizes by allowing better sharing of artifacts between different Python packages that rely on the same C++ components from other packages.

Contributes to rapidsai/build-planning#33

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Ray Douglass (https://github.com/raydouglass)
  - Divye Gala (https://github.com/divyegala)

URL: #2264

.github/workflows/build.yaml

@@ -67,42 +67,46 @@ jobs:
       node_type: "gpu-v100-latest-1"
       run_script: "ci/build_docs.sh"
       sha: ${{ inputs.sha }}
-  wheel-build-pylibraft:
+  wheel-build-cpp:
+    needs: checks
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]
     with:
+      matrix_filter: group_by([.ARCH, (.CUDA_VER|split(".")|map(tonumber)|.[0])]) | map(max_by(.PY_VER|split(".")|map(tonumber)))
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
       sha: ${{ inputs.sha }}
       date: ${{ inputs.date }}
-      script: ci/build_wheel_pylibraft.sh
-  wheel-publish-pylibraft:
-    needs: wheel-build-pylibraft
+      script: ci/build_wheel_cpp.sh
+  wheel-build-python:
+    needs: wheel-build-cpp
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-publish[email protected]
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-build[email protected]
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
       sha: ${{ inputs.sha }}
       date: ${{ inputs.date }}
-      package-name: pylibraft
-  wheel-build-raft-dask:
-    needs: wheel-publish-pylibraft
+      script: ci/build_wheel_python.sh
+  wheel-publish-cpp:
+    needs: wheel-build-cpp
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-build[email protected]
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-publish[email protected]
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
       sha: ${{ inputs.sha }}
       date: ${{ inputs.date }}
-      script: ci/build_wheel_raft_dask.sh
-  wheel-publish-raft-dask:
-    needs: wheel-build-raft-dask
+      package-name: raft
+      package-type: cpp
+  wheel-publish-python:
+    needs: wheel-build-python
     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: raft_dask
+      package-name: raft
+      package-type: python

.github/workflows/pr.yaml

@@ -19,9 +19,9 @@ jobs:
       - conda-python-build
       - conda-python-tests
       - docs-build
-      - wheel-build-pylibraft
+      - wheel-build-cpp
+      - wheel-build-python
       - wheel-tests-pylibraft
-      - wheel-build-raft-dask
       - wheel-tests-raft-dask
       - devcontainer
     secrets: inherit
@@ -74,29 +74,30 @@ jobs:
       arch: "amd64"
       container_image: "rapidsai/ci-conda:latest"
       run_script: "ci/build_docs.sh"
-  wheel-build-pylibraft:
+  wheel-build-cpp:
     needs: checks
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]
     with:
+      matrix_filter: group_by([.ARCH, (.CUDA_VER|split(".")|map(tonumber)|.[0])]) | map(max_by(.PY_VER|split(".")|map(tonumber)))
       build_type: pull-request
-      script: ci/build_wheel_pylibraft.sh
-  wheel-tests-pylibraft:
-    needs: wheel-build-pylibraft
+      script: ci/build_wheel_cpp.sh
+  wheel-build-python:
+    needs: wheel-build-cpp
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-test[email protected]
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-build[email protected]
     with:
       build_type: pull-request
-      script: ci/test_wheel_pylibraft.sh
-  wheel-build-raft-dask:
-    needs: wheel-tests-pylibraft
+      script: ci/build_wheel_python.sh
+  wheel-tests-pylibraft:
+    needs: wheel-build-python
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-build[email protected]
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-test[email protected]
     with:
       build_type: pull-request
-      script: "ci/build_wheel_raft_dask.sh"
+      script: ci/test_wheel_pylibraft.sh
   wheel-tests-raft-dask:
-    needs: wheel-build-raft-dask
+    needs: wheel-build-python
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]
     with:

build.sh

@@ -381,14 +381,6 @@ 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')
-
 # If clean given, run it prior to any other steps
 if (( ${CLEAN} == 1 )); then
     # If the dirs to clean are mounted dirs in a container, the
@@ -495,13 +487,13 @@ fi
 
 # Build and (optionally) install the pylibraft Python package
 if (( ${NUMARGS} == 0 )) || hasArg pylibraft; then
-    SKBUILD_CMAKE_ARGS="${SKBUILD_EXTRA_CMAKE_ARGS}" \
+    CMAKE_ARGS="${EXTRA_CMAKE_ARGS}" \
         python -m pip install --no-build-isolation --no-deps ${REPODIR}/python/pylibraft
 fi
 
 # Build and (optionally) install the raft-dask Python package
 if (( ${NUMARGS} == 0 )) || hasArg raft-dask; then
-    SKBUILD_CMAKE_ARGS="${SKBUILD_EXTRA_CMAKE_ARGS}" \
+    CMAKE_ARGS="${EXTRA_CMAKE_ARGS}" \
         python -m pip install --no-build-isolation --no-deps ${REPODIR}/python/raft-dask
 fi
 

ci/build_wheel_cpp.sh

@@ -0,0 +1,46 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+package_name="libraft"
+package_dir="python/libraft"
+
+source rapids-configure-sccache
+source rapids-date-string
+
+version=$(rapids-generate-version)
+git_commit=$(git rev-parse HEAD)
+
+RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
+
+# everywhere except in the final wheel name.
+PACKAGE_CUDA_SUFFIX="-${RAPIDS_PY_CUDA_SUFFIX}"
+
+# Patch project metadata files to include the CUDA version suffix and version override.
+pyproject_file="${package_dir}/pyproject.toml"
+version_file="${package_dir}/${package_name}/_version.py"
+
+sed -i "s/name = \"${package_name}\"/name = \"${package_name}${PACKAGE_CUDA_SUFFIX}\"/g" ${pyproject_file}
+echo "${version}" > VERSION
+sed -i "/^__git_commit__ / s/= .*/= \"${git_commit}\"/g" ${version_file}
+
+# For nightlies we want to ensure that we're pulling in alphas as well. The
+# easiest way to do so is to augment the spec with a constraint containing a
+# min alpha version that doesn't affect the version bounds but does allow usage
+# of alpha versions for that dependency without --pre
+alpha_spec=''
+if ! rapids-is-release-build; then
+    alpha_spec=',>=0.0.0a0'
+fi
+
+sed -r -i "s/librmm(.*)\"/librmm${PACKAGE_CUDA_SUFFIX}\1${alpha_spec}\"/g" ${pyproject_file}
+
+cd "${package_dir}"
+
+python -m pip wheel . -w dist -vvv --no-deps --disable-pip-version-check
+
+mkdir -p final_dist
+python -m auditwheel repair -w final_dist dist/*
+
+RAPIDS_PY_WHEEL_NAME="raft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 cpp final_dist

ci/build_wheel_python.sh

@@ -0,0 +1,84 @@
+#!/bin/bash
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+# Clear out system ucx files to ensure that we're getting ucx from the wheel
+# when building raft-dask.
+rm -rf /usr/lib64/ucx
+rm -rf /usr/lib64/libuc*
+
+source rapids-configure-sccache
+source rapids-date-string
+
+version=$(rapids-generate-version)
+git_commit=$(git rev-parse HEAD)
+
+RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
+
+# This is the version of the suffix with a preceding hyphen. It's used
+# everywhere except in the final wheel name.
+PACKAGE_CUDA_SUFFIX="-${RAPIDS_PY_CUDA_SUFFIX}"
+
+echo "${version}" > VERSION
+
+# For nightlies we want to ensure that we're pulling in alphas as well. The
+# easiest way to do so is to augment the spec with a constraint containing a
+# min alpha version that doesn't affect the version bounds but does allow usage
+# of alpha versions for that dependency without --pre
+alpha_spec=''
+if ! rapids-is-release-build; then
+    alpha_spec=',>=0.0.0a0'
+fi
+
+CPP_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="raft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 cpp /tmp/libraft_dist)
+PYTHON_WHEELHOUSE="${PWD}/dist/"
+PYTHON_AUDITED_WHEELHOUSE="${PWD}/final_dist/"
+WHEELHOUSES=("${PYTHON_WHEELHOUSE}" "${CPP_WHEELHOUSE}")
+mkdir -p "${PYTHON_AUDITED_WHEELHOUSE}"
+
+FIND_LINKS=""
+# Iterate over the array
+for wheelhouse in "${WHEELHOUSES[@]}"; do
+    FIND_LINKS+="--find-links ${wheelhouse} "
+done
+
+
+build_wheel () {
+    local package_name="${1}"
+    local underscore_package_name=$(echo "${package_name}" | tr "-" "_")
+
+    local package_dir="python/${package_name}"
+    local pyproject_file="${package_dir}/pyproject.toml"
+    local version_file="${package_dir}/${underscore_package_name}/_version.py"
+
+    sed -i "s/name = \"${package_name}\"/name = \"${package_name}${PACKAGE_CUDA_SUFFIX}\"/g" ${pyproject_file}
+    sed -i "/^__git_commit__ / s/= .*/= \"${git_commit}\"/g" ${version_file}
+
+    sed -r -i "s/libucx(.*)\"/libucx${PACKAGE_CUDA_SUFFIX}\1${alpha_spec}\"/g" ${pyproject_file}
+
+    for dep in rmm libraft pylibraft ucx-py distributed-ucxx; do
+        sed -r -i "s/${dep}==(.*)\"/${dep}${PACKAGE_CUDA_SUFFIX}==\1${alpha_spec}\"/g" ${pyproject_file}
+    done
+
+    # dask-cuda & rapids-dask-dependency don't get a suffix, but they do get an alpha spec.
+    for dep in dask-cuda rapids-dask-dependency; do
+        sed -r -i "s/${dep}==(.*)\"/${dep}==\1${alpha_spec}\"/g" ${pyproject_file}
+    done
+
+    if [[ $PACKAGE_CUDA_SUFFIX == "-cu12" ]]; then
+        sed -i "s/cuda-python[<=>\.,0-9a]*/cuda-python>=12.0,<13.0a0/g" ${pyproject_file}
+        sed -i "s/cupy-cuda11x/cupy-cuda12x/g" ${pyproject_file}
+    fi
+
+    pushd "${package_dir}"
+
+    python -m pip wheel . -w "${PYTHON_WHEELHOUSE}" -vvv --no-deps --disable-pip-version-check ${FIND_LINKS}
+    popd
+}
+
+build_wheel pylibraft
+build_wheel raft-dask
+
+python -m auditwheel repair -w "${PYTHON_AUDITED_WHEELHOUSE}" --exclude libraft.so --exclude "libucp.so.0" "${PYTHON_WHEELHOUSE}"/*
+RAPIDS_PY_WHEEL_NAME="raft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 python "${PYTHON_AUDITED_WHEELHOUSE}"

ci/test_wheel_pylibraft.sh

@@ -3,11 +3,12 @@
 
 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
 
-# echo to expand wildcard before adding `[extra]` requires for pip
-python -m pip install $(echo ./dist/pylibraft*.whl)[test]
+WHEELHOUSE="${PWD}/dist/"
+RAPIDS_PY_WHEEL_NAME="raft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 cpp "${WHEELHOUSE}"
+RAPIDS_PY_WHEEL_NAME="raft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 python "${WHEELHOUSE}"
 
-python -m pytest ./python/pylibraft/pylibraft/test
+python -m pip install "pylibraft-${RAPIDS_PY_CUDA_SUFFIX}[test]>=0.0.0a0" --find-links "${WHEELHOUSE}"
+
+python -m pytest python/pylibraft/pylibraft/test