Allow reuse of resolved Python test matrix snapshots

.github/actionlint.yaml

@@ -12,6 +12,10 @@ config-variables:
   - TELEMETRY_ENABLED
 
 paths:
+  .github/workflows/resolve-python-test-matrices.yaml:
+    ignore:
+      # https://github.com/rhysd/actionlint/issues/711
+      - 'specifying action "\$/.github/actions/prepare-matrix-dispatch" in invalid format because ref is missing. available formats are "{owner}/{repo}@{ref}" or "{owner}/{repo}/{path}@{ref}"'
   .github/workflows/compute-matrix.yaml:
     ignore:
       # https://github.com/rhysd/actionlint/issues/711

.github/actions/prepare-matrix-dispatch/action.yaml

@@ -8,6 +8,10 @@ inputs:
   action-ref:
     description: Action ref.
     required: true
+  resolved-matrix:
+    description: When supplied, only check out the implementation; the caller already resolved its matrix snapshot.
+    required: false
+    default: ''
 
 outputs:
   matrix:
@@ -26,4 +30,5 @@ runs:
         persist-credentials: false
     - id: prepare-matrix
       name: Prepare Build Matrix
+      if: inputs.resolved-matrix == ''
       uses: ./shared-workflows/.github/actions/prepare-matrix

.github/actions/prepare-matrix/compute.sh

