Enable Dask tests with UCX-Py/UCXX in CI

ci/build_wheel.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 
-set -euo pipefail
+set -xeuo pipefail
 
 package_name="cuml"
 package_dir="python"
@@ -14,6 +14,8 @@ git_commit=$(git rev-parse HEAD)
 
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
 
+raft_dask_wheelhouse=$(RAPIDS_PY_WHEEL_NAME="raft_dask_${RAPIDS_PY_CUDA_SUFFIX}" rapids-get-pr-wheel-artifact raft 1983 python)
+
 # This is the version of the suffix with a preceding hyphen. It's used
 # everywhere except in the final wheel name.
 PACKAGE_CUDA_SUFFIX="-${RAPIDS_PY_CUDA_SUFFIX}"
@@ -34,7 +36,7 @@ if ! rapids-is-release-build; then
     alpha_spec=',>=0.0.0a0'
 fi
 
-for dep in cudf pylibraft raft-dask rmm; do
+for dep in cudf distributed-ucxx pylibraft raft-dask rmm; do
     sed -r -i "s/${dep}==(.*)\"/${dep}${PACKAGE_CUDA_SUFFIX}==\1${alpha_spec}\"/g" ${pyproject_file}
 done
 
@@ -54,7 +56,8 @@ SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DDISABLE_DEPRECATION_WARNINGS=ON;-DC
     -w dist \
     -vvv \
     --no-deps \
-    --disable-pip-version-check
+    --disable-pip-version-check \
+    --find-links ${raft_dask_wheelhouse}
 
 mkdir -p final_dist
 python -m auditwheel repair -w final_dist dist/*

ci/release/update-version.sh

@@ -22,10 +22,11 @@ CURRENT_SHORT_TAG=${CURRENT_MAJOR}.${CURRENT_MINOR}
 NEXT_MAJOR=$(echo $NEXT_FULL_TAG | awk '{split($0, a, "."); print a[1]}')
 NEXT_MINOR=$(echo $NEXT_FULL_TAG | awk '{split($0, a, "."); print a[2]}')
 NEXT_SHORT_TAG=${NEXT_MAJOR}.${NEXT_MINOR}
-NEXT_UCX_PY_VERSION="$(curl -sL https://version.gpuci.io/rapids/${NEXT_SHORT_TAG}).*"
+NEXT_UCX_PY_SHORT_TAG="$(curl -sL https://version.gpuci.io/rapids/${NEXT_SHORT_TAG})"
 
 # Need to distutils-normalize the original version
 NEXT_SHORT_TAG_PEP440=$(python -c "from setuptools.extern import packaging; print(packaging.version.Version('${NEXT_SHORT_TAG}'))")
+NEXT_UCX_PY_SHORT_TAG_PEP440=$(python -c "from setuptools.extern import packaging; print(packaging.version.Version('${NEXT_UCX_PY_SHORT_TAG}'))")
 
 echo "Preparing release $CURRENT_TAG => $NEXT_FULL_TAG"
 
@@ -43,6 +44,7 @@ sed_runner "s/rmm==.*\",/rmm==${NEXT_SHORT_TAG_PEP440}.*\",/g" python/pyproject.
 sed_runner "s/cudf==.*\",/cudf==${NEXT_SHORT_TAG_PEP440}.*\",/g" python/pyproject.toml
 sed_runner "s/pylibraft==.*\",/pylibraft==${NEXT_SHORT_TAG_PEP440}.*\",/g" python/pyproject.toml
 sed_runner "s/raft-dask==.*\",/raft-dask==${NEXT_SHORT_TAG_PEP440}.*\",/g" python/pyproject.toml
+sed_runner "s/distributed-ucxx==.*\",/distributed-ucxx==${NEXT_UCX_PY_SHORT_TAG_PEP440}.*\",/g" python/pyproject.toml
 
 DEPENDENCIES=(
   cudf
@@ -64,6 +66,9 @@ for FILE in dependencies.yaml conda/environments/*.yaml; do
   for DEP in "${DEPENDENCIES[@]}"; do
     sed_runner "/-.* ${DEP}\(-cu[[:digit:]]\{2\}\)\{0,1\}==/ s/==.*/==${NEXT_SHORT_TAG_PEP440}.*/g" "${FILE}"
   done
+  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
 
 sed_runner "s|/branch-[^/]*/|/branch-${NEXT_SHORT_TAG}/|g" README.md

ci/run_cuml_dask_pytests.sh

@@ -4,4 +4,14 @@
 # Support invoking run_cuml_dask_pytests.sh outside the script directory
 cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cuml/tests/dask
 
