From ee053ddcffe4f7b427b27e8e52e433027588271f Mon Sep 17 00:00:00 2001
From: Leo Fang <leo80042@gmail.com>
Date: Thu, 9 Jan 2025 00:23:15 -0500
Subject: [PATCH] add test jobs against CUDA wheels

---
 .github/workflows/build-and-test.yml           |  3 +++
 .../{test_local_ctk.yml => test_wheel.yml}     | 18 +++++++++++++++---
 2 files changed, 18 insertions(+), 3 deletions(-)
 rename .github/workflows/{test_local_ctk.yml => test_wheel.yml} (94%)

diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml
index ccb620b4..3be3b88e 100644
--- a/.github/workflows/build-and-test.yml
+++ b/.github/workflows/build-and-test.yml
@@ -200,6 +200,9 @@ jobs:
           - "11.8.0"
         runner:
           - default
+        local_ctk:
+          - 1  # use mini CTK
+          - 0  # use CTK wheels
         include:
           - host-platform: linux-64
             python-version: "3.12"
diff --git a/.github/workflows/test_local_ctk.yml b/.github/workflows/test_wheel.yml
similarity index 94%
rename from .github/workflows/test_local_ctk.yml
rename to .github/workflows/test_wheel.yml
index 4c9bcd7c..2e93724e 100644
--- a/.github/workflows/test_local_ctk.yml
+++ b/.github/workflows/test_wheel.yml
@@ -1,4 +1,4 @@
-name: "CI: Test against local CTK"
+name: "CI: Test wheels"
 
 on:
   workflow_call:
@@ -15,6 +15,9 @@ on:
       cuda-version:
         type: string
         required: true
+      local_ctk:
+        type: string
+        required: true
       runner:
         type: string
         required: true
@@ -138,6 +141,7 @@ jobs:
           AGENT_TOOLSDIRECTORY: "/opt/hostedtoolcache"
 
       - name: Set up mini CTK
+        if: ${{ inputs.local_ctk == '1' }}
         uses: ./.github/actions/fetch_ctk
         continue-on-error: false
         with:
@@ -150,7 +154,11 @@ jobs:
           ls $CUDA_PATH
 
           pushd "${CUDA_BINDINGS_ARTIFACTS_DIR}"
-          pip install *.whl
+          if [[ "${{ inputs.local_ctk }}" == 1 ]]; then
+            pip install *.whl
+          else
+            pip install $(ls *.whl)[all]
+          fi
           popd
 
           pushd ./cuda_bindings
@@ -171,7 +179,11 @@ jobs:
           # If mismatch: cuda.bindings is installed from the backport branch.
           if [[ "${SKIP_CUDA_BINDINGS_TEST}" == 1 ]]; then
             pushd "${CUDA_BINDINGS_ARTIFACTS_DIR}"
-            pip install *.whl
+            if [[ "${{ inputs.local_ctk }}" == 1 ]]; then
+              pip install *.whl
+            else
+              pip install $(ls *.whl)[all]
+            fi
             popd
           fi
           TEST_CUDA_MAJOR="$(cut -d '.' -f 1 <<< ${{ inputs.cuda-version }})"