Revert "Revert "Build C++ wheel (#2264)" (#2305)"

This reverts commit ea54c79.

.github/workflows/build.yaml

@@ -67,41 +67,45 @@ jobs:
       node_type: "gpu-v100-latest-1"
       run_script: "ci/build_docs.sh"
       sha: ${{ inputs.sha }}
-  wheel-build-pylibraft:
+  wheel-build-cpp:
     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:
+      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