-python -m pytest --cache-clear "$@" .
+# Disable non-UCXX tests for wheels tests
+if [[ -v CONDA_PREFIX ]]; then
+    rapids-logger "pytest cuml-dask (No UCX-Py/UCXX)"
+    python -m pytest --cache-clear "$@" .
+
+    rapids-logger "pytest cuml-dask (UCX-Py only)"
+    python -m pytest --cache-clear --run_ucx "$@" .
+fi
+
+rapids-logger "pytest cuml-dask (UCXX only)"
+python -m pytest --cache-clear --run_ucxx "$@" .

ci/test_python_common.sh

@@ -8,13 +8,14 @@ set -euo pipefail
 rapids-logger "Downloading artifacts from previous jobs"
 CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
 PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
+RAFT_CHANNEL=$(rapids-get-pr-conda-artifact raft 1983 python)
 
 rapids-logger "Generate Python testing dependencies"
 rapids-dependency-file-generator \
   --output conda \
   --file_key test_python \
   --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \
-  --prepend-channels "${CPP_CHANNEL};${PYTHON_CHANNEL}" | tee env.yaml
+  --prepend-channels "${CPP_CHANNEL};${PYTHON_CHANNEL};${RAFT_CHANNEL}" | tee env.yaml
 
 rapids-mamba-retry env create --yes -f env.yaml -n test
 

ci/test_python_dask.sh

@@ -11,7 +11,7 @@ EXITCODE=0
 trap "EXITCODE=1" ERR
 set +e
 
-rapids-logger "pytest cuml-dask"
+# Run tests (no UCX-Py/UCXX)
 ./ci/run_cuml_dask_pytests.sh \
   --junitxml="${RAPIDS_TESTS_DIR}/junit-cuml-dask.xml" \
   --cov-config=../../../.coveragerc \

ci/test_wheel.sh

@@ -6,14 +6,15 @@ set -euo pipefail
 mkdir -p ./dist
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
 RAPIDS_PY_WHEEL_NAME="cuml_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist
+raft_dask_wheelhouse=$(RAPIDS_PY_WHEEL_NAME="raft_dask_${RAPIDS_PY_CUDA_SUFFIX}" rapids-get-pr-wheel-artifact raft 1983 python)
 
 # On arm also need to install CMake because treelite needs to be compiled (no wheels available for arm).
 if [[ "$(arch)" == "aarch64" ]]; then
     python -m pip install cmake
 fi
 
 # echo to expand wildcard before adding `[extra]` requires for pip
-python -m pip install $(echo ./dist/cuml*.whl)[test]
+python -m pip install $(echo ./dist/cuml*.whl)[test] --find-links dist/ --find-link "${raft_dask_wheelhouse}"
 
 RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}
 mkdir -p "${RAPIDS_TESTS_DIR}"

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

@@ -19,6 +19,7 @@ dependencies:
 - dask-cuda==24.6.*
 - dask-cudf==24.6.*
 - dask-ml
+- distributed-ucxx==0.38.*
 - doxygen=1.9.1
 - gcc_linux-64=11.*
 - graphviz

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

@@ -21,6 +21,7 @@ dependencies:
 - dask-cuda==24.6.*
 - dask-cudf==24.6.*
 - dask-ml
+- distributed-ucxx==0.38.*
 - doxygen=1.9.1
 - gcc_linux-64=11.*
 - graphviz

cpp/cmake/thirdparty/get_cumlprims_mg.cmake

@@ -1,5 +1,5 @@
 #=============================================================================
