Added cuml, raft-dask

Author: TomAugspurger

.gitignore

@@ -2,3 +2,4 @@ dask
 distributed
 cudf
 dask-cuda
+packages

scripts/install.sh

@@ -11,29 +11,59 @@ DASK_VERSION=main
 uv pip install --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple \
   --overrides=requirements/overrides.txt \
   --prerelease allow \
+  "cuml-${RAPIDS_PY_CUDA_SUFFIX}[test]" \
   "cudf-${RAPIDS_PY_CUDA_SUFFIX}" \
   "dask-cudf-${RAPIDS_PY_CUDA_SUFFIX}" \
+  "raft-dask-${RAPIDS_PY_CUDA_SUFFIX}" \
   "ucx-py-${RAPIDS_PY_CUDA_SUFFIX}" \
   "ucxx-${RAPIDS_PY_CUDA_SUFFIX}" \
   "scipy" \
   "dask-cuda"
 
+# packages holds all the downstream and upstream dependencies.
+# we want to avoid directories with the same name as packages
+# in the working directory
+mkdir -p packages
+
 # Clone cudf repo for tests
 CUDF_VERSION="branch-25.04"
+
 cudf_commit=$(./scripts/check-version.py cudf)
 
 if [ ! -d "cudf" ]; then
     echo "Cloning cudf@{$CUDF_VERSION}"
-    git clone https://github.com/rapidsai/cudf.git --branch $CUDF_VERSION
+    git clone https://github.com/rapidsai/cudf.git --branch $CUDF_VERSION packages
 fi
 
-pushd cudf
+pushd packages/cudf
 git checkout $cudf_commit
 popd
 
+cuml_commit=$(./scripts/check-version.py cuml)
+
+if [ ! -d "cuml" ]; then
+    echo "Cloning cuml@{$CUDF_VERSION}"
+    git clone https://github.com/rapidsai/cuml.git --branch $CUDF_VERSION packages
+fi
+
+pushd packages/cuml
+git checkout $cuml_commit
+popd
+
+raft_commit=$(./scripts/check-version.py raft_dask)
+
+if [ ! -d "raft" ]; then
+    echo "Cloning raft@{$CUDF_VERSION}"
+    git clone https://github.com/rapidsai/raft.git --branch $CUDF_VERSION packages
+fi
+
+pushd packages/raft
+git checkout $raft_commit
+popd
+
 if [ ! -d "dask-cuda" ]; then
     echo "Cloning cudf@{$CUDF_VERSION}"
-    git clone https://github.com/rapidsaicudf_commit/dask-cuda.git --branch $CUDF_VERSION
+    git clone https://github.com/rapidsaicudf_commit/dask-cuda.git --branch $CUDF_VERSION packages
 fi
 
 # Clone dask-cuda for tests
@@ -42,32 +72,32 @@ fi
 
 # dask_cuda_commit=$(./scripts/check-version.py dask_cuda)
 
-pushd dask-cuda
+pushd packages/dask-cuda
 git checkout $CUDF_VERSION
 popd
 
 # depth needs to be sufficient to reach the last tag, so that the package
 # versions are set correctly
 if [ ! -d "dask" ]; then
     echo "Cloning dask@{$DASK_VERSION}"
-    git clone https://github.com/dask/dask --depth 100 --branch $DASK_VERSION
+    git clone https://github.com/dask/dask --depth 100 --branch $DASK_VERSION packages
 fi
 
 if [ ! -d "distributed" ]; then
     echo "Cloning dask@{$DASK_VERSION}"
-    git clone https://github.com/dask/distributed --depth 100 --branch $DASK_VERSION
+    git clone https://github.com/dask/distributed --depth 100 --branch $DASK_VERSION packages
 fi
 
-pushd dask
+pushd packages/dask
 git checkout $DASK_VERSION
 popd
 
-pushd distributed
+pushd packages/distributed
 git checkout $DASK_VERSION
 popd
 
 # Finally, ensure that
-uv pip install --no-deps -e ./dask ./distributed
+uv pip install --no-deps -e ./packages/dask ./packages/distributed
 
 echo "[Setup done]"
 uv pip list

scripts/test.sh

@@ -2,20 +2,27 @@
 # SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION & AFFILIATES.
 
 if [ $# -eq 0 ]; then
+    run_cuml=true
     run_dask=true
     run_dask_cuda=true
     run_dask_cudf=true
     run_distributed=true
+    run_raft_dask=true
 else
+    run_cuml=false
     run_dask=false
     run_dask_cuda=false
     run_dask_cudf=false
     run_distributed=false
+    run_raft_dask=false
 fi
 
 # Parse command-line arguments
 while [[ $# -gt 0 ]]; do
     case "$1" in
+        --cuml-only)
+            run_cuml=true
+            ;;
         --dask-only)
             run_dask=true
             ;;
@@ -28,6 +35,9 @@ while [[ $# -gt 0 ]]; do
         --distributed-only)
             run_distributed=true
             ;;
+        --raft-dask-only)
+            run_raft_dask=true
+            ;;
         --help)
             echo "Usage: $0 [--dask-only] [--distributed-only] [--dask-cuda-only] [--dask-cudf-only]"
             exit 0
@@ -42,65 +52,81 @@ done
 
 exit_code=0;
 
+# --- cuml ---
+if $run_cuml; then
+
+    echo "[testing cuml]"
+    pytest -v packages/cuml/python/cuml/cuml/tests/dask
+
+    if [[ $? -ne 0 ]]; then
+        exit_code=1
+    fi
+
+fi
+
+
 # --- dask-cudf ---
 if $run_dask_cudf; then
 
     echo "[testing dask-cudf]"
-    pushd cudf/python/dask_cudf || exit 1
-    pytest -v dask_cudf
+    pytest -v packages/cudf/python/dask_cudf
 
     if [[ $? -ne 0 ]]; then
         exit_code=1
     fi
 
-    popd || exit 1
-
 fi
 
 # --- dask-cuda ---
 
 if $run_dask_cuda; then
     echo "[testing dask-cuda]"
-    pushd dask-cuda/dask_cuda/tests || exit 1
-    pytest -v .
+    pytest -v packages/dask-cuda/dask_cuda/tests
 
     if [[ $? -ne 0 ]]; then
         exit_code=1
     fi
 
-    popd || exit 1
 fi
 
 # --- dask ---
 
 if $run_dask; then
 
     echo "[testing dask]"
-    pushd dask/dask || exit 1
-    pytest -v -m gpu .
+    pytest -v -m gpu packages/dask/dask
 
     if [[ $? -ne 0 ]]; then
         exit_code=1
     fi
 
-    popd || exit 1
+fi
+
+# --- raft-dask ---
+if $run_raft_dask; then
+
+    echo "[testing raft-dask]"
+    pytest -v packages/raft/python/raft-dask/raft_dask/tests
+
+    if [[ $? -ne 0 ]]; then
+        exit_code=1
+    fi
 
 fi
 
+
+
 # --- distributed ---
 
 if $run_distributed; then
 
     echo "[testing distributed]"
-    pushd distributed || exit 1
-    pytest -v -m gpu --runslow distributed
+    pytest -v -m gpu --runslow packages/distributed/distributed
 
     if [[ $? -ne 0 ]]; then
         exit_code=1
     fi
 
-    popd || exit 1
-
 fi
 
 exit $exit_code