add test jobs against CUDA wheels

NVIDIA · Jan 9, 2025 · ee053dd · ee053dd
1 parent 5015410
commit ee053dd
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 3 deletions.
diff --git 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 → .github/workflows/test_wheel.yml b/.github/workflows/test_local_ctk.yml → .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 }})"