From 2164a69abef5815cdf59ad15de3cb3f2c28a7d22 Mon Sep 17 00:00:00 2001
From: Leo Fang <leo80042@gmail.com>
Date: Sat, 11 Jan 2025 18:27:32 -0500
Subject: [PATCH] Test against CUDA wheels

---
 .github/workflows/build-and-test.yml | 37 ++++++++++++++++++----------
 1 file changed, 24 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml
index def159f5..3a249eda 100644
--- a/.github/workflows/build-and-test.yml
+++ b/.github/workflows/build-and-test.yml
@@ -158,6 +158,9 @@ jobs:
         cuda-version:
           # Note: this is for test-time only.
           - "11.8.0"
+        local-ctk:
+          - 1  # use mini CTK
+          - 0  # use CTK wheels
         runner:
           - default
         include:
@@ -165,7 +168,7 @@ jobs:
             python-version: "3.12"
             cuda-version: "11.8.0"
             runner: H100
-    name: Test (${{ matrix.host-platform }}, Python ${{ matrix.python-version }}, CUDA ${{ matrix.cuda-version }}, Runner ${{ matrix.runner }})
+    name: Test (${{ matrix.host-platform }}, Python ${{ matrix.python-version }}, CUDA ${{ matrix.cuda-version }}, Runner ${{ matrix.runner }}, ${{ (matrix.local-ctk == '1' && 'local CTK') || 'CTK wheels' }})
     # The build stage could fail but we want the CI to keep moving.
     if: ${{ github.repository_owner == 'nvidia' && always() }}
     permissions:
@@ -239,6 +242,7 @@ jobs:
           AGENT_TOOLSDIRECTORY: "/opt/hostedtoolcache"
 
       - name: Set up mini CTK
+        if: ${{ matrix.local-ctk == '1' }}
         uses: ./.github/actions/fetch_ctk
         continue-on-error: false
         with:
@@ -249,22 +253,29 @@ jobs:
         if: ${{ env.SKIP_CUDA_BINDINGS_TEST == '0' }}
         shell: bash --noprofile --norc -xeuo pipefail {0}
         run: |
-          ls $CUDA_PATH
-
           pushd "${CUDA_BINDINGS_ARTIFACTS_DIR}"
-          pip install *.whl
+          if [[ "${{ matrix.local-ctk }}" == 1 ]]; then
+            ls $CUDA_PATH
+            pip install *.whl
+          else
+            pip install $(ls *.whl)[all]
+          fi
           popd
 
           pushd ./cuda_bindings
           pip install -r requirements.txt
           pytest -rxXs tests/
-          if [[ "${{ matrix.host-platform }}" == linux* ]]; then
-            # cython tests require gcc
-            apt install -y build-essential
-            bash tests_cython/build_tests.sh
-          elif [[ "${{ matrix.host-platform }}" == win* ]]; then
-            # TODO: enable this once win-64 runners are up 
-            exit 1
-          fi  
-          pytest -rxXs tests_cython
+          # It is a bit convoluted to run the Cython tests against CTK wheels,
+          # so let's just skip them.
+          if [[ "${{ matrix.local-ctk }}" == 1 ]]; then
+            if [[ "${{ matrix.host-platform }}" == linux* ]]; then
+              # cython tests require gcc
+              apt install -y build-essential
+              bash tests_cython/build_tests.sh
+            elif [[ "${{ matrix.host-platform }}" == win* ]]; then
+              # TODO: enable this once win-64 runners are up 
+              exit 1
+            fi  
+            pytest -rxXs tests_cython
+          fi
           popd