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: #1416

.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/[email protected]
   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/[email protected]
+    with:
+      build_type: pull-request
+      script: "ci/test_wheel.sh"

.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/[email protected]
+    with:
+      build_type: nightly
+      branch: ${{ inputs.branch }}
+      date: ${{ inputs.date }}
+      sha: ${{ inputs.sha }}
+      script: ci/test_wheel.sh

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

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"

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,21 +180,21 @@ 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.*"
               cuda_suffixed: "true"
             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:

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"