Test against CUDA wheels

Author: leofang

.github/workflows/build-and-test.yml

@@ -158,14 +158,17 @@ 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:
           - host-platform: linux-64
             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,30 @@ 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