From a1d78742a2e26ee04daafadadd91f425c7d0a340 Mon Sep 17 00:00:00 2001
From: Gil Forsyth <gforsyth@users.noreply.github.com>
Date: Mon, 6 Jan 2025 16:02:07 -0500
Subject: [PATCH] Add wheel tests to CI (#1416)

Adding this now that wheels are available

- **deps(kvikio): add kvikio to CUDA version matrices**
- **test(wheels): enable wheel tests in CI**

Resolves #1344

Authors:
  - Gil Forsyth (https://github.com/gforsyth)

Approvers:
  - Peter Andreas Entschev (https://github.com/pentschev)
  - James Lamb (https://github.com/jameslamb)

URL: https://github.com/rapidsai/dask-cuda/pull/1416
---
 .github/workflows/pr.yaml   |  8 ++++++++
 .github/workflows/test.yaml |  9 +++++++++
 ci/build_wheel.sh           |  2 +-
 ci/test_wheel.sh            | 19 +++++++++++++++++++
 dependencies.yaml           | 12 ++----------
 pyproject.toml              |  8 --------
 6 files changed, 39 insertions(+), 19 deletions(-)
 create mode 100755 ci/test_wheel.sh

diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml
index 1193161a..ef24aab8 100644
--- a/.github/workflows/pr.yaml
+++ b/.github/workflows/pr.yaml
@@ -17,6 +17,7 @@ jobs:
       - conda-python-tests
       - docs-build
       - wheel-build
+      - wheel-tests
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-25.02
   checks:
@@ -53,3 +54,10 @@ jobs:
       # This selects "ARCH=amd64 + the latest supported Python + CUDA".
       matrix_filter: map(select(.ARCH == "amd64")) | max_by([(.PY_VER|split(".")|map(tonumber)), (.CUDA_VER|split(".")|map(tonumber))]) | [.]
       script: "ci/build_wheel.sh"
+  wheel-tests:
+    needs: wheel-build
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-25.02
+    with:
+      build_type: pull-request
+      script: "ci/test_wheel.sh"
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 237a38a0..c6e2e852 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -22,3 +22,12 @@ jobs:
       branch: ${{ inputs.branch }}
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
+  wheel-tests:
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-25.02
+    with:
+      build_type: nightly
+      branch: ${{ inputs.branch }}
+      date: ${{ inputs.date }}
+      sha: ${{ inputs.sha }}
+      script: ci/test_wheel.sh
diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index 760e46e3..741aeab9 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -10,4 +10,4 @@ rapids-generate-version > ./VERSION
 python -m pip wheel . -w dist -v --no-deps --disable-pip-version-check
 ./ci/validate_wheel.sh dist
 
-RAPIDS_PY_WHEEL_NAME="dask-cuda" rapids-upload-wheels-to-s3 dist
+RAPIDS_PY_WHEEL_NAME="dask-cuda" RAPIDS_PY_WHEEL_PURE="1" rapids-upload-wheels-to-s3 dist
diff --git a/ci/test_wheel.sh b/ci/test_wheel.sh
new file mode 100755
index 00000000..f883a735
--- /dev/null
+++ b/ci/test_wheel.sh
@@ -0,0 +1,19 @@
+#!/bin/bash
+# Copyright (c) 2025, NVIDIA CORPORATION.
+
+set -eou pipefail
+
+RAPIDS_PY_WHEEL_NAME="dask-cuda" RAPIDS_PY_WHEEL_PURE="1" rapids-download-wheels-from-s3 python ./dist
+
+# Install cuda-suffixed dependencies b/c while `dask-cuda` has no cuda suffix, the test dependencies do
+rapids-dependency-file-generator \
+    --output requirements \
+    --file-key "test_python" \
+    --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION};cuda_suffixed=true" \
+| tee /tmp/requirements-test.txt
+
+rapids-logger "Installing test dependencies"
+# echo to expand wildcard
+python -m pip install -v --prefer-binary -r /tmp/requirements-test.txt $(echo ./dist/dask_cuda*.whl)
+
+python -m pytest ./dask_cuda/tests -k "not ucxx"
diff --git a/dependencies.yaml b/dependencies.yaml
index 7b4d9434..4b2fe2a0 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -47,14 +47,6 @@ files:
       table: project
     includes:
       - run_python
-  py_test:
-    output: pyproject
-    pyproject_dir: .
-    extras:
-      table: project.optional-dependencies
-      key: test
-    includes:
-      - test_python
   py_docs:
     output: pyproject
     pyproject_dir: .
@@ -188,14 +180,13 @@ dependencies:
               - numactl-devel-cos7-aarch64
       - output_types: [requirements, pyproject]
         matrices:
-          # kvikio should be added to the CUDA-version-specific matrices once there are wheels available
-          # ref: https://github.com/rapidsai/kvikio/pull/369
           - matrix:
               cuda: "12.*"
               cuda_suffixed: "true"
             packages:
               - cudf-cu12==25.2.*,>=0.0.0a0
               - dask-cudf-cu12==25.2.*,>=0.0.0a0
+              - kvikio-cu12==25.2.*,>=0.0.0a0
               - ucx-py-cu12==0.42.*,>=0.0.0a0
           - matrix:
               cuda: "11.*"
@@ -203,6 +194,7 @@ dependencies:
             packages:
               - cudf-cu11==25.2.*,>=0.0.0a0
               - dask-cudf-cu11==25.2.*,>=0.0.0a0
+              - kvikio-cu11==25.2.*,>=0.0.0a0
               - ucx-py-cu11==0.42.*,>=0.0.0a0
           - matrix:
             packages:
diff --git a/pyproject.toml b/pyproject.toml
index 7469a345..f1d67c2a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,14 +49,6 @@ docs = [
     "sphinx-click>=2.7.1",
     "sphinx-rtd-theme>=0.5.1",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit dependencies.yaml and run `rapids-dependency-file-generator`.
-test = [
-    "cudf==25.2.*,>=0.0.0a0",
-    "dask-cudf==25.2.*,>=0.0.0a0",
-    "kvikio==25.2.*,>=0.0.0a0",
-    "pytest",
-    "pytest-cov",
-    "ucx-py==0.42.*,>=0.0.0a0",
-] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit dependencies.yaml and run `rapids-dependency-file-generator`.
 
 [project.urls]
 Homepage = "https://github.com/rapidsai/dask-cuda"