add test jobs against CUDA wheels

leofang · leofang · commit 7f36568f3d87 · 2025-01-09T10:30:06.000-05:00
diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml
@@ -198,14 +198,17 @@ jobs:
           - "12.6.2"
           - "12.0.1"
           - "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: "12.6.2"
             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' && !cancelled() }}
     permissions:
@@ -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
@@ -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-wheel.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:
@@ -147,31 +151,43 @@ jobs:
       - name: Run cuda.bindings tests
         if: ${{ env.SKIP_CUDA_BINDINGS_TEST == '0' }}
         run: |
-          ls $CUDA_PATH
-
           pushd "${CUDA_BINDINGS_ARTIFACTS_DIR}"
-          pip install *.whl
+          if [[ "${{ inputs.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 [[ "${{ inputs.host-platform }}" == linux* ]]; then
-            bash tests/cython/build_tests.sh
-          elif [[ "${{ inputs.host-platform }}" == win* ]]; then
-            # TODO: enable this once win-64 runners are up 
-            exit 1
-          fi  
-          pytest -rxXs tests/cython
-          popd
+
+          # It is a bit convoluted to run the Cython tests against CTK wheels,
+          # so let's just skip them.
+          if [[ "${{ inputs.local-ctk }}" == 1 ]]; then
+            if [[ "${{ inputs.host-platform }}" == linux* ]]; then
+              bash tests/cython/build_tests.sh
+            elif [[ "${{ inputs.host-platform }}" == win* ]]; then
+              # TODO: enable this once win-64 runners are up 
+              exit 1
+            fi  
+            pytest -rxXs tests/cython
+            popd
+          fi
 
       - name: Run cuda.core tests
         run: |
           # If build/test majors match: cuda.bindings is installed in the previous step.
           # 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 }})"
diff --git a/cuda_core/tests/requirements-cu11.txt b/cuda_core/tests/requirements-cu11.txt
@@ -1,3 +1,4 @@
 pytest
 # TODO: remove this hack once cupy has a cp313 build
 cupy-cuda11x; python_version < "3.13"
+nvidia-cuda-runtime-cu11  # headers consumed by CuPy
diff --git a/cuda_core/tests/requirements-cu12.txt b/cuda_core/tests/requirements-cu12.txt
@@ -1,3 +1,4 @@
 pytest
 # TODO: remove this hack once cupy has a cp313 build
 cupy-cuda12x; python_version < "3.13"
+nvidia-cuda-runtime-cu12  # headers consumed by CuPy
