use CUDA wheels on CUDA 11

ci/build_wheel_cuml.sh

@@ -6,60 +6,31 @@ set -euo pipefail
 package_name="cuml"
 package_dir="python/cuml"
 
-# TODO(jameslamb): remove this when https://github.com/rapidsai/raft/pull/2531 is merged
-source ./ci/use_wheels_from_prs.sh
-
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
 
 # Download the libcuml wheel built in the previous step and make it
 # available for pip to find.
 RAPIDS_PY_WHEEL_NAME="libcuml_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 cpp /tmp/libcuml_dist
+echo "libcuml-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo /tmp/libcuml_dist/libcuml_*.whl)" > /tmp/constraints.txt
+export PIP_CONSTRAINT="/tmp/constraints.txt"
 
-rapids-logger "Generating build requirements"
-
-rapids-dependency-file-generator \
-  --output requirements \
-  --file-key "py_build_${package_name}" \
-  --file-key "py_rapids_build_${package_name}" \
-  --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION};cuda_suffixed=true" \
-| tee /tmp/requirements-build.txt
-
-rapids-logger "Installing build requirements"
-python -m pip install \
-    -v \
-    --prefer-binary \
-    -r /tmp/requirements-build.txt \
-    /tmp/libcuml_dist/libcuml_*.whl
-
-# 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
+# TODO(jameslamb): remove this when https://github.com/rapidsai/raft/pull/2531 is merged
+source ./ci/use_wheels_from_prs.sh
 
 EXCLUDE_ARGS=(
   --exclude "libcuml++.so"
   --exclude "libcumlprims_mg.so"
   --exclude "libcuvs.so"
+  --exclude "libcublas.so.*"
+  --exclude "libcublasLt.so.*"
+  --exclude "libcufft.so.*"
+  --exclude "libcurand.so.*"
+  --exclude "libcusolver.so.*"
+  --exclude "libcusparse.so.*"
+  --exclude "libnvJitLink.so.*"
 )
 
-case "${RAPIDS_CUDA_VERSION}" in
-  12.*)
-    EXCLUDE_ARGS+=(
-      --exclude "libcublas.so.12"
-      --exclude "libcublasLt.so.12"
-      --exclude "libcufft.so.11"
-      --exclude "libcurand.so.10"
-      --exclude "libcusolver.so.11"
-      --exclude "libcusparse.so.12"
-      --exclude "libnvJitLink.so.12"
-    )
-    EXTRA_CMAKE_ARGS="-DUSE_CUDA_MATH_WHEELS=ON"
-    ;;
-  11.*)
-    EXTRA_CMAKE_ARGS="-DUSE_CUDA_MATH_WHEELS=OFF"
-    ;;
-esac
-
-export SKBUILD_CMAKE_ARGS="-DDISABLE_DEPRECATION_WARNINGS=ON;-DCPM_cumlprims_mg_SOURCE=${GITHUB_WORKSPACE}/cumlprims_mg/;${EXTRA_CMAKE_ARGS};-DSINGLEGPU=OFF;-DUSE_LIBCUML_WHEEL=ON"
+export SKBUILD_CMAKE_ARGS="-DDISABLE_DEPRECATION_WARNINGS=ON;-DCPM_cumlprims_mg_SOURCE=${GITHUB_WORKSPACE}/cumlprims_mg/;-DSINGLEGPU=OFF;-DUSE_LIBCUML_WHEEL=ON"
 ./ci/build_wheel.sh "${package_name}" "${package_dir}"
 
 mkdir -p ${package_dir}/final_dist

ci/build_wheel_libcuml.sh

@@ -36,34 +36,16 @@ EXCLUDE_ARGS=(
   --exclude "libcumlprims_mg.so"
   --exclude "libcuvs.so"
   --exclude "libraft.so"
+  --exclude "libcublas.so.*"
+  --exclude "libcublasLt.so.*"
+  --exclude "libcufft.so.*"
+  --exclude "libcurand.so.*"
+  --exclude "libcusolver.so.*"
+  --exclude "libcusparse.so.*"
+  --exclude "libnvJitLink.so.*"
 )
 
-# Avoid ever vendoring CUDA libraries into wheels.
-#
-# On CUDA 11 builds these excludes should technically be unnecessary because
-# there cuml and its dependencies statically link again these libraries, but
-# this is here unconditionally to ensure those wheels don't accidentally pick up
-# these libraries transitively.
-EXCLUDE_ARGS+=(
-  --exclude "libcublas.so.12"
-  --exclude "libcublasLt.so.12"
-  --exclude "libcufft.so.11"
-  --exclude "libcurand.so.10"
-  --exclude "libcusolver.so.11"
-  --exclude "libcusparse.so.12"
-  --exclude "libnvJitLink.so.12"
-)
-
-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="-DDISABLE_DEPRECATION_WARNINGS=ON;-DCPM_cumlprims_mg_SOURCE=${GITHUB_WORKSPACE}/cumlprims_mg/;${EXTRA_CMAKE_ARGS}"
+export SKBUILD_CMAKE_ARGS="-DDISABLE_DEPRECATION_WARNINGS=ON;-DCPM_cumlprims_mg_SOURCE=${GITHUB_WORKSPACE}/cumlprims_mg/"
 ./ci/build_wheel.sh "${package_name}" "${package_dir}"
 
 mkdir -p ${package_dir}/final_dist

ci/use_wheels_from_prs.sh

