ci.yaml

name: CI

on:
  pull_request:
    paths-ignore:
      - '**.md'
  workflow_dispatch:
    inputs:
      CUDA_IMAGE:
        type: string
        description: 'Base CUDA image, e.g. nvidia/cuda:X.Y.Z-devel-ubuntu22.04'
        required: false
        default: 'latest'
      SRC_JAX:
        description: 'JAX source: <repo>#<branch|tag|commit>'
        type: string
        required: true
        default: 'https://github.com/google/jax.git#main'
      SRC_XLA:
        description: 'XLA source: <repo>#<branch|tag|commit>'
        type: string
        required: true
        default: 'https://github.com/openxla/xla.git#main'
      SRC_TE:
        description: 'TE source: <repo>#<branch|tag|commit>'
        type: string
        required: true
        # TODO: This is a temporary pinning of TE as the API in TE no longer matches the TE patch
        #   This should be reverted to main ASAP
        default: 'https://github.com/NVIDIA/TransformerEngine.git#7976bd003fcf084dd068069b92a9a79b1743316a'
      SRC_T5X:
        description: 'T5X source: <repo>#<branch|tag|commit>'
        type: string
        required: true
        default: 'https://github.com/google-research/t5x.git#main'
      SRC_PAXML:
        description: 'Paxml source: <repo>#<branch|tag|commit>'
        type: string
        required: true
        default: 'https://github.com/google/paxml.git#main'
      SRC_PRAXIS:
        description: 'Praxis source: <repo>#<branch|tag|commit>'
        type: string
        required: true
        default: 'https://github.com/google/praxis.git#main'

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

permissions:
  contents: read  # to fetch code
  actions:  write # to cancel previous workflows
  packages: write # to upload container

jobs:

  amd64:
    uses: ./.github/workflows/_ci.yaml
    with:
      ARCHITECTURE: amd64
      CUDA_IMAGE: ${{ inputs.CUDA_IMAGE || 'latest' }}
      SRC_JAX: ${{ inputs.SRC_JAX || 'https://github.com/google/jax.git#main' }}
      SRC_XLA: ${{ inputs.SRC_XLA || 'https://github.com/openxla/xla.git#main'}}
      SRC_TE: ${{ inputs.SRC_TE || 'https://github.com/NVIDIA/TransformerEngine.git#7976bd003fcf084dd068069b92a9a79b1743316a'}}
      SRC_T5X: ${{ inputs.SRC_T5X || 'https://github.com/google-research/t5x.git#main'}}
      SRC_PAXML: ${{ inputs.SRC_PAXML || 'https://github.com/google/paxml.git#main'}}
      SRC_PRAXIS: ${{ inputs.SRC_PRAXIS || 'https://github.com/google/praxis.git#main'}}
    secrets: inherit
 
  arm64:
    uses: ./.github/workflows/_ci.yaml
    with:
      ARCHITECTURE: arm64
      CUDA_IMAGE: ${{ inputs.CUDA_IMAGE || 'latest' }}
      SRC_JAX: ${{ inputs.SRC_JAX || 'https://github.com/google/jax.git#main' }}
      SRC_XLA: ${{ inputs.SRC_XLA || 'https://github.com/openxla/xla.git#main'}}
      SRC_TE: ${{ inputs.SRC_TE || 'https://github.com/NVIDIA/TransformerEngine.git#7976bd003fcf084dd068069b92a9a79b1743316a'}}
      SRC_T5X: ${{ inputs.SRC_T5X || 'https://github.com/google-research/t5x.git#main'}}
      SRC_PAXML: ${{ inputs.SRC_PAXML || 'https://github.com/google/paxml.git#main'}}
      SRC_PRAXIS: ${{ inputs.SRC_PRAXIS || 'https://github.com/google/praxis.git#main'}}
    secrets: inherit

  finalize:
    needs: [amd64, arm64]
    if: always()
    uses: ./.github/workflows/_finalize.yaml
    with:
      PUBLISH_BADGE: false
    secrets: inherit