Enable New bazel presubmits for pull requests.

At this time these workflows will currently be non-blocking for submission. PiperOrigin-RevId: 703311362
jax-ml · Dec 10, 2024 · 9fd37b7 · 9fd37b7
1 parent 9c32fe8
commit 9fd37b7
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/.github/workflows/bazel_cpu_rbe.yml b/.github/workflows/bazel_cpu_rbe.yml
@@ -11,6 +11,9 @@ on:
         options:
         - 'yes'
         - 'no'
+  pull_request:
+    branches:
+      - main
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
@@ -21,15 +24,18 @@ jobs:
     if: github.event.repository.fork == false
     strategy:
       matrix:
-        runner: ["linux-x86-n2-16", "linux-arm64-t2a-16"]
+        runner: ["linux-x86-n2-16", "linux-arm64-c4a-16"]
+        enable-x_64: [1, 0]
 
     runs-on: ${{ matrix.runner }}
-    # TODO(b/369382309): Replace Linux Arm64 container with the ml-build container once it is available
     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/linux-arm64-arc-container: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_ENABLE_X64: ${{ matrix.enable-x_64 }}
+
+    name: "Bazel CPU tests (${{ matrix.runner }}, Python 3.12, x64=${{ matrix.enable-x_64 }})"
 
     steps:
       - uses: actions/checkout@v3

diff --git a/.github/workflows/bazel_gpu_rbe.yml b/.github/workflows/bazel_gpu_rbe.yml
@@ -11,6 +11,9 @@ on:
         options:
         - 'yes'
         - 'no'
+  pull_request:
+    branches:
+      - main
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
@@ -22,12 +25,16 @@ jobs:
     strategy:
       matrix:
         runner: ["linux-x86-n2-16"]
+        enable-x_64: [1, 0]
 
     runs-on: ${{ matrix.runner }}
     container: 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest'
 
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
+      JAXCI_ENABLE_X64: ${{ matrix.enable-x_64 }}
+
+    name: "Bazel single accelerator GPU tests (${{ matrix.runner }}, Python 3.12, x64=${{ matrix.enable-x_64 }})"
 
     steps:
       - uses: actions/checkout@v3