update workflows after sync to upstream

Author: nitins17

.github/workflows/asan.yaml

@@ -50,6 +50,9 @@ jobs:
           make -j64
           make install
           ${GITHUB_WORKSPACE}/cpythonasan/bin/python3 -m venv ${GITHUB_WORKSPACE}/venv
+      # Halt for testing
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@main
       - name: Install JAX test requirements
         env:
           ASAN_OPTIONS: detect_leaks=0

.github/workflows/bazel_cpu_rbe.yml

@@ -1,9 +1,9 @@
 name: CI - Bazel CPU tests (RBE)
 
 on:
-  # pull_request:
-  #   branches:
-  #     - main
+  pull_request:
+    branches:
+      - main
   workflow_dispatch:
     inputs:
       halt-for-connection:
@@ -29,21 +29,24 @@ jobs:
       fail-fast: false # don't cancel all jobs on failure
       matrix:
         runner: ["windows-x86-n2-64", "linux-x86-n2-16", "linux-arm64-t2a-16"]
+        enable-x_64: [1, 0]
 
     runs-on: ${{ matrix.runner }}
     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') ||
                    (contains(matrix.runner, 'windows-x86') && null) }}
 
-
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
+      JAXCI_ENABLE_X64: ${{ matrix.enable-x_64 }}
 
+    name: "Run Bazel CPU tests (x64=${{ matrix.enable-x_64 }}"
+      
     steps:
-      - uses: actions/checkout@v3
-      - name: Wait For Connection
-        uses: google-ml-infra/actions/ci_connection@main
-        with:
-          halt-dispatch-input: ${{ inputs.halt-for-connection }}
-      - name: Run Bazel CPU Tests with RBE
-        run: ./ci/run_bazel_test_cpu_rbe.sh
+    - uses: actions/checkout@v3
+    - name: Wait For Connection
+      uses: google-ml-infra/actions/ci_connection@main
+      with:
+        halt-dispatch-input: ${{ inputs.halt-for-connection }}
+    - name: Run Bazel CPU Tests with RBE
+      run: ./ci/run_bazel_test_cpu_rbe.sh

.github/workflows/bazel_gpu_non_rbe.yml

@@ -32,13 +32,18 @@ jobs:
       upload_destination_prefix: '${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'
 
   run_bazel_tests:
-    name: "Run Bazel GPU tests (single accelerator and multi-accelerator tests, non-RBE)"
     needs: build_artifacts
     runs-on: "linux-x86-g2-48-l4-4gpu"
     container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
-
+    strategy:
+      fail-fast: false # don't cancel all jobs on failure
+      matrix:
+        enable-x_64: [1, 0]
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: 3.11
+      JAXCI_ENABLE_X64: ${{ matrix.enable-x_64 }}
+
+    name: "Run Bazel GPU tests (single accelerator and multi-accelerator tests, non-RBE), x64=${{ matrix.enable-x_64 }}"
 
     steps:
       - uses: actions/checkout@v3

.github/workflows/bazel_gpu_rbe.yml

@@ -1,9 +1,9 @@
 name: CI - Bazel GPU tests (RBE)
 
 on:
-  # pull_request:
-  #   branches:
-  #     - main
+  pull_request:
+    branches:
+      - main
   workflow_dispatch:
     inputs:
       halt-for-connection:
@@ -24,13 +24,17 @@ 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: "Run Bazel GPU tests (x64=${{ matrix.enable-x_64 }}"
+  
     steps:
       - uses: actions/checkout@v3
       - name: Wait For Connection

.github/workflows/build_artifacts.yml

@@ -159,6 +159,7 @@ jobs:
 
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
+      JAXCI_BUILD_ARTIFACT_WITH_RBE: 1
       JAXCI_CLONE_MAIN_XLA: "${{ inputs.clone_main_xla }}"
 
     steps: