bazel_cpu_rbe.yml

name: CI - Bazel CPU tests (RBE)

on:
  pull_request:
      branches:
        - main
  workflow_dispatch:

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

jobs:
  run_tests:
    if: !github.event.repository.fork
    strategy:
      matrix:
        runner: ["linux-x86-n2-64", "linux-arm64-t2a-48"]

    runs-on: ${{ matrix.runner }}
    container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
                   # TODO(b/369382309): Replace this with the ml-build container once it is available for Linux Arm64
                   (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: Run Bazel CPU Tests with RBE
        run: ./ci/run_bazel_test.sh "ci/envs/run_tests/bazel_cpu_rbe.env"