@@ -11,7 +11,7 @@ LIBCUVS_CHANNEL=$(
   RAPIDS_PY_WHEEL_NAME="libcuvs_${RAPIDS_PY_CUDA_SUFFIX}" rapids-get-pr-wheel-artifact cuvs 594 cpp "${CUVS_COMMIT:0:7}"
 )
 
-cat > ./constraints.txt <<EOF
+cat >> ./constraints.txt <<EOF
 cuvs-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo ${CUVS_CHANNEL}/cuvs_*.whl)
 libcuvs-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo ${LIBCUVS_CHANNEL}/libcuvs_*.whl)
 EOF

ci/validate_wheel.sh

@@ -10,6 +10,7 @@ RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}"
 
 # some packages are much larger on CUDA 11 than on CUDA 12
 PYDISTCHECK_ARGS=()
+# TODO(jameslamb): update these thresholds
 if [[ "${package_dir}" == "python/libcuml" ]]; then
     if [[ "${RAPIDS_CUDA_MAJOR}" == "11" ]]; then
         PYDISTCHECK_ARGS+=(

dependencies.yaml

@@ -400,6 +400,11 @@ dependencies:
               cuda: "11.*"
               use_cuda_wheels: "true"
             packages:
+              - nvidia-cublas-cu11
+              - nvidia-cufft-cu11
+              - nvidia-curand-cu11
+              - nvidia-cusparse-cu11
+              - nvidia-cusolver-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/cuml/CMakeLists.txt

@@ -37,13 +37,11 @@ project(
 # - User Options  --------------------------------------------------------------
 option(CUML_UNIVERSAL "Build all cuML Python components." ON)
 option(SINGLEGPU "Disable all mnmg components and comms libraries" OFF)
-option(USE_CUDA_MATH_WHEELS "Use the CUDA math wheels instead of the system libraries" OFF)
 option(USE_LIBCUML_WHEEL "Use libcuml wheel to provide some dependencies" OFF)
 
 # todo: use CMAKE_MESSAGE_CONTEXT for prefix for logging.
 # https://github.com/rapidsai/cuml/issues/4843
 message(VERBOSE "CUML_PY: Build only cuML CPU Python components.: ${CUML_CPU}")
-message(VERBOSE "CUML_PY: Use CUDA math wheels instead of system libraries: ${USE_CUDA_MATH_WHEELS}")
 message(VERBOSE "CUML_PY: Disabling all mnmg components and comms libraries: ${SINGLEGPU}")
 
 set(CUML_ALGORITHMS "ALL" CACHE STRING "Choose which algorithms are built cuML. Can specify individual algorithms or groups in a semicolon-separated list.")

python/libcuml/CMakeLists.txt

@@ -25,13 +25,6 @@ project(
   LANGUAGES CXX CUDA
 )
 
-################################################################################
-# - User Options  --------------------------------------------------------------
-option(USE_CUDA_MATH_WHEELS "Use the CUDA math wheels instead of the system libraries" OFF)
-
-################################################################################
-# - Process User Options  ------------------------------------------------------
-
 # Check if cuml 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 cuml package.
 find_package(cuml "${RAPIDS_VERSION}")
@@ -80,29 +73,17 @@ set(CUML_EXCLUDE_TREELITE_FROM_ALL ON)
 
 # --- CUDA --- #
 set(CUDA_STATIC_RUNTIME ON)
-
-# Link to the CUDA wheels with shared libraries for CUDA 12+
-find_package(CUDAToolkit REQUIRED)
-if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.0)
-  set(CUDA_STATIC_MATH_LIBRARIES OFF)
-else()
-  if(USE_CUDA_MATH_WHEELS)
-    message(FATAL_ERROR "Cannot use CUDA math wheels with CUDA < 12.0")
-  endif()
-  set(CUDA_STATIC_MATH_LIBRARIES ON)
-endif()
+set(CUDA_STATIC_MATH_LIBRARIES OFF)
 
 add_subdirectory(../../cpp cuml-cpp)
 
-if(NOT CUDA_STATIC_MATH_LIBRARIES AND USE_CUDA_MATH_WHEELS)
-  # assumes libcuml++ is installed 2 levels deep, e.g. site-packages/cuml/lib64/libcuml++.so
-  set(rpaths
-    "$ORIGIN/../../nvidia/cublas/lib"
-    "$ORIGIN/../../nvidia/cufft/lib"
-    "$ORIGIN/../../nvidia/curand/lib"
-    "$ORIGIN/../../nvidia/cusolver/lib"
-    "$ORIGIN/../../nvidia/cusparse/lib"
-    "$ORIGIN/../../nvidia/nvjitlink/lib"
-  )
-  set_property(TARGET ${CUML_CPP_TARGET} PROPERTY INSTALL_RPATH ${rpaths} APPEND)
-endif()
+# assumes libcuml++ is installed 2 levels deep, e.g. site-packages/cuml/lib64/libcuml++.so
+set(rpaths
+  "$ORIGIN/../../nvidia/cublas/lib"
+  "$ORIGIN/../../nvidia/cufft/lib"
+  "$ORIGIN/../../nvidia/curand/lib"
+  "$ORIGIN/../../nvidia/cusolver/lib"
+  "$ORIGIN/../../nvidia/cusparse/lib"
+  "$ORIGIN/../../nvidia/nvjitlink/lib"
+)
+set_property(TARGET ${CUML_CPP_TARGET} PROPERTY INSTALL_RPATH ${rpaths} APPEND)