# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
name: CI
permissions:
  contents: read
on:
  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'
  pull_request:
  push:
    branches:
      - main

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
  cancel-in-progress: ${{ github.ref != 'main' }}

jobs:
  Tests:
    strategy:
      # Don't fail fast - want to see results for all builds even if one fails.
      fail-fast: false
      matrix:
        job_info: [
          {
            pool: "linux-x86-n2-16",
            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
            name: "XLA Linux x86 CPU",
            repo: "openxla/xla",
          },
          {
            pool: "windows-x86-n2-16",
            container: null,
            name: "XLA Windows x86 CPU",
            repo: "openxla/xla",
          },
          {
            pool: "linux-x86-n2-16",
            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
            name: "XLA Linux x86 CPU Bzlmod",
            repo: "openxla/xla",
          },
          {
            pool: "linux-x86-n2-16",
            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
            name: "XLA Linux X86 GPU ONEAPI",
            repo: "openxla/xla",
          },
          {
            pool: "linux-x86-n2-16",
            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
            name: "XLA Linux x86 GPU Hermetic ROCm",
            repo: "openxla/xla",
          },
          {
            pool: "linux-x86-g2-16-l4-1gpu",
            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
            name: "XLA Linux x86 GPU L4",
            repo: "openxla/xla",
          },
          {
            pool: "linux-arm64-c4a-16",
            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64:latest",
            name: "XLA Linux ARM64 CPU",
            repo: "openxla/xla",
          },
          {
            pool: "linux-x86-n2-16",
            container: "gcr.io/tensorflow-sigs/build:latest-python3.11",
            name: "JAX Linux x86 CPU Bzlmod",
            repo: "jax-ml/jax",
          },
          {
            pool: "windows-x86-n2-16",
            container: null,
            name: "JAX Windows x86 CPU",
            repo: "jax-ml/jax",
          },
          {
            pool: "linux-x86-g2-16-l4-1gpu",
            container: "gcr.io/tensorflow-sigs/build:latest-python3.12",
            name: "JAX Linux x86 GPU L4",
            repo: "jax-ml/jax",
          },
          {
            pool: "linux-x86-n2-16",
            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
            name: "TensorFlow Linux x86 CPU",
            repo: "tensorflow/tensorflow",
          },
          {
            pool: "linux-x86-g2-16-l4-1gpu",
            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
            name: "TensorFlow Linux x86 GPU L4",
            repo: "tensorflow/tensorflow",
          },
        ]
    name: ${{ matrix.job_info.name }}
    runs-on: ${{ matrix.job_info.pool }}
    container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
    defaults:
      run:
        shell: bash
    steps:
      - name: "Checking out openxla/xla"
        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
        with:
          path: "openxla/xla"
          persist-credentials: false
      - name: Checking out ${{ matrix.job_info.repo }}
        if: ${{ matrix.job_info.repo != 'openxla/xla' }}
        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
        with:
          repository: ${{ matrix.job_info.repo }}
          path: ${{ matrix.job_info.repo }}
          persist-credentials: false
      - name: Checking out rules_ml_toolchain # TODO(alekstheod):  remove once jax is migrated to the latest rules_ml_toolchain changes
        run: |
          git clone https://github.com/google-ml-infra/rules_ml_toolchain.git openxla/rules_ml_toolchain
          cd openxla/rules_ml_toolchain
          git checkout b11745590f513ec55b32e2d126073576fde18c71
      - name: "Wait For Connection"
        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
        with:
          halt-dispatch-input: ${{ inputs.halt-for-connection }}
      - name: "Run build.py"
        working-directory: ${{ matrix.job_info.repo }}
        timeout-minutes: 60
        run: |
          # Shellcheck doesn't understand GitHub Actions expressions and incorrectly
          # thinks this comparison will always fail.
          # shellcheck disable=SC2193
          if [[ "${{ matrix.job_info.pool }}" == *"windows"* ]]; then
            python "$GITHUB_WORKSPACE"\\openxla\\xla\\build_tools\\ci\\build.py --build="${{ matrix.job_info.name }}_github_actions"
          else
            "$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions"
          fi