-# Copyright (c) 2021-2022, NVIDIA CORPORATION.
+# Copyright (c) 2021-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -64,7 +64,7 @@ endfunction()
 ###
 find_and_configure_cumlprims_mg(VERSION          ${CUML_MIN_VERSION_cumlprims_mg}
                                 FORK       rapidsai
-                                PINNED_TAG branch-${CUML_BRANCH_VERSION_cumlprims_mg}
+                                PINNED_TAG pull-request/186
                                 BUILD_STATIC     ${CUML_USE_CUMLPRIMS_MG_STATIC}
                                 EXCLUDE_FROM_ALL ${CUML_EXCLUDE_CUMLPRIMS_MG_FROM_ALL}
                                 # When PINNED_TAG above doesn't match cuml,

cpp/cmake/thirdparty/get_raft.cmake

@@ -83,7 +83,7 @@ endfunction()
 # CPM_raft_SOURCE=/path/to/local/raft
 find_and_configure_raft(VERSION          ${CUML_MIN_VERSION_raft}
       FORK             rapidsai
-      PINNED_TAG       branch-${CUML_BRANCH_VERSION_raft}
+      PINNED_TAG       pull-request/1983
       EXCLUDE_FROM_ALL ${CUML_EXCLUDE_RAFT_FROM_ALL}
       # When PINNED_TAG above doesn't match cuml,
       # force local raft clone in build directory

dependencies.yaml

@@ -9,6 +9,7 @@ files:
       - common_build
       - cuda
       - cuda_version
+      - depends_on_distributed_ucxx
       - docs
       - py_build
       - py_run
@@ -54,6 +55,7 @@ files:
     output: none
     includes:
       - cuda_version
+      - depends_on_distributed_ucxx
       - py_version
       - test_cuml
       - test_python
@@ -84,6 +86,7 @@ files:
       table: project.optional-dependencies
       key: test
     includes:
+      - depends_on_distributed_ucxx
       - test_python
 channels:
   - rapidsai
@@ -441,3 +444,25 @@ dependencies:
           - pandas
           - *scikit_learn
           - seaborn
+  depends_on_distributed_ucxx:
+    common:
+      - output_types: conda
+        packages:
+          # UCXX is not currently a hard-dependency thus only installed during tests,
+          # this will change in the future.
+          - &distributed_ucxx_conda distributed-ucxx==0.38.*
+      - output_types: requirements
+        packages:
+          # pip recognizes the index as a global option for the requirements.txt file
+          - --extra-index-url=https://pypi.nvidia.com
+          - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
+    specific:
+      - output_types: [requirements, pyproject]
+        matrices:
+          - matrix: {cuda: "12.*"}
+            packages:
+              - distributed-ucxx-cu12==0.38.*
+          - matrix: {cuda: "11.*"}
+            packages:
+              - distributed-ucxx-cu11==0.38.*
+          - {matrix: null, packages: [*distributed_ucxx_conda]}

python/cuml/tests/dask/conftest.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2023, NVIDIA CORPORATION.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION.
 
 import pytest
 
@@ -34,18 +34,8 @@ def client(cluster):
 
 @pytest.fixture(scope="module")
 def ucx_cluster():
-    initialize.initialize(
-        create_cuda_context=True,
-        enable_tcp_over_ucx=enable_tcp_over_ucx,
-        enable_nvlink=enable_nvlink,
-        enable_infiniband=enable_infiniband,
-    )
     cluster = LocalCUDACluster(
         protocol="ucx",
-        enable_tcp_over_ucx=enable_tcp_over_ucx,
-        enable_nvlink=enable_nvlink,
-        enable_infiniband=enable_infiniband,
-        worker_class=IncreasedCloseTimeoutNanny,
     )
     yield cluster
     cluster.close()
@@ -57,3 +47,62 @@ def ucx_client(ucx_cluster):
     client = Client(ucx_cluster)
     yield client
     client.close()
+
+
+@pytest.fixture(scope="module")
+def ucxx_cluster():
+    cluster = LocalCUDACluster(
+        protocol="ucxx",
+        worker_class=IncreasedCloseTimeoutNanny,
+    )
+    yield cluster
+    cluster.close()
+
+
+@pytest.fixture(scope="function")
+def ucxx_client(ucxx_cluster):
+    pytest.importorskip("distributed_ucxx")
+
+    client = Client(ucxx_cluster)
+    yield client
+    client.close()
+
+
+def pytest_addoption(parser):
+    group = parser.getgroup("Dask cuML Custom Options")
+
+    group.addoption(
+        "--run_ucx", action="store_true", help="run _only_ UCX-Py tests"
+    )
+
+    group.addoption(
+        "--run_ucxx", action="store_true", help="run _only_ UCXX tests"
+    )
+
+
+def pytest_collection_modifyitems(config, items):
+    if config.getoption("--run_ucx"):
+        skip_others = pytest.mark.skip(
+            reason="only runs when --run_ucx is not specified"
+        )
+        for item in items:
+            if "ucx" not in item.keywords:
+                item.add_marker(skip_others)
+    else:
+        skip_ucx = pytest.mark.skip(reason="requires --run_ucx to run")
+        for item in items:
+            if "ucx" in item.keywords:
+                item.add_marker(skip_ucx)
+
+    if config.getoption("--run_ucxx"):
+        skip_others = pytest.mark.skip(
+            reason="only runs when --run_ucxx is not specified"
+        )
+        for item in items:
+            if "ucxx" not in item.keywords:
+                item.add_marker(skip_others)
+    else:
+        skip_ucxx = pytest.mark.skip(reason="requires --run_ucxx to run")
+        for item in items:
+            if "ucxx" in item.keywords:
+                item.add_marker(skip_ucxx)