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

---
 .github/workflows/build-and-test.yml          | 10 +++++---
 .github/workflows/build-docs.yml              |  6 ++---
 .../{test_local_ctk.yml => test-wheel.yml}    | 24 ++++++++++++++-----
 3 files changed, 28 insertions(+), 12 deletions(-)
 rename .github/workflows/{test_local_ctk.yml => test-wheel.yml} (93%)

diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml
index ccb620b4..a298b65e 100644
--- a/.github/workflows/build-and-test.yml
+++ b/.github/workflows/build-and-test.yml
@@ -198,6 +198,9 @@ jobs:
           - "12.6.2"
           - "12.0.1"
           - "11.8.0"
+        local-ctk:
+          - 1  # use mini CTK
+          - 0  # use CTK wheels
         runner:
           - default
         include:
@@ -214,12 +217,13 @@ jobs:
       - build
     secrets: inherit
     uses:
-      ./.github/workflows/test_local_ctk.yml
+      ./.github/workflows/test_wheel.yml
     with:
       host-platform: ${{ matrix.host-platform }}
       python-version: ${{ matrix.python-version }}
-      build_ctk_ver: ${{ needs.build.outputs.BUILD_CTK_VER }}
+      build-ctk-ver: ${{ needs.build.outputs.BUILD_CTK_VER }}
       cuda-version: ${{ matrix.cuda-version }}
+      local-ctk: ${{ matrix.local-ctk}}
       runner: ${{ matrix.runner }}
 
   doc:
@@ -236,4 +240,4 @@ jobs:
     uses:
       ./.github/workflows/build-docs.yml
     with:
-      build_ctk_ver: ${{ needs.build.outputs.BUILD_CTK_VER }}
+      build-ctk-ver: ${{ needs.build.outputs.BUILD_CTK_VER }}
diff --git a/.github/workflows/build-docs.yml b/.github/workflows/build-docs.yml
index cafb1fc9..f88e38bf 100644
--- a/.github/workflows/build-docs.yml
+++ b/.github/workflows/build-docs.yml
@@ -3,7 +3,7 @@ name: "CI: Build and update docs"
 on:
   workflow_call:
     inputs:
-      build_ctk_ver:
+      build-ctk-ver:
         type: string
         required: true
 
@@ -53,7 +53,7 @@ jobs:
         continue-on-error: false
         with:
           host-platform: linux-64
-          cuda-version: ${{ inputs.build_ctk_ver }}
+          cuda-version: ${{ inputs.build-ctk-ver }}
 
       - name: Set environment variables
         run: |
@@ -65,7 +65,7 @@ jobs:
           echo "CUDA_CORE_ARTIFACT_BASENAME=${CUDA_CORE_ARTIFACT_BASENAME}" >> $GITHUB_ENV
           echo "CUDA_CORE_ARTIFACT_NAME=${CUDA_CORE_ARTIFACT_BASENAME}-${{ github.sha }}" >> $GITHUB_ENV
           echo "CUDA_CORE_ARTIFACTS_DIR=$(realpath "$REPO_DIR/cuda_core/dist")" >> $GITHUB_ENV
-          CUDA_BINDINGS_ARTIFACT_BASENAME="cuda-bindings-python${PYTHON_VERSION_FORMATTED}-cuda${{ inputs.build_ctk_ver }}-linux-64"
+          CUDA_BINDINGS_ARTIFACT_BASENAME="cuda-bindings-python${PYTHON_VERSION_FORMATTED}-cuda${{ inputs.build-ctk-ver }}-linux-64"
           echo "CUDA_BINDINGS_ARTIFACT_BASENAME=${CUDA_BINDINGS_ARTIFACT_BASENAME}" >> $GITHUB_ENV
           echo "CUDA_BINDINGS_ARTIFACT_NAME=${CUDA_BINDINGS_ARTIFACT_BASENAME}-${{ github.sha }}" >> $GITHUB_ENV
           echo "CUDA_BINDINGS_ARTIFACTS_DIR=$(realpath "$REPO_DIR/cuda_bindings/dist")" >> $GITHUB_ENV
diff --git a/.github/workflows/test_local_ctk.yml b/.github/workflows/test-wheel.yml
similarity index 93%
rename from .github/workflows/test_local_ctk.yml
rename to .github/workflows/test-wheel.yml
index 4c9bcd7c..97c62a17 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:
@@ -9,12 +9,15 @@ on:
       python-version:
         type: string
         required: true
-      build_ctk_ver:
+      build-ctk-ver:
         type: string
         required: true
       cuda-version:
         type: string
         required: true
+      local-ctk:
+        type: string
+        required: true
       runner:
         type: string
         required: true
@@ -55,7 +58,7 @@ jobs:
             REPO_DIR=$(cygpath -w $PWD)
           fi
 
-          BUILD_CUDA_MAJOR="$(cut -d '.' -f 1 <<< ${{ inputs.build_ctk_ver }})"
+          BUILD_CUDA_MAJOR="$(cut -d '.' -f 1 <<< ${{ inputs.build-ctk-ver }})"
           TEST_CUDA_MAJOR="$(cut -d '.' -f 1 <<< ${{ inputs.cuda-version }})"
           if [[ $BUILD_CUDA_MAJOR != $TEST_CUDA_MAJOR ]]; then
             SKIP_CUDA_BINDINGS_TEST=1
@@ -69,7 +72,7 @@ jobs:
           echo "CUDA_CORE_ARTIFACT_BASENAME=${CUDA_CORE_ARTIFACT_BASENAME}" >> $GITHUB_ENV
           echo "CUDA_CORE_ARTIFACT_NAME=${CUDA_CORE_ARTIFACT_BASENAME}-${{ github.sha }}" >> $GITHUB_ENV
           echo "CUDA_CORE_ARTIFACTS_DIR=$(realpath "$REPO_DIR/cuda_core/dist")" >> $GITHUB_ENV
-          CUDA_BINDINGS_ARTIFACT_BASENAME="cuda-bindings-python${PYTHON_VERSION_FORMATTED}-cuda${{ inputs.build_ctk_ver }}-${{ inputs.host-platform }}"
+          CUDA_BINDINGS_ARTIFACT_BASENAME="cuda-bindings-python${PYTHON_VERSION_FORMATTED}-cuda${{ inputs.build-ctk-ver }}-${{ inputs.host-platform }}"
           echo "CUDA_BINDINGS_ARTIFACT_BASENAME=${CUDA_BINDINGS_ARTIFACT_BASENAME}" >> $GITHUB_ENV
           echo "CUDA_BINDINGS_ARTIFACT_NAME=${CUDA_BINDINGS_ARTIFACT_BASENAME}-${{ github.sha }}" >> $GITHUB_ENV
           echo "CUDA_BINDINGS_ARTIFACTS_DIR=$(realpath "$REPO_DIR/cuda_bindings/dist")" >> $GITHUB_ENV
@@ -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 }})"