build_and_test_maxtext.yml

# Copyright 2025 Google LLC

# 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

#     https://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.

# This workflow will build maxtext python package and run tests.

name: MaxText Package Tests

on:
  pull_request:
  workflow_call:
  workflow_dispatch:
  schedule:
    # Run the job every 4 hours
    - cron:  '0 */4 * * *'

concurrency:
  # Dedup pull requests (canceling previous runs of the same workflow for same PR), and scheduled runs but nothing else
  group: >
    ${{
      github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) ||
      github.event_name == 'schedule' && format('{0}-schedule', github.workflow) ||
      github.run_id
    }}
  cancel-in-progress: true

permissions:
  contents: read
jobs:
  doc_only_check:
    name: Check for Documentation-Only Changes
    runs-on: ubuntu-latest
    outputs:
      run_tests: ${{ steps.check.outputs.run_tests }}
      run_notebooks: ${{ steps.check.outputs.run_notebooks }}
    steps:
      - uses: actions/checkout@v4
        with:
          fetch-depth: 0
      - name: Check if only documentation changed
        id: check
        run: |
          if [ "${{ github.event_name }}" != "pull_request" ]; then
            echo "Not a pull request, running all tests"
            echo "run_tests=true" >> $GITHUB_OUTPUT
            echo "run_notebooks=true" >> $GITHUB_OUTPUT
            exit 0
          fi

          git fetch origin ${GITHUB_BASE_REF}
          CHANGED_FILES=$(git diff --name-only origin/${GITHUB_BASE_REF}...HEAD)

          echo "Changed files:"
          echo "$CHANGED_FILES"

          if [ -z "$CHANGED_FILES" ]; then
            echo "No files changed"
            echo "run_tests=true" >> $GITHUB_OUTPUT
            echo "run_notebooks=true" >> $GITHUB_OUTPUT
            exit 0
          fi

          # Check for source code changes (anything not .md and not .ipynb)
          if echo "$CHANGED_FILES" | grep -v -E '\.(md|ipynb)$' > /dev/null; then
            echo "Source code files changed, enabling unit tests."
            echo "run_tests=true" >> $GITHUB_OUTPUT
          else
            echo "No source code changes, skipping unit tests."
            echo "run_tests=false" >> $GITHUB_OUTPUT
          fi

          # Check for notebook (.ipynb) changes specifically
          if echo "$CHANGED_FILES" | grep '\.ipynb$' > /dev/null; then
            echo "Notebook files changed, enabling notebook run."
            echo "run_notebooks=true" >> $GITHUB_OUTPUT
          else
            echo "No notebook changes, skipping notebook run."
            echo "run_notebooks=false" >> $GITHUB_OUTPUT
          fi

          exit 0

  build_and_upload_maxtext_package:
    needs: doc_only_check
    # Run if either tests or notebooks need to run
    if: |
      needs.doc_only_check.outputs.run_tests == 'true' || 
      needs.doc_only_check.outputs.run_notebooks == 'true'
    uses: ./.github/workflows/build_package.yml
    with:
      device_type: tpu
      device_name: v4-8
      cloud_runner: linux-x86-n2-16-buildkit

  maxtext_jupyter_notebooks:
    needs: build_and_upload_maxtext_package
    if: needs.doc_only_check.outputs.run_notebooks == 'true'
    uses: ./.github/workflows/run_jupyter_notebooks.yml
    strategy:
        fail-fast: false
        matrix:
          image_type: ["py312"]
    with:
      device_type: tpu
      device_name: v6e-4
      image_type: ${{ matrix.image_type }}
      cloud_runner: linux-x86-ct6e-180-4tpu
      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
    secrets:
      HF_TOKEN: ${{ secrets.HF_TOKEN }}

  maxtext_cpu_unit_tests:
    needs: build_and_upload_maxtext_package
    if: needs.doc_only_check.outputs.run_tests == 'true'
    uses: ./.github/workflows/run_tests_against_package.yml
    strategy:
        fail-fast: false # don't cancel all jobs on failure
        matrix:
          image_type: ["py312"]
          worker_group: [1, 2]
    with:
      device_type: cpu
      device_name: X64
      cloud_runner: linux-x86-n2-16
      image_type: ${{ matrix.image_type }}
      pytest_marker: 'cpu_only'
      xla_python_client_mem_fraction: 0.75
      tf_force_gpu_allow_growth: false
      container_resource_option: "--privileged"
      is_scheduled_run: ${{ github.event_name == 'schedule' }}
      worker_group: ${{ matrix.worker_group }}
      total_workers: 2
      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}

  maxtext_tpu_unit_tests:
    needs: build_and_upload_maxtext_package
    if: needs.doc_only_check.outputs.run_tests == 'true'
    uses: ./.github/workflows/run_tests_against_package.yml
    strategy:
        fail-fast: false
        matrix:
          image_type: ["py312"]
    with:
      device_type: tpu
      device_name: v6e-4
      image_type: ${{ matrix.image_type }}
      cloud_runner: linux-x86-ct6e-180-4tpu
      pytest_marker: 'not cpu_only and not gpu_only and not integration_test'
      xla_python_client_mem_fraction: 0.75
      tf_force_gpu_allow_growth: false
      container_resource_option: "--privileged"
      is_scheduled_run: ${{ github.event_name == 'schedule' }}
      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}

  maxtext_tpu_integration_tests:
    needs: build_and_upload_maxtext_package
    if: needs.doc_only_check.outputs.run_tests == 'true'
    uses: ./.github/workflows/run_tests_against_package.yml
    strategy:
        fail-fast: false
        matrix:
          image_type: ["py312"]
    with:
      device_type: tpu
      device_name: v6e-4
      image_type: ${{ matrix.image_type }}
      cloud_runner: linux-x86-ct6e-180-4tpu
      pytest_marker: 'not cpu_only and not gpu_only and integration_test'
      xla_python_client_mem_fraction: 0.75
      tf_force_gpu_allow_growth: false
      container_resource_option: "--privileged"
      is_scheduled_run: ${{ github.event_name == 'schedule' }}
      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}

  maxtext_tpu_pathways_unit_tests:
    needs: build_and_upload_maxtext_package
    if: needs.doc_only_check.outputs.run_tests == 'true'
    uses: ./.github/workflows/run_pathways_tests.yml
    strategy:
        fail-fast: false
        matrix:
          image_type: ["py312"]
    with:
      device_type: tpu
      device_name: v6e-4
      image_type: ${{ matrix.image_type }}
      cloud_runner: linux-x86-ct6e-180-4tpu
      pytest_marker: 'not cpu_only and not gpu_only and not integration_test'
      xla_python_client_mem_fraction: 0.75
      tf_force_gpu_allow_growth: false
      container_resource_option: "--privileged"
      is_scheduled_run: ${{ github.event_name == 'schedule' }}
      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}

  maxtext_tpu_pathways_integration_tests:
    needs: build_and_upload_maxtext_package
    if: needs.doc_only_check.outputs.run_tests == 'true'
    uses: ./.github/workflows/run_pathways_tests.yml
    strategy:
        fail-fast: false
        matrix:
          image_type: ["py312"]
    with:
      device_type: tpu
      device_name: v6e-4
      image_type: ${{ matrix.image_type }}
      cloud_runner: linux-x86-ct6e-180-4tpu
      pytest_marker: 'not cpu_only and not gpu_only and integration_test'
      xla_python_client_mem_fraction: 0.75
      tf_force_gpu_allow_growth: false
      container_resource_option: "--privileged"
      is_scheduled_run: ${{ github.event_name == 'schedule' }}
      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}

  maxtext_gpu_unit_tests:
    needs: build_and_upload_maxtext_package
    if: needs.doc_only_check.outputs.run_tests == 'true'
    uses: ./.github/workflows/run_tests_against_package.yml
    strategy:
        fail-fast: false
        matrix:
          image_type: ["py312"]
          cuda: ["cuda12"]
    with:
      device_type: ${{ matrix.cuda }}
      device_name: a100-40gb-4
      image_type: ${{ matrix.image_type }}
      cloud_runner: linux-x86-a2-48-a100-4gpu
      pytest_marker: 'not cpu_only and not tpu_only and not integration_test'
      xla_python_client_mem_fraction: 0.65
      tf_force_gpu_allow_growth: true
      container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged"
      is_scheduled_run: ${{ github.event_name == 'schedule' }}
      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}

  maxtext_gpu_integration_tests:
    needs: build_and_upload_maxtext_package
    if: needs.doc_only_check.outputs.run_tests == 'true'
    uses: ./.github/workflows/run_tests_against_package.yml
    strategy:
        fail-fast: false
        matrix:
          image_type: ["py312"]
          cuda: ["cuda12"]
    with:
      device_type: ${{ matrix.cuda }}
      device_name: a100-40gb-4
      image_type: ${{ matrix.image_type }}
      cloud_runner: linux-x86-a2-48-a100-4gpu
      pytest_marker: 'not cpu_only and not tpu_only and integration_test'
      xla_python_client_mem_fraction: 0.65
      tf_force_gpu_allow_growth: true
      container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged"
      is_scheduled_run: ${{ github.event_name == 'schedule' }}
      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}

  all_tests_passed:
    name: All Required Tests Passed
    needs: [doc_only_check, build_and_upload_maxtext_package, maxtext_cpu_unit_tests, maxtext_tpu_unit_tests, maxtext_tpu_integration_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, maxtext_gpu_unit_tests, maxtext_gpu_integration_tests]
    if: always()
    runs-on: ubuntu-latest
    steps:
      - name: Check test results
        run: |
          # If doc-only, all tests should be skipped
          if [ "${NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_TESTS}" == "false" ]; then
            echo "Documentation-only changes detected, tests were skipped"
            exit 0
          fi

          # Otherwise, check that build and all tests passed or were skipped
          echo "Build result: ${NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT}"
          echo "CPU tests: ${NEEDS_MAXTEXT_CPU_UNIT_TESTS_RESULT}"
          echo "TPU tests: ${NEEDS_MAXTEXT_TPU_UNIT_TESTS_RESULT}"
          echo "TPU integration: ${NEEDS_MAXTEXT_TPU_INTEGRATION_TESTS_RESULT}"
          echo "TPU pathways: ${NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT}"
          echo "TPU pathways integration: ${NEEDS_MAXTEXT_TPU_PATHWAYS_INTEGRATION_TESTS_RESULT}"
          echo "GPU tests: ${NEEDS_MAXTEXT_GPU_UNIT_TESTS_RESULT}"
          echo "GPU integration: ${NEEDS_MAXTEXT_GPU_INTEGRATION_TESTS_RESULT}"

          # Fail only if any job failed or was cancelled (skipped is OK)
          if [ "${{ contains(needs.*.result, 'failure') }}" == "true" ] || [ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]; then
            echo "One or more jobs failed or were cancelled"
            exit 1
          fi

          echo "All required tests passed successfully"
        env:
          NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_TESTS: ${{ needs.doc_only_check.outputs.run_tests }}
          NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT: ${{ needs.build_and_upload_maxtext_package.result }}
          NEEDS_MAXTEXT_CPU_UNIT_TESTS_RESULT: ${{ needs.maxtext_cpu_unit_tests.result }}
          NEEDS_MAXTEXT_TPU_UNIT_TESTS_RESULT: ${{ needs.maxtext_tpu_unit_tests.result }}
          NEEDS_MAXTEXT_TPU_INTEGRATION_TESTS_RESULT: ${{ needs.maxtext_tpu_integration_tests.result }}
          NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT: ${{ needs.maxtext_tpu_pathways_unit_tests.result }}
          NEEDS_MAXTEXT_TPU_PATHWAYS_INTEGRATION_TESTS_RESULT: ${{ needs.maxtext_tpu_pathways_integration_tests.result }}
          NEEDS_MAXTEXT_GPU_UNIT_TESTS_RESULT: ${{ needs.maxtext_gpu_unit_tests.result }}
          NEEDS_MAXTEXT_GPU_INTEGRATION_TESTS_RESULT: ${{ needs.maxtext_gpu_integration_tests.result }}

  all_notebooks_passed:
    name: All Notebooks Passed
    needs: [doc_only_check, build_and_upload_maxtext_package, maxtext_jupyter_notebooks]
    if: always()
    runs-on: ubuntu-latest
    steps:
      - name: Check notebooks results
        run: |
          if [ "${NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_NOTEBOOKS}" == "false" ]; then
            echo "Non-notebook changes detected, runs were skipped"
            exit 0
          fi

          # Otherwise, check that build and notebooks run passed or were skipped
          echo "Build result: ${NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT}"
          echo "Jupyter Notebooks result: ${NEEDS_MAXTEXT_JUPYTER_NOTEBOOKS_RESULT}"

          # Fail only if any job failed or was cancelled (skipped is OK)
          if [ "${{ contains(needs.*.result, 'failure') }}" == "true" ] || [ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]; then
            echo "One or more jobs failed or were cancelled"
            exit 1
          fi

          echo "All required notebooks passed successfully"
        env:
          NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_NOTEBOOKS: ${{ needs.doc_only_check.outputs.run_notebooks }}
          NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT: ${{ needs.build_and_upload_maxtext_package.result }}
          NEEDS_MAXTEXT_JUPYTER_NOTEBOOKS_RESULT: ${{ needs.maxtext_jupyter_notebooks.result }}

  notify_failure:
    name: Notify failed build # creates an issue or modifies last open existing issue for failed build
    needs: [maxtext_jupyter_notebooks, maxtext_cpu_unit_tests, maxtext_tpu_unit_tests, maxtext_tpu_integration_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, maxtext_gpu_unit_tests, maxtext_gpu_integration_tests]
    if: ${{ always() }}
    runs-on: ubuntu-latest
    permissions:
      issues: write
    steps:
    - name: Check whether one of the jobs failed
      if: ${{ contains(needs.*.result, 'failure') && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }}
      uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b  # v1.2.0
      with:
        github-token: ${{ secrets.GITHUB_TOKEN }}