@@ -0,0 +1,87 @@
+#!/bin/bash
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+set -euo pipefail
+
+RESOLVED_MATRIX=${RESOLVED_MATRIX:-}
+MATRIX_FILTER=${MATRIX_FILTER:-.}
+PURE_WHEEL=${PURE_WHEEL:-false}
+PURE_CONDA=${PURE_CONDA:-false}
+
+case "$BUILD_TYPE" in
+  branch|nightly|pull-request|release-candidate) ;;
+  *) echo "::error::Invalid build_type: $BUILD_TYPE"; exit 1 ;;
+esac
+
+if [[ -n "$RESOLVED_MATRIX" ]]; then
+  if ! jq -en '
+    $ENV.RESOLVED_MATRIX | fromjson
+    | type == "object" and (keys == ["include"])
+      and (.include | type == "array" and length > 0
+        and all(.[]; type == "object"))
+  ' >/dev/null; then
+    echo '::error::resolved-matrix must be a nonempty {"include": [...]} matrix'
+    exit 1
+  fi
+  BASE_MATRIX=$(jq -cn '$ENV.RESOLVED_MATRIX | fromjson | .include')
+else
+  if ! jq -en '$ENV.MATRIX | fromjson | has($ENV.MATRIX_NAME)' >/dev/null; then
+    echo "::error::Invalid matrix_name: $MATRIX_NAME"
+    exit 1
+  fi
+  if [[ "$MATRIX_TYPE" == auto ]]; then
+    case "$BUILD_TYPE" in
+      branch|release-candidate) MATRIX_TYPE=nightly ;;
+      *) MATRIX_TYPE=$BUILD_TYPE ;;
+    esac
+  fi
+  export MATRIX_TYPES
+  MATRIX_TYPES=$(jq -cn --arg types "$MATRIX_TYPE" '$types | split(",")')
+  if ! jq -en '
+    ($ENV.MATRIX | fromjson)[$ENV.MATRIX_NAME] as $matrix
+    | $ENV.MATRIX_TYPES | fromjson
+    | length > 0 and all(.[]; . as $type | $matrix | has($type))
+  ' >/dev/null; then
+    echo "::error::Invalid matrix_type: $MATRIX_TYPE"
+    exit 1
+  fi
+  BASE_MATRIX=$(jq -cn '
+    ($ENV.MATRIX | fromjson)[$ENV.MATRIX_NAME] as $matrix
+    | [$ENV.MATRIX_TYPES | fromjson | .[] as $type | $matrix[$type][]]
+    | reduce .[] as $entry ([]; if index($entry) == null then . + [$entry] else . end)
+  ')
+fi
+export BASE_MATRIX
+
+# Preserve the existing default selectors for build workflows.
+if [[ "$MATRIX_FILTER" == . ]]; then
+  if [[ "$PURE_WHEEL" == true ]]; then
+    MATRIX_FILTER='map(select(.ARCH == "amd64")) | group_by(.CUDA_VER) | map(max_by(.PY_VER | split(".") | map(tonumber)))'
+  elif [[ "$PURE_CONDA" == true ]]; then
+    MATRIX_FILTER='map(select(.ARCH == "amd64")) | sort_by(.CUDA_VER, .PY_VER) | [last]'
+  elif [[ "$PURE_CONDA" == cuda_major ]]; then
+    MATRIX_FILTER='map(select(.ARCH == "amd64")) | group_by(.CUDA_VER|split(".")|map(tonumber)|.[0]) | map(max_by([(.PY_VER|split(".")|map(tonumber)), (.CUDA_VER|split(".")|map(tonumber))]))'
+  fi
+fi
+
+if ! RESULT=$(jq -ecn "
+  \$ENV.BASE_MATRIX | fromjson | ${MATRIX_FILTER}
+  | if type == \"array\" and length > 0 and all(.[]; type == \"object\")
+    then {include: .} else error(\"Empty or invalid matrix\") end
+"); then
+  echo '::error::Empty or invalid matrix'
+  exit 1
+fi
+
+if [[ "$MATRIX_NAME" == wheels-test || "$MATRIX_NAME" == conda-python-tests ]]; then
+  if ! jq -e '
+    .include | all(.[];
+      . as $entry | ["ARCH", "CUDA_VER", "PY_VER", "LINUX_VER", "GPU", "DRIVER", "DEPENDENCIES"]
+      | all(.[]; $entry[.] | type == "string" and length > 0))
+  ' <<< "$RESULT" >/dev/null; then
+    echo '::error::Python test matrix entries must contain all runner and environment fields'
+    exit 1
+  fi
+fi
+
+echo "matrix=$RESULT" | tee --append "$GITHUB_OUTPUT"

.github/workflows/compute-matrix.yaml

@@ -6,45 +6,35 @@ on:
         required: true
         type: string
       matrix_name:
-        description: "One of: [conda-cpp-build, conda-cpp-tests, conda-python-build, conda-python-tests, wheels-build, wheels-test]"
+        description: Name of the workflow matrix in prepare-matrix/matrix.yaml.
         required: true
         type: string
       matrix_type:
-        description: "One or more comma-separated matrix types defined for this workflow's matrix in .github/actions/prepare-matrix/matrix.yaml, e.g. 'pull-request' or 'nightly'. Defaults to 'auto', which must be used alone and means 'choose a value based on what's provided via build_type'."
-        required: false
+        description: Comma-separated matrix types, or auto to select by build_type.
         type: string
-        default: "auto"
+        default: auto
       matrix_filter:
-        description: |
-          jq expression which modifies the matrix.
-          For example, 'map(select(.ARCH == "amd64"))' to achieve "only run amd64 jobs".
+        description: jq expression which modifies the matrix entry array.
         type: string
         default: "."
+      resolved-matrix:
+        description: Optional nonempty JSON include matrix. Reuses the caller's snapshot instead of resolving current definitions; matrix_filter still applies.
+        type: string
+        default: ''
       pure_wheel:
-        description: "One of [true, false], true if the wheel is not dependent on operating system, Python minor version, or CPU architecture"
-        required: false
+        description: Whether the wheel is independent of OS, Python, and architecture.
         type: boolean
         default: false
       pure-conda:
-        required: false
+        description: "One of: [true, false, cuda_major]"
         type: string
         default: "false"
-        description: |
-          One of [true, false, cuda_major]. 'true' if the conda package is not
-          dependent on operating system, Python minor version, CPU architecture,
-          or CUDA version. 'cuda_major' if the package is not dependent on
-          operating system, Python minor version, or CPU architecture, but is
-          dependent on CUDA major version.
     outputs:
       matrix:
         value: ${{ jobs.compute-matrix.outputs.matrix }}
 jobs:
   compute-matrix:
     runs-on: ubuntu-latest
-    env:
-      BUILD_TYPE: ${{ inputs.build_type }}
-      MATRIX_NAME: ${{ inputs.matrix_name }}
-      MATRIX_TYPE: ${{ inputs.matrix_type }}
     outputs:
       matrix: ${{ steps.compute-matrix.outputs.matrix }}
     steps:
@@ -54,74 +44,16 @@ jobs:
         with:
           action-repo: ${{ github.action_repository }}
           action-ref: ${{ github.action_ref }}
-      - name: Validate Inputs
-        env:
-          MATRIX: ${{ steps.prepare-matrix.outputs.matrix }}
-        run: |
-          if [[ "$BUILD_TYPE" != "branch" ]] && [[ "$BUILD_TYPE" != "nightly" ]] && [[ "$BUILD_TYPE" != "pull-request" ]] && [[ "$BUILD_TYPE" != "release-candidate" ]]; then
-              echo "::error::Invalid build_type! Must be one of 'branch', 'nightly', 'pull-request', or 'release-candidate'."
-              exit 1
-          fi
-          if ! jq -n -e '$ENV.MATRIX | fromjson | has($ENV.MATRIX_NAME)'; then
-              echo "::error::Invalid matrix_name! Must be one of '$(jq -n -r "\$ENV.MATRIX | fromjson | keys | join(\"', '\")")'."
-              exit 1
-          fi
-          if [[ "$MATRIX_TYPE" != "auto" ]] && ! jq -n -e '($ENV.MATRIX | fromjson)[$ENV.MATRIX_NAME] as $matrix | $ENV.MATRIX_TYPE | split(",") | length > 0 and all(. as $matrix_type | $matrix | has($matrix_type))'; then
-              echo "::error::Invalid matrix_type! Must be 'auto' or a comma-separated list of matrix types defined for '${MATRIX_NAME}': $(jq -n -r '($ENV.MATRIX | fromjson)[$ENV.MATRIX_NAME] | keys | join(", ")')."
-              exit 1
-          fi
+          resolved-matrix: ${{ inputs.resolved-matrix }}
       - name: Compute Build Matrix
         id: compute-matrix
         env:
+          BUILD_TYPE: ${{ inputs.build_type }}
+          MATRIX_NAME: ${{ inputs.matrix_name }}
+          MATRIX_TYPE: ${{ inputs.matrix_type }}
           MATRIX: ${{ steps.prepare-matrix.outputs.matrix }}
+          RESOLVED_MATRIX: ${{ inputs.resolved-matrix }}
           MATRIX_FILTER: ${{ inputs.matrix_filter }}
           PURE_WHEEL: ${{ inputs.pure_wheel }}
           PURE_CONDA: ${{ inputs.pure-conda }}
-        run: |
-          set -eo pipefail
-
-          # only overwrite MATRIX_TYPE if it was set to 'auto'
-          if [[ "${MATRIX_TYPE}" == "auto" ]]; then
-            if [[ "${BUILD_TYPE}" == "branch" || "${BUILD_TYPE}" == "release-candidate" ]]; then
-              # Use the nightly matrix for branch tests
-              MATRIX_TYPE="nightly"
-            else
-              MATRIX_TYPE="${BUILD_TYPE}"
-            fi
-          fi
-          MATRIX_TYPES="$(jq -cn --arg matrix_types "${MATRIX_TYPE}" '$matrix_types | split(",")')"
-          export MATRIX_TYPES
-
-          # When pure-wheel is true and matrix_filter is default, override to build one wheel per CUDA version with amd64 and the latest PY_VER
-          if [ "${PURE_WHEEL}" = "true" ] && [ "${MATRIX_FILTER}" = "." ]; then
-            MATRIX_FILTER="map(select(.ARCH == \"amd64\")) | group_by(.CUDA_VER) | map(max_by(.PY_VER | split(\".\") | map(tonumber)))"
-          fi
-
-          # When pure-conda is true and matrix_filter is default, override to build one conda package with amd64, latest CUDA_VER, and the latest PY_VER
-          if [ "${PURE_CONDA}" = "true" ] && [ "${MATRIX_FILTER}" = "." ]; then
-            MATRIX_FILTER="map(select(.ARCH == \"amd64\")) | sort_by(.CUDA_VER, .PY_VER) | [last]"
-          elif [ "${PURE_CONDA}" = "cuda_major" ] && [ "${MATRIX_FILTER}" = "." ]; then
-            MATRIX_FILTER="map(select(.ARCH == \"amd64\")) | group_by(.CUDA_VER|split(\".\")|map(tonumber)|.[0]) | map(max_by([(.PY_VER|split(\".\")|map(tonumber)), (.CUDA_VER|split(\".\")|map(tonumber))]))"
-          fi
-
-          if ! MATRIX="$(
-            jq -n -c "
-              \$ENV.MATRIX | fromjson as \$matrix
-              | [
-                  \$ENV.MATRIX_TYPES | fromjson
-                  | .[] as \$matrix_type
-                  | \$matrix[\$ENV.MATRIX_NAME][\$matrix_type][]
-                ]
-              | reduce .[] as \$entry (
-                  [];
-                  if index(\$entry) == null then . + [\$entry] else . end
-                )
-              | ${MATRIX_FILTER}
-              | if (. | length) > 0 then {include: .} else \"\" | halt_error(1) end
-            "
-          )"; then
-            echo "::error::Empty matrix"
-            exit 1
-          fi
-
-          echo "matrix=${MATRIX}" | tee --append "${GITHUB_OUTPUT}"
+        run: bash shared-workflows/.github/actions/prepare-matrix/compute.sh

.github/workflows/conda-python-tests.yaml

@@ -37,6 +37,10 @@ on:
           For example, 'map(select(.ARCH == "amd64"))' to achieve "only run amd64 jobs".
         type: string
         default: "."
+      resolved-matrix:
+        description: Optional JSON include matrix from a previously resolved snapshot. matrix_filter still applies.
+        type: string
+        default: ''
       continue-on-error:
         description: |
           If true, job failures do not result in workflow failures (useful for implementing
@@ -196,6 +200,7 @@ jobs:
       matrix_name: conda-python-tests
       matrix_type: ${{ inputs.matrix_type }}
       matrix_filter: ${{ inputs.matrix_filter }}
+      resolved-matrix: ${{ inputs.resolved-matrix }}
   tests:
     name: ${{ matrix.CUDA_VER }}, ${{ matrix.PY_VER }}, ${{ matrix.ARCH }}, ${{ matrix.LINUX_VER }}, ${{ matrix.GPU }}, ${{ matrix.DRIVER }}-driver, ${{ matrix.DEPENDENCIES }}-deps
     needs: compute-matrix
@@ -263,7 +268,7 @@ jobs:
           extra_attributes: "rapids.PACKAGER=conda,rapids.CUDA_VER=${{ matrix.CUDA_VER }},rapids.PY_VER=${{ matrix.PY_VER }},rapids.ARCH=${{ matrix.ARCH }},rapids.LINUX_VER=${{ matrix.LINUX_VER }},rapids.GPU=${{ matrix.GPU }},rapids.DRIVER=${{ matrix.DRIVER }},rapids.DEPENDENCIES=${{ matrix.DEPENDENCIES }}"
 
       - name: Setup proxy cache
-        uses: nv-gha-runners/setup-proxy-cache@main
+        uses: nv-gha-runners/setup-proxy-cache@eadc57871cde0c92a223d0f486ab669a2360f978 # main
         continue-on-error: true
         with:
           enable-apt: true

.github/workflows/resolve-python-test-matrices.yaml

@@ -0,0 +1,46 @@
+name: Resolve Python test matrix snapshot
+on:
+  workflow_call:
+    inputs:
+      build_type:
+        type: string
+        required: true
+      matrix_type:
+        type: string
+        default: auto
+    outputs:
+      wheels-test:
+        value: ${{ jobs.resolve.outputs.wheels-test }}
+      conda-python-tests:
+        value: ${{ jobs.resolve.outputs.conda-python-tests }}
+permissions:
+  contents: read
+jobs:
+  resolve:
+    runs-on: ubuntu-latest
+    outputs:
+      wheels-test: ${{ steps.wheels.outputs.matrix }}
+      conda-python-tests: ${{ steps.conda.outputs.matrix }}
+    steps:
+      - name: Prepare matrix definitions once
+        id: prepare
+        uses: $/.github/actions/prepare-matrix-dispatch
+        with:
+          action-repo: ${{ github.action_repository }}
+          action-ref: ${{ github.action_ref }}
+      - name: Resolve wheel test matrix
+        id: wheels
+        env:
+          BUILD_TYPE: ${{ inputs.build_type }}
+          MATRIX_TYPE: ${{ inputs.matrix_type }}
+          MATRIX_NAME: wheels-test
+          MATRIX: ${{ steps.prepare.outputs.matrix }}
+        run: bash shared-workflows/.github/actions/prepare-matrix/compute.sh
+      - name: Resolve Conda Python test matrix
+        id: conda
+        env:
+          BUILD_TYPE: ${{ inputs.build_type }}
+          MATRIX_TYPE: ${{ inputs.matrix_type }}
+          MATRIX_NAME: conda-python-tests
+          MATRIX: ${{ steps.prepare.outputs.matrix }}
+        run: bash shared-workflows/.github/actions/prepare-matrix/compute.sh

.github/workflows/test-python-test-matrices.yaml

@@ -0,0 +1,40 @@
+name: Python test matrix resolution
+on:
+  pull_request:
+    paths:
+      - .github/actions/prepare-matrix/**
+      - .github/actions/prepare-matrix-dispatch/**
+      - .github/workflows/compute-matrix.yaml
+      - .github/workflows/resolve-python-test-matrices.yaml
+      - .github/workflows/test-python-test-matrices.yaml
+      - .github/workflows/wheels-test.yaml
+      - .github/workflows/conda-python-tests.yaml
+      - .pre-commit-config.yaml
+      - ci/test_compute_matrix.py
+  push:
+    branches: [main]
+    paths:
+      - .github/actions/prepare-matrix/**
+      - .github/actions/prepare-matrix-dispatch/**
+      - .github/workflows/compute-matrix.yaml
+      - .github/workflows/resolve-python-test-matrices.yaml
+      - .github/workflows/test-python-test-matrices.yaml
+      - .github/workflows/wheels-test.yaml
+      - .github/workflows/conda-python-tests.yaml
+      - .pre-commit-config.yaml
+      - ci/test_compute_matrix.py
+permissions:
+  contents: read
+jobs:
+  matrix-tests:
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          persist-credentials: false
+      - name: Test matrix resolution and snapshot reuse
+        run: |
+          python3 -m venv .matrix-test-env
+          .matrix-test-env/bin/python -m pip install PyYAML==6.0.3
+          .matrix-test-env/bin/python -m unittest discover -s ci -p test_compute_matrix.py

.github/workflows/wheels-test.yaml

@@ -36,6 +36,10 @@ on:
           For example, 'map(select(.ARCH == "amd64"))' to achieve "only run amd64 jobs".
         type: string
         default: "."
+      resolved-matrix:
+        description: Optional JSON include matrix from a previously resolved snapshot. matrix_filter still applies.
+        type: string
+        default: ''
       container-options:
         description: |
           Command-line arguments passed to 'docker run' when starting the container this workflow runs in.
@@ -202,6 +206,7 @@ jobs:
       matrix_name: wheels-test
       matrix_type: ${{ inputs.matrix_type }}
       matrix_filter: ${{ inputs.matrix_filter }}
+      resolved-matrix: ${{ inputs.resolved-matrix }}
   test:
     name: ${{ matrix.CUDA_VER }}, ${{ matrix.PY_VER }}, ${{ matrix.ARCH }}, ${{ matrix.LINUX_VER }}, ${{ matrix.GPU }}, ${{ matrix.DRIVER }}-driver, ${{ matrix.DEPENDENCIES }}-deps
     needs: compute-matrix
@@ -269,7 +274,7 @@ jobs:
           build_workflow_name: ${{ inputs.build_workflow_name }}
 
       - name: Setup proxy cache
-        uses: nv-gha-runners/setup-proxy-cache@main
+        uses: nv-gha-runners/setup-proxy-cache@eadc57871cde0c92a223d0f486ab669a2360f978 # main
         continue-on-error: true
         with:
           enable-apt: true