diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh index 1a871c82d2..e3ea712dd8 100755 --- a/ci/build_wheel.sh +++ b/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" diff --git a/ci/test_wheel_raft_dask.sh b/ci/test_wheel_raft_dask.sh index e38b278d05..0e2350e8d8 100755 --- a/ci/test_wheel_raft_dask.sh +++ b/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 diff --git a/dependencies.yaml b/dependencies.yaml index b00fa3a9b1..1389454b3f 100644 --- a/dependencies.yaml +++ b/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: diff --git a/python/raft-dask/raft_dask/common/CMakeLists.txt b/python/raft-dask/raft_dask/common/CMakeLists.txt index 1279d5d501..0ef7734cd6 100644 --- a/python/raft-dask/raft_dask/common/CMakeLists.txt +++ b/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()