Add UCXX support

ci/build_wheel.sh

@@ -1,12 +1,16 @@
 #!/bin/bash
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 
 set -euo pipefail
 
 package_name=$1
 package_dir=$2
 underscore_package_name=$(echo "${package_name}" | tr "-" "_")
 
+# Clear out system ucx files to ensure that we're getting ucx from the wheel.
+rm -rf /usr/lib64/ucx
+rm -rf /usr/lib64/libuc*
+
 source rapids-configure-sccache
 source rapids-date-string
 
@@ -38,9 +42,11 @@ fi
 
 if [[ ${package_name} == "raft-dask" ]]; then
     sed -r -i "s/pylibraft==(.*)\"/pylibraft${PACKAGE_CUDA_SUFFIX}==\1${alpha_spec}\"/g" ${pyproject_file}
+    sed -r -i "s/libucx(.*)\"/libucx${PACKAGE_CUDA_SUFFIX}\1${alpha_spec}\"/g" ${pyproject_file}
     sed -r -i "s/ucx-py==(.*)\"/ucx-py${PACKAGE_CUDA_SUFFIX}==\1${alpha_spec}\"/g" ${pyproject_file}
     sed -r -i "s/rapids-dask-dependency==(.*)\"/rapids-dask-dependency==\1${alpha_spec}\"/g" ${pyproject_file}
     sed -r -i "s/dask-cuda==(.*)\"/dask-cuda==\1${alpha_spec}\"/g" ${pyproject_file}
+    sed -r -i "s/distributed-ucxx==(.*)\"/distributed-ucxx${PACKAGE_CUDA_SUFFIX}==\1${alpha_spec}\"/g" ${pyproject_file}
 else
     sed -r -i "s/rmm(.*)\"/rmm${PACKAGE_CUDA_SUFFIX}\1${alpha_spec}\"/g" ${pyproject_file}
 fi
@@ -56,6 +62,6 @@ cd "${package_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 dist/*
+python -m auditwheel repair -w final_dist --exclude "libucp.so.0" dist/*
 
 RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 final_dist

ci/release/update-version.sh

@@ -37,6 +37,8 @@ function sed_runner() {
 }
 
 sed_runner "s/set(RAPIDS_VERSION .*)/set(RAPIDS_VERSION \"${NEXT_SHORT_TAG}\")/g" cpp/template/cmake/thirdparty/fetch_rapids.cmake
+sed_runner 's/'"find_and_configure_ucxx(VERSION .*"'/'"find_and_configure_ucxx(VERSION  ${NEXT_UCX_PY_SHORT_TAG_PEP440}"'/g' python/raft-dask/cmake/thirdparty/get_ucxx.cmake
+sed_runner 's/'"branch-.*"'/'"branch-${NEXT_UCX_PY_SHORT_TAG_PEP440}"'/g' python/raft-dask/cmake/thirdparty/get_ucxx.cmake
 
 # Centralized version file update
 echo "${NEXT_FULL_TAG}" > VERSION
@@ -50,7 +52,7 @@ DEPENDENCIES=(
   rmm-cu11
   rmm-cu12
   rapids-dask-dependency
-  # ucx-py is handled separately below
+  # ucx-py and ucxx are handled separately below
 )
 for FILE in dependencies.yaml conda/environments/*.yaml; do
   for DEP in "${DEPENDENCIES[@]}"; do
@@ -59,6 +61,10 @@ for FILE in dependencies.yaml conda/environments/*.yaml; do
   sed_runner "/-.* ucx-py==/ s/==.*/==${NEXT_UCX_PY_SHORT_TAG_PEP440}\.*/g" ${FILE};
   sed_runner "/-.* ucx-py-cu11==/ s/==.*/==${NEXT_UCX_PY_SHORT_TAG_PEP440}\.*/g" ${FILE};
   sed_runner "/-.* ucx-py-cu12==/ s/==.*/==${NEXT_UCX_PY_SHORT_TAG_PEP440}\.*/g" ${FILE};
+  sed_runner "/-.* libucxx==/ s/==.*/==${NEXT_UCX_PY_SHORT_TAG_PEP440}\.*/g" ${FILE};
+  sed_runner "/-.* distributed-ucxx==/ s/==.*/==${NEXT_UCX_PY_SHORT_TAG_PEP440}\.*/g" ${FILE};
+  sed_runner "/-.* distributed-ucxx-cu11==/ s/==.*/==${NEXT_UCX_PY_SHORT_TAG_PEP440}\.*/g" ${FILE};
+  sed_runner "/-.* distributed-ucxx-cu12==/ s/==.*/==${NEXT_UCX_PY_SHORT_TAG_PEP440}\.*/g" ${FILE};
 done
 for FILE in python/*/pyproject.toml; do
   for DEP in "${DEPENDENCIES[@]}"; do
@@ -68,6 +74,7 @@ for FILE in python/*/pyproject.toml; do
 done
 
 sed_runner "/^ucx_py_version:$/ {n;s/.*/  - \"${NEXT_UCX_PY_VERSION}\"/}" conda/recipes/raft-dask/conda_build_config.yaml
