use dynamic CUDA wheels on CUDA 11 (#2548)

Contributes to rapidsai/build-planning#137

Follow-up to #2531 .

See the linked issue for many more details, but in short... using a dynamically-loaded libraft which has statically-linked cuBLAS causes issues for other libraries.

There are now aarch64 CUDA 11 wheels for cuBLAS and other CUDA libraries, so it's possible to have RAFT wheels dynamically link against them. This PR does that.

## Notes for Reviewers

This has other side benefits in addition to fixing runtime issues... it also simplifies the wheel-building scripts and CMake, and makes CUDA 11 wheels noticeably smaller 😊

Authors:
  - James Lamb (https://github.com/jameslamb)

Approvers:
  - Bradley Dice (https://github.com/bdice)

URL: #2548

ci/build_wheel.sh

@@ -21,24 +21,15 @@ rapids-generate-version > ./VERSION
 
 cd "${package_dir}"
 
-case "${RAPIDS_CUDA_VERSION}" in
-  12.*)
-    EXCLUDE_ARGS=(
-      --exclude "libcublas.so.12"
-      --exclude "libcublasLt.so.12"
-      --exclude "libcurand.so.10"
-      --exclude "libcusolver.so.11"
-      --exclude "libcusparse.so.12"
-      --exclude "libnvJitLink.so.12"
-      --exclude "libucp.so.0"
-    )
-  ;;
-  11.*)
-    EXCLUDE_ARGS=(
-      --exclude "libucp.so.0"
-    )
-  ;;
-esac
+EXCLUDE_ARGS=(
+  --exclude "libcublas.so.*"
+  --exclude "libcublasLt.so.*"
+  --exclude "libcurand.so.*"
+  --exclude "libcusolver.so.*"
+  --exclude "libcusparse.so.*"
+  --exclude "libnvJitLink.so.*"
+  --exclude "libucp.so.*"
+)
 
 if [[ ${package_name} != "libraft" ]]; then
     EXCLUDE_ARGS+=(

ci/build_wheel_libraft.sh

@@ -28,16 +28,5 @@ 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

ci/validate_wheel.sh

@@ -9,27 +9,12 @@ package_name=$3
 
 RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}"
 
-# some packages are much larger on CUDA 11 than on CUDA 12
-PYDISTCHECK_ARGS=()
-if [[ "${package_name}" == "libraft" ]]; then
-    if [[ "${RAPIDS_CUDA_MAJOR}" == "11" ]]; then
-        PYDISTCHECK_ARGS+=(
-            --max-allowed-size-compressed '750M'
-        )
-    else
-        PYDISTCHECK_ARGS+=(
-            --max-allowed-size-compressed '100M'
-        )
-    fi
-fi
-
 cd "${package_dir}"
 
 rapids-logger "validate packages with 'pydistcheck'"
 
 pydistcheck \
     --inspect \
-    "${PYDISTCHECK_ARGS[@]}" \
     "$(echo ${wheel_dir_relative_path}/*.whl)"
 
 rapids-logger "validate packages with 'twine'"

dependencies.yaml

@@ -358,11 +358,14 @@ dependencies:
               - nvidia-curand-cu12
               - nvidia-cusolver-cu12
               - nvidia-cusparse-cu12
-          # CUDA 11 does not provide wheels, so use the system libraries instead
           - matrix:
               cuda: "11.*"
               use_cuda_wheels: "true"
             packages:
+              - nvidia-cublas-cu11
+              - nvidia-curand-cu11
+              - nvidia-cusolver-cu11
+              - nvidia-cusparse-cu11
           # if use_cuda_wheels=false is provided, do not add dependencies on any CUDA wheels
           # (e.g. for DLFW and pip devcontainers)
           - matrix:

python/libraft/CMakeLists.txt

@@ -22,8 +22,6 @@ project(
   LANGUAGES CXX
 )
 
-option(USE_CUDA_MATH_WHEELS "Use the CUDA math wheels instead of the system libraries" OFF)
-
 # Check if raft is already available. If so, it is the user's responsibility to ensure that the
 # CMake package is also available at build time of the Python raft package.
 find_package(raft "${RAPIDS_VERSION}")
@@ -35,14 +33,8 @@ endif()
 unset(raft_FOUND)
 
 # --- CUDA --- #
-find_package(CUDAToolkit REQUIRED)
 set(CUDA_STATIC_RUNTIME ON)
-set(CUDA_STATIC_MATH_LIBRARIES ON)
-if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.0)
-  set(CUDA_STATIC_MATH_LIBRARIES OFF)
-elseif(USE_CUDA_MATH_WHEELS)
-  message(FATAL_ERROR "Cannot use CUDA math wheels with CUDA < 12.0")
-endif()
+set(CUDA_STATIC_MATH_LIBRARIES OFF)
 
 # --- RAFT ---#
 set(BUILD_TESTS OFF)
@@ -52,14 +44,13 @@ set(RAFT_COMPILE_LIBRARY ON)
 
 add_subdirectory(../../cpp raft-cpp)
 
-if(NOT CUDA_STATIC_MATH_LIBRARIES AND USE_CUDA_MATH_WHEELS)
-  set_property(
-    TARGET raft_lib
-    PROPERTY INSTALL_RPATH
-             "$ORIGIN/../nvidia/cublas/lib"
-             "$ORIGIN/../nvidia/curand/lib"
-             "$ORIGIN/../nvidia/cusolver/lib"
-             "$ORIGIN/../nvidia/cusparse/lib"
-             "$ORIGIN/../nvidia/nvjitlink/lib"
-  )
-endif()
+# assumes libraft.so is installed 2 levels deep, e.g. site-packages/libraft/lib64/libraft.so
+set_property(
+  TARGET raft_lib
+  PROPERTY INSTALL_RPATH
+           "$ORIGIN/../../nvidia/cublas/lib"
+           "$ORIGIN/../../nvidia/curand/lib"
+           "$ORIGIN/../../nvidia/cusolver/lib"
+           "$ORIGIN/../../nvidia/cusparse/lib"
+           "$ORIGIN/../../nvidia/nvjitlink/lib"
+)

python/libraft/pyproject.toml

@@ -110,6 +110,8 @@ matrix-entry = "cuda_suffixed=true;use_cuda_wheels=true"
 
 [tool.pydistcheck]
 select = [
-    # NOTE: size threshold is managed via CLI args in CI scripts
     "distro-too-large-compressed",
 ]
+
+# PyPI limit is 100 MiB, fail CI before we get too close to that
+max_allowed_size_compressed = '75M'