Merge pull request #11 from TomAugspurger/tom/dask-cuda

Added dask-cuda

Author: TomAugspurger

.gitignore

@@ -1,3 +1,4 @@
 dask
 distributed
 cudf
+dask-cuda

scripts/run.sh

@@ -19,27 +19,36 @@ pip install --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/
   "cudf-${RAPIDS_PY_CUDA_SUFFIX}" \
   "dask-cudf-${RAPIDS_PY_CUDA_SUFFIX}" \
   "ucx-py-${RAPIDS_PY_CUDA_SUFFIX}" \
+  "ucxx-${RAPIDS_PY_CUDA_SUFFIX}" \
   "scipy" \
   "dask-cuda"
 
 
-# Clone cudf repo (need dask_cudf tests)
+# Clone cudf repo for tests
 CUDF_VERSION="branch-25.04"
-echo "Cloning cudf@{$CUDF_VERSION}"
 
 if [ ! -d "cudf" ]; then
+    echo "Cloning cudf@{$CUDF_VERSION}"
     git clone https://github.com/rapidsai/cudf.git --branch $CUDF_VERSION
 fi
 
-echo "Installing dask@{$DASK_VERSION}"
+# Clone dask-cuda for tests
+
+if [ ! -d "dask-cuda" ]; then
+    echo "Cloning cudf@{$CUDF_VERSION}"
+    git clone https://github.com/rapidsai/dask-cuda.git --branch $CUDF_VERSION
+fi
+
 
 # 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
 fi
 
 if [ ! -d "distributed" ]; then
+    echo "Cloning dask@{$DASK_VERSION}"
     git clone https://github.com/dask/distributed --depth 100 --branch $DASK_VERSION
 fi
 

scripts/test.sh

@@ -1,25 +1,106 @@
 #!/usr/bin/env bash
 # SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION & AFFILIATES.
 
-echo "[testing dask-cudf]"
-pushd cudf/python/dask_cudf || exit 1
-pytest dask_cudf -v
-dask_cudf_status=$?
-popd || exit 1
-
-echo "[testing dask]"
-pushd dask || exit 1
-pytest dask -v -m gpu
-dask_status=$?
-popd || exit 1
-
-echo "[testing distributed]"
-pushd distributed || exit 1
-pytest distributed -v -m gpu --runslow
-distributed_status=$?
-popd || exit 1
-
-if [ $dask_cudf_status -ne 0 ] || [ $dask_status -ne 0 ] || [ $distributed_status -ne 0 ]; then
-    echo "Tests faild"
-    exit 1
+if [ $# -eq 0 ]; then
+    run_dask=true
+    run_dask_cuda=true
+    run_dask_cudf=true
+    run_distributed=true
+else
+    run_dask=false
+    run_dask_cuda=false
+    run_dask_cudf=false
+    run_distributed=false
 fi
+
+# Parse command-line arguments
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --dask-only)
+            run_dask=true
+            ;;
+        --dask-cuda-only)
+            run_dask_cuda=true
+            ;;
+        --dask-cudf-only)
+            run_dask_cudf=true
+            ;;
+        --distributed-only)
+            run_distributed=true
+            ;;
+        --help)
+            echo "Usage: $0 [--dask-only] [--distributed-only] [--dask-cuda-only] [--dask-cudf-only]"
+            exit 0
+            ;;
+        *)
+            echo "Unknown option: $1"
+            exit 1
+            ;;
+    esac
+    shift
+done
+
+exit_code=0;
+
+# --- dask-cudf ---
+if $run_dask_cudf; then
+
+    echo "[testing dask-cudf]"
+    pushd cudf/python/dask_cudf || exit 1
+    pytest -v 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 .
+
+    if [[ $? -ne 0 ]]; then
+        exit_code=1
+    fi
+
+    popd || exit 1
+fi
+
+# --- dask ---
+
+if $run_dask; then
+
+    echo "[testing dask]"
+    pushd dask || exit 1
+    pytest -v -m gpu dask
+
+    if [[ $? -ne 0 ]]; then
+        exit_code=1
+    fi
+
+    popd || exit 1
+
+fi
+
+# --- dask ---
+
+if $run_distributed; then
+
+    echo "[testing distributed]"
+    pushd distributed || exit 1
+    pytest -v -m gpu --runslow distributed
+
+    if [[ $? -ne 0 ]]; then
+        exit_code=1
+    fi
+
+    popd || exit 1
+
+fi
+
+exit $exit_code