bazel_cpu_rbe.yml

name: CI - Bazel CPU tests (RBE)

# TODO(srnitin): Do not submit without removing pull_request event. 
on:
  pull_request:
      branches:
        - main
  workflow_dispatch:
    inputs:
      halt-for-connection:
        description: 'Should this workflow run wait for a remote connection?'
        type: choice
        required: true
        default: 'no'
        options:
        - 'yes'
        - 'no'

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
  cancel-in-progress: true

jobs:
  run_tests:
    if: github.event.repository.fork == false
    strategy:
      matrix:
        runner: ["linux-x86-n2-16", "linux-arm64-t2a-16"]

    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') }}

    env:
      JAXCI_HERMETIC_PYTHON_VERSION: "3.12"

    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.sh "ci/envs/run_tests/bazel_cpu_rbe.env"