diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh index 62d93a668e..e7ae52f33a 100755 --- a/ci/build_wheel.sh +++ b/ci/build_wheel.sh @@ -20,10 +20,29 @@ 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 + # Hardcode the output 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 --exclude "libucp.so.0" 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 final_dist diff --git a/ci/build_wheel_pylibraft.sh b/ci/build_wheel_pylibraft.sh index 895c311f46..ce9f0ed172 100755 --- a/ci/build_wheel_pylibraft.sh +++ b/ci/build_wheel_pylibraft.sh @@ -3,7 +3,16 @@ set -euo pipefail +case "${RAPIDS_CUDA_VERSION}" in + 12.*) + EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=ON" + ;; + 11.*) + EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=OFF" + ;; +esac + # Set up skbuild options. Enable sccache in skbuild config options -export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF" +export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF${EXTRA_CMAKE_ARGS}" ci/build_wheel.sh pylibraft python/pylibraft diff --git a/dependencies.yaml b/dependencies.yaml index 630b5a50c9..92c6d98414 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -83,6 +83,7 @@ files: extras: table: project includes: + - cuda_wheels - run_pylibraft py_test_pylibraft: output: pyproject @@ -377,6 +378,36 @@ dependencies: - *libcusolver114 - *libcusparse_dev114 - *libcusparse114 + cuda_wheels: + specific: + - output_types: pyproject + matrices: + - matrix: + cuda: "12.*" + use_cuda_wheels: "true" + packages: + - nvidia-cublas-cu12 + - 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: + # if use_cuda_wheels=false is provided, do not add dependencies on any CUDA wheels + # (e.g. for DLFW and pip devcontainers) + - matrix: + use_cuda_wheels: "false" + packages: + # if no matching matrix selectors passed, list the unsuffixed packages + # (just as a source of documentation, as this populates pyproject.toml in source control) + - matrix: + packages: + - nvidia-cublas + - nvidia-curand + - nvidia-cusolver + - nvidia-cusparse depends_on_cupy: common: diff --git a/python/pylibraft/CMakeLists.txt b/python/pylibraft/CMakeLists.txt index 6cbe8e4cbf..c286d3debf 100644 --- a/python/pylibraft/CMakeLists.txt +++ b/python/pylibraft/CMakeLists.txt @@ -30,6 +30,7 @@ project( option(FIND_RAFT_CPP "Search for existing RAFT C++ installations before defaulting to local files" ON ) +option(USE_CUDA_MATH_WHEELS "Use the CUDA math wheels instead of the system libraries" OFF) # If the user requested it we attempt to find RAFT. if(FIND_RAFT_CPP) @@ -48,15 +49,32 @@ endif() include(rapids-cython-core) if(NOT raft_FOUND) + find_package(CUDAToolkit REQUIRED) + set(BUILD_TESTS OFF) set(BUILD_PRIMS_BENCH OFF) set(BUILD_ANN_BENCH OFF) set(RAFT_COMPILE_LIBRARY ON) 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() add_subdirectory(../../cpp raft-cpp EXCLUDE_FROM_ALL) + 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() + # When building the C++ libraries from source we must copy libraft.so alongside the # pairwise_distance and random Cython libraries TODO: when we have a single 'compiled' raft # library, we shouldn't need this diff --git a/python/pylibraft/pyproject.toml b/python/pylibraft/pyproject.toml index 4c1d8e8fa1..9a826e53c6 100644 --- a/python/pylibraft/pyproject.toml +++ b/python/pylibraft/pyproject.toml @@ -33,6 +33,10 @@ requires-python = ">=3.9" dependencies = [ "cuda-python", "numpy>=1.23,<2.0a0", + "nvidia-cublas", + "nvidia-curand", + "nvidia-cusolver", + "nvidia-cusparse", "rmm==24.10.*,>=0.0.0a0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. classifiers = [ @@ -124,7 +128,7 @@ requires = [ "rmm==24.10.*,>=0.0.0a0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. dependencies-file = "../../dependencies.yaml" -matrix-entry = "cuda_suffixed=true" +matrix-entry = "cuda_suffixed=true;use_cuda_wheels=true" [tool.pytest.ini_options] filterwarnings = [