Add in some testing matrix updates (simple version)

cryos · cryos · commit 965e7a1ca1a0 · 2025-04-17T14:40:00.000-04:00
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -21,7 +21,7 @@ jobs:
         host-platform:
           - linux-64
           - linux-aarch64
-          - win-64
+          #- win-64
     name: Build ${{ matrix.host-platform }}
     if: ${{ github.repository_owner == 'nvidia' }}
     secrets: inherit
@@ -30,3 +30,26 @@ jobs:
     with:
       host-platform: ${{ matrix.host-platform }}
       cuda-version: "12.8.0"
+
+  test-linux:
+    strategy:
+      fail-fast: false
+      # TODO: add driver version here
+      matrix:
+        host-platform:
+          - linux-64
+          - linux-aarch64
+    name: Test ${{ matrix.host-platform }}
+    if: ${{ github.repository_owner == 'nvidia' }}
+    permissions:
+      contents: read  # This is required for actions/checkout
+    needs:
+      - build
+    secrets: inherit
+    uses:
+      ./.github/workflows/test-wheel-linux.yml
+    with:
+      build-type: pull-request
+      host-platform: ${{ matrix.host-platform }}
+      local-ctk: 1
+      build-ctk-ver: 12.8.0
diff --git a/.github/workflows/test-wheel-linux.yml b/.github/workflows/test-wheel-linux.yml
@@ -3,32 +3,81 @@ name: "CI: Test wheels"
 on:
   workflow_call:
     inputs:
-      host-platform:
+      build-type:
         type: string
         required: true
-      python-version:
+      host-platform:
         type: string
         required: true
       build-ctk-ver:
         type: string
         required: true
-      cuda-version:
-        type: string
-        required: true
       local-ctk:
         type: string
         required: true
-      runner:
-        type: string
-        required: true
 
 jobs:
+  compute-matrix:
+    runs-on: ubuntu-latest
+    env:
+      BUILD_TYPE: ${{ inputs.build-type }}
+      ARCH: ${{ (inputs.host-platform == 'linux-64' && 'amd64') ||
+                (inputs.host-platform == 'linux-aarch64' && 'arm64') }}
+    outputs:
+      MATRIX: ${{ steps.compute-matrix.outputs.MATRIX }}
+    steps:
+      - name: Validate Test Type
+        run: |
+          if [[ "$BUILD_TYPE" != "pull-request" ]] && [[ "$BUILD_TYPE" != "nightly" ]] && [[ "$BUILD_TYPE" != "branch" ]]; then
+              echo "Invalid build type! Must be one of 'nightly', 'pull-request', or 'branch'."
+              exit 1
+          fi
+      - name: Compute Python Test Matrix
+        id: compute-matrix
+        run: |
+          set -eo pipefail
+          # Please keep the matrices sorted in ascending order by the following:
+          #
+          #     [PY_VER, CUDA_VER, LINUX_VER, GPU, DRIVER]
+          #
+          gpu="l4"
+          if [[ "${ARCH}" == "arm64" ]]; then
+            gpu="a100"
+          fi
+          export MATRICES="
+            pull-request:
+              - { ARCH=${ARCH}, PY_VER: '3.9',  CUDA_VER: '11.8.0', LINUX_VER: 'rockylinux8', GPU: ${gpu}, DRIVER: 'earliest' }
+              - { ARCH=${ARCH}, PY_VER: '3.9',  CUDA_VER: '12.0.1', LINUX_VER: 'ubuntu24.04', GPU: ${gpu}, DRIVER: 'latest' }
+              - { ARCH=${ARCH}, PY_VER: '3.13', CUDA_VER: '12.8.0', LINUX_VER: 'ubuntu22.04', GPU: ${gpu}, DRIVER: 'latest' }
+            nightly:
+              - { ARCH=${ARCH}, PY_VER: '3.9',  CUDA_VER: '11.8.0', LINUX_VER: 'rockylinux8', GPU: ${gpu}, DRIVER: 'earliest' }
+              - { ARCH=${ARCH}, PY_VER: '3.13', CUDA_VER: '12.8.0', LINUX_VER: 'ubuntu22.04', GPU: ${gpu}, DRIVER: 'latest' }
+          "
+
+          # Use the nightly matrix for branch tests
+          MATRIX_TYPE="${BUILD_TYPE}"
+          if [[ "${MATRIX_TYPE}" == "branch" ]]; then
+            MATRIX_TYPE="nightly"
+          fi
+          export MATRIX_TYPE
+          TEST_MATRIX=$(yq -n 'env(MATRICES) | .[strenv(MATRIX_TYPE)]')
+          export TEST_MATRIX
+
+          MATRIX="$(
+            yq -n -o json 'env(TEST_MATRIX)' | \
+            jq -c '${{ inputs.matrix_filter }} | if (. | length) > 0 then {include: .} else "Error: Empty matrix\n" | halt_error(1) end'
+          )"
+
+          echo "MATRIX=${MATRIX}" | tee --append "${GITHUB_OUTPUT}"
+
   test:
