Use NCCL wheels from PyPI for CUDA 12 builds

ci/build_wheel.sh

@@ -38,6 +38,14 @@ if [[ ${package_name} != "libraft" ]]; then
     )
 fi
 
+RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}"
+if [[ ${RAPIDS_CUDA_MAJOR} == "12" ]]; then
+    EXCLUDE_ARGS+=(
+      --exclude "libnccl.so.*"
+    )
+    export SKBUILD_CMAKE_ARGS="-DUSE_NCCL_RUNTIME_WHEEL=ON"
+fi
+
 sccache --zero-stats
 
 rapids-logger "Building '${package_name}' wheel"

ci/test_wheel_raft_dask.sh

@@ -3,6 +3,12 @@
 
 set -euo pipefail
 
+# Delete system libnccl.so to ensure the wheel is used
+RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}"
+if [[ ${RAPIDS_CUDA_MAJOR} == "12" ]]; then
+  rm -rf /usr/lib64/libnccl*
+fi
+
 mkdir -p ./dist
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")"
 RAPIDS_PY_WHEEL_NAME="libraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 cpp ./local-libraft-dep

dependencies.yaml

@@ -16,6 +16,7 @@ files:
       - depends_on_distributed_ucxx
       - depends_on_rmm
       - depends_on_rapids_logger
+      - depends_on_nccl
       - develop
       - docs
       - rapids_build_skbuild
@@ -78,6 +79,7 @@ files:
       - build_common
       - depends_on_librmm
       - depends_on_rapids_logger
+      - depends_on_nccl
   py_run_libraft:
     output: pyproject
     pyproject_dir: python/libraft
@@ -87,6 +89,7 @@ files:
       - cuda_wheels
       - depends_on_librmm
       - depends_on_rapids_logger
+      - depends_on_nccl
   py_build_pylibraft:
     output: pyproject
     pyproject_dir: python/pylibraft
@@ -146,6 +149,7 @@ files:
       - depends_on_libraft
       - depends_on_librmm
       - depends_on_ucx_build
+      - depends_on_nccl
   py_run_raft_dask:
     output: pyproject
     pyproject_dir: python/raft-dask
@@ -154,6 +158,7 @@ files:
     includes:
       - depends_on_distributed_ucxx
       - depends_on_libraft
+      - depends_on_nccl
       - run_raft_dask
   py_test_raft_dask:
     output: pyproject
@@ -192,7 +197,6 @@ dependencies:
           - c-compiler
           - cxx-compiler
           - libucxx==0.44.*,>=0.0.0a0
-          - nccl>=2.19
     specific:
       - output_types: conda
         matrices:
@@ -700,3 +704,18 @@ dependencies:
           - matrix: null
             packages:
               - libucx>=1.15.0
+  depends_on_nccl:
+    common:
+      - output_types: conda
+        packages:
+          - &nccl_unsuffixed nccl>=2.19
+    specific:
+      - output_types: [pyproject, requirements]
+        matrices:
+          - matrix:
+              cuda: "12.*"
+              cuda_suffixed: "true"
+            packages:
+              - nvidia-nccl-cu12>=2.19
+          - matrix:
+            packages:

python/raft-dask/raft_dask/common/CMakeLists.txt

@@ -12,8 +12,21 @@
 # the License.
 # =============================================================================
 
+option(USE_NCCL_RUNTIME_WHEEL "Use the NCCL wheel at runtime instead of the system library" OFF)
+
 set(cython_sources comms_utils.pyx nccl.pyx)
 set(linked_libraries raft::raft raft::distributed)
 rapids_cython_create_modules(
   SOURCE_FILES "${cython_sources}" LINKED_LIBRARIES "${linked_libraries}" CXX
 )
+
+if(USE_NCCL_RUNTIME_WHEEL)
+  set(rpaths "$ORIGIN/../../nvidia/nccl/lib")
+  foreach(tgt IN LISTS RAPIDS_CYTHON_CREATED_TARGETS)
+    set_property(
+      TARGET ${tgt}
+      PROPERTY INSTALL_RPATH ${rpaths}
+      APPEND
+    )
+  endforeach()
+endif()