+sed_runner "/^ucxx_version:$/ {n;s/.*/  - \"${NEXT_UCX_PY_VERSION}\"/}" conda/recipes/raft-dask/conda_build_config.yaml
 
 for FILE in .github/workflows/*.yaml; do
   sed_runner "/shared-workflows/ s/@.*/@branch-${NEXT_SHORT_TAG}/g" "${FILE}"

ci/test_python.sh

@@ -59,5 +59,23 @@ rapids-logger "pytest raft-dask"
   --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/raft-dask-coverage.xml" \
   --cov-report=term
 
+rapids-logger "pytest raft-dask (ucx-py only)"
+./ci/run_raft_dask_pytests.sh \
+  --junitxml="${RAPIDS_TESTS_DIR}/junit-raft-dask-ucx.xml" \
+  --cov-config=../.coveragerc \
+  --cov=raft_dask \
+  --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/raft-dask-ucx-coverage.xml" \
+  --cov-report=term \
+  --run_ucx
+
+rapids-logger "pytest raft-dask (ucxx only)"
+./ci/run_raft_dask_pytests.sh \
+  --junitxml="${RAPIDS_TESTS_DIR}/junit-raft-dask-ucxx.xml" \
+  --cov-config=../.coveragerc \
+  --cov=raft_dask \
+  --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/raft-dask-ucxx-coverage.xml" \
+  --cov-report=term \
+  --run_ucxx
+
 rapids-logger "Test script exiting with value: $EXITCODE"
 exit ${EXITCODE}

ci/test_wheel_raft_dask.sh

@@ -11,7 +11,13 @@ RAPIDS_PY_WHEEL_NAME="raft_dask_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels
 RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./local-pylibraft-dep
 python -m pip install --no-deps ./local-pylibraft-dep/pylibraft*.whl
 
-# echo to expand wildcard before adding `[extra]` requires for pip
-python -m pip install $(echo ./dist/raft_dask*.whl)[test]
+python -m pip install "raft_dask-${RAPIDS_PY_CUDA_SUFFIX}[test]>=0.0.0a0" --find-links dist/
 
-python -m pytest ./python/raft-dask/raft_dask/test
+# rapids-logger "pytest raft-dask"
+# python -m pytest ./python/raft-dask/raft_dask/test
+
+# rapids-logger "pytest raft-dask (ucx-py only)"
+# python -m pytest ./python/raft-dask/raft_dask/test --run_ucx
+
+rapids-logger "pytest raft-dask (ucxx only)"
+python -m pytest ./python/raft-dask/raft_dask/test --run_ucxx

conda/environments/all_cuda-118_arch-aarch64.yaml

@@ -21,6 +21,7 @@ dependencies:
 - cxx-compiler
 - cython>=3.0.0
 - dask-cuda==24.6.*
+- distributed-ucxx==0.38.*
 - doxygen>=1.8.20
 - gcc_linux-aarch64=11.*
 - graphviz
@@ -34,6 +35,7 @@ dependencies:
 - libcusolver=11.4.1.48
 - libcusparse-dev=11.7.5.86
 - libcusparse=11.7.5.86
+- libucxx==0.38.*
 - nccl>=2.9.9
 - ninja
 - numba>=0.57

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -21,6 +21,7 @@ dependencies:
 - cxx-compiler
 - cython>=3.0.0
 - dask-cuda==24.6.*
+- distributed-ucxx==0.38.*
 - doxygen>=1.8.20
 - gcc_linux-64=11.*
 - graphviz
@@ -34,6 +35,7 @@ dependencies:
 - libcusolver=11.4.1.48
 - libcusparse-dev=11.7.5.86
 - libcusparse=11.7.5.86
+- libucxx==0.38.*
 - nccl>=2.9.9
 - ninja
 - numba>=0.57

conda/environments/all_cuda-122_arch-aarch64.yaml

@@ -22,6 +22,7 @@ dependencies:
 - cxx-compiler
 - cython>=3.0.0
 - dask-cuda==24.6.*
+- distributed-ucxx==0.38.*
 - doxygen>=1.8.20
 - gcc_linux-aarch64=11.*
 - graphviz