+    needs: compute-matrix
+    strategy:
+      fail-fast: false
+      matrix: ${{ fromJSON(needs.compute-matrix.outputs.MATRIX) }}
+    runs-on: "linux-${{ matrix.ARCH }}-gpu-${{ matrix.GPU }}-${{ matrix.DRIVER }}-1"
     # The build stage could fail but we want the CI to keep moving.
     if: ${{ github.repository_owner == 'nvidia' && !cancelled() }}
-    runs-on: ${{ (inputs.runner == 'default' && inputs.host-platform == 'linux-64' && 'linux-amd64-gpu-v100-latest-1') ||
-                 (inputs.runner == 'default' && inputs.host-platform == 'linux-aarch64' && 'linux-arm64-gpu-a100-latest-1') ||
-                 (inputs.runner == 'H100' && 'linux-amd64-gpu-h100-latest-1') }}
     # Our self-hosted runners require a container
     # TODO: use a different (nvidia?) container
     container:
@@ -50,7 +99,7 @@ jobs:
 
       - name: Set environment variables
         run: |
-          PYTHON_VERSION_FORMATTED=$(echo '${{ inputs.python-version }}' | tr -d '.')
+          PYTHON_VERSION_FORMATTED=$(echo '${{ matrix.PY_VER }}' | tr -d '.')
           if [[ "${{ inputs.host-platform }}" == linux* ]]; then
             REPO_DIR=$(pwd)
           elif [[ "${{ inputs.host-platform }}" == win* ]]; then
@@ -59,14 +108,14 @@ jobs:
           fi
 
           BUILD_CUDA_MAJOR="$(cut -d '.' -f 1 <<< ${{ inputs.build-ctk-ver }})"
-          TEST_CUDA_MAJOR="$(cut -d '.' -f 1 <<< ${{ inputs.cuda-version }})"
+          TEST_CUDA_MAJOR="$(cut -d '.' -f 1 <<< ${{ matrix.CUDA_VER }})"
           if [[ $BUILD_CUDA_MAJOR != $TEST_CUDA_MAJOR ]]; then
             SKIP_CUDA_BINDINGS_TEST=1
             SKIP_CUDA_CORE_CYTHON_TEST=0
           else
             SKIP_CUDA_BINDINGS_TEST=0
             BUILD_CUDA_MINOR="$(cut -d '.' -f 2 <<< ${{ inputs.build-ctk-ver }})"
-            TEST_CUDA_MINOR="$(cut -d '.' -f 2 <<< ${{ inputs.cuda-version }})"
+            TEST_CUDA_MINOR="$(cut -d '.' -f 2 <<< ${{ matrix.CUDA_VER }})"
             if [[ $BUILD_CUDA_MINOR != $TEST_CUDA_MINOR ]]; then
               SKIP_CUDA_CORE_CYTHON_TEST=1
             else
@@ -164,10 +213,10 @@ jobs:
           pwd
           ls -lahR $CUDA_CORE_ARTIFACTS_DIR
 
-      - name: Set up Python ${{ inputs.python-version }}
+      - name: Set up Python ${{ matrix.PY_VER }}
         uses: actions/setup-python@v5
         with:
-          python-version: ${{ inputs.python-version }}
+          python-version: ${{ matrix.PY_VER }}
         env:
           # we use self-hosted runners on which setup-python behaves weirdly...
           AGENT_TOOLSDIRECTORY: "/opt/hostedtoolcache"
@@ -178,7 +227,7 @@ jobs:
         continue-on-error: false
         with:
           host-platform: ${{ inputs.host-platform }}
-          cuda-version: ${{ inputs.cuda-version }}
+          cuda-version: ${{ matrix.CUDA_VER }}
 
       - name: Run cuda.bindings tests
         if: ${{ env.SKIP_CUDA_BINDINGS_TEST == '0' }}
@@ -222,7 +271,7 @@ jobs:
             fi
             popd
           fi
-          TEST_CUDA_MAJOR="$(cut -d '.' -f 1 <<< ${{ inputs.cuda-version }})"
+          TEST_CUDA_MAJOR="$(cut -d '.' -f 1 <<< ${{ matrix.CUDA_VER }})"
           pushd "${CUDA_CORE_ARTIFACTS_DIR}"
           pip install $(ls *.whl)["cu${TEST_CUDA_MAJOR}"]
           popd
