Update the matrix for the RBE presubmits

PiperOrigin-RevId: 721900273
jax-ml · Jan 31, 2025 · ef9046d · ef9046d
1 parent 3f73758
commit ef9046d
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/.github/workflows/bazel_cpu_rbe.yml b/.github/workflows/bazel_cpu_rbe.yml
@@ -26,14 +26,20 @@ jobs:
     container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
                    (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest') }}
     env:
-      JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
+      JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
       JAXCI_ENABLE_X64: ${{ matrix.enable-x_64 }}
 # Begin Presubmit Naming Check - name modification requires internal check to be updated
     strategy:
       matrix:
+        python: ["3.10", "3.13"]
         runner: ["linux-x86-n2-16", "linux-arm64-c4a-16"]
         enable-x_64: [1, 0]
-    name: "Bazel CPU tests (${{ matrix.runner }}, Python 3.12, x64=${{ matrix.enable-x_64 }})"
+        exclude:
+            - python: "3.10"
+              enable-x_64: 1
+            - python: "3.13"
+              enable-x_64: 0
+    name: "Bazel CPU tests (${{ matrix.runner }}, Python ${{ matrix.python }}, x64=${{ matrix.enable-x_64 }})"
 # End Presubmit Naming Check github-cpu-presubmits
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2

diff --git a/.github/workflows/bazel_cuda_rbe.yml b/.github/workflows/bazel_cuda_rbe.yml
@@ -25,14 +25,20 @@ jobs:
     runs-on: ${{ matrix.runner }}
     container: 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest'
     env:
-      JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
+      JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
       JAXCI_ENABLE_X64: ${{ matrix.enable-x_64 }}
 # Begin Presubmit Naming Check - name modification requires internal check to be updated
     strategy:
       matrix:
+        python: ["3.10", "3.13"]
         runner: ["linux-x86-n2-16"]
         enable-x_64: [1, 0]
-    name: "Bazel single accelerator CUDA tests (${{ matrix.runner }}, Python 3.12, x64=${{ matrix.enable-x_64 }})"
+        exclude:
+            - python: "3.10"
+              enable-x_64: 1
+            - python: "3.13"
+              enable-x_64: 0
+    name: "Bazel single accelerator CUDA tests (${{ matrix.runner }}, Python ${{ matrix.python }}, x64=${{ matrix.enable-x_64 }})"
 # End Presubmit Naming Check github-cuda-presubmits
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2