@@ -31,6 +32,7 @@ dependencies:
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
+- libucxx==0.38.*
 - nccl>=2.9.9
 - ninja
 - numba>=0.57

conda/environments/all_cuda-122_arch-x86_64.yaml

@@ -22,6 +22,7 @@ dependencies:
 - cxx-compiler
 - cython>=3.0.0
 - dask-cuda==24.6.*
+- distributed-ucxx==0.38.*
 - doxygen>=1.8.20
 - gcc_linux-64=11.*
 - graphviz
@@ -31,6 +32,7 @@ dependencies:
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
+- libucxx==0.38.*
 - nccl>=2.9.9
 - ninja
 - numba>=0.57

conda/environments/bench_ann_cuda-118_arch-aarch64.yaml

@@ -30,6 +30,7 @@ dependencies:
 - libcusolver=11.4.1.48
 - libcusparse-dev=11.7.5.86
 - libcusparse=11.7.5.86
+- libucxx==0.38.*
 - matplotlib
 - nccl>=2.9.9
 - ninja

conda/environments/bench_ann_cuda-118_arch-x86_64.yaml

@@ -30,6 +30,7 @@ dependencies:
 - libcusolver=11.4.1.48
 - libcusparse-dev=11.7.5.86
 - libcusparse=11.7.5.86
+- libucxx==0.38.*
 - matplotlib
 - nccl>=2.9.9
 - ninja

conda/environments/bench_ann_cuda-120_arch-aarch64.yaml

@@ -27,6 +27,7 @@ dependencies:
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
+- libucxx==0.38.*
 - matplotlib
 - nccl>=2.9.9
 - ninja

conda/environments/bench_ann_cuda-120_arch-x86_64.yaml

@@ -27,6 +27,7 @@ dependencies:
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
+- libucxx==0.38.*
 - matplotlib
 - nccl>=2.9.9
 - ninja

conda/recipes/raft-dask/conda_build_config.yaml

@@ -16,11 +16,11 @@ c_stdlib:
 c_stdlib_version:
   - "2.17"
 
-ucx_version:
-  - ">=1.15.0,<1.16.0"
-
 ucx_py_version:
   - "0.38.*"
 
+ucxx_version:
+  - "0.38.*"
+
 cmake_version:
   - ">=3.26.4"

conda/recipes/raft-dask/meta.yaml

@@ -56,9 +56,9 @@ requirements:
     - rmm ={{ minor_version }}
     - scikit-build-core >=0.7.0
     - setuptools
-    - ucx {{ ucx_version }}
-    - ucx-proc=*=gpu
     - ucx-py {{ ucx_py_version }}
+    - libucxx {{ ucxx_version }}
+    - ucxx {{ ucxx_version }}
   run:
     {% if cuda_major == "11" %}
     - cudatoolkit
@@ -73,9 +73,8 @@ requirements:
     - pylibraft {{ version }}
     - python x.x
     - rmm ={{ minor_version }}
-    - ucx {{ ucx_version }}
-    - ucx-proc=*=gpu
     - ucx-py {{ ucx_py_version }}
+    - ucxx {{ ucxx_version }}
 
 tests:
   requirements:

cpp/CMakeLists.txt

@@ -650,12 +650,21 @@ rapids_find_generate_module(
   INSTALL_EXPORT_SET raft-distributed-exports
 )
 
-rapids_export_package(BUILD ucx raft-distributed-exports)
-rapids_export_package(INSTALL ucx raft-distributed-exports)
+rapids_export_package(
+  BUILD ucxx raft-distributed-exports COMPONENTS ucxx python GLOBAL_TARGETS ucxx::ucxx ucxx::python
+)
+rapids_export_package(
+  INSTALL ucxx raft-distributed-exports COMPONENTS ucxx python GLOBAL_TARGETS ucxx::ucxx
+                                                                              ucxx::python
+)
 rapids_export_package(BUILD NCCL raft-distributed-exports)
 rapids_export_package(INSTALL NCCL raft-distributed-exports)
 
-target_link_libraries(raft_distributed INTERFACE ucx::ucp NCCL::NCCL)
+# ucx is a requirement for raft_distributed, but its config is not safe to be found multiple times,
+# so rather than exporting a package dependency on it above we rely on consumers to find it
+# themselves. Once https://github.com/rapidsai/ucxx/issues/173 is resolved we can export it above
+# again.
+target_link_libraries(raft_distributed INTERFACE ucx::ucp ucxx::ucxx NCCL::NCCL)
 
 # ##################################################################################################
 # * install targets-----------------------------------------------------------