Abhinav/elastic snapshot

.github/CODEOWNERS

@@ -1,11 +1,11 @@
-* @abhinavclemson @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @gagika @shralex @SurbhiJainUSC @hengtaoguo @A9isha @aireenmei @NuojCheng @jiangjy1982 @suexu1025 @NicoGrande @jesselu-google @dipannita08 @igorts-git @darisoy @xibinliu
+* @abhinavclemson @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @gagika @shralex @SurbhiJainUSC @hengtaoguo @A9isha @aireenmei @NuojCheng @jiangjy1982 @suexu1025 @NicoGrande @jesselu-google @dipannita08 @igorts-git @darisoy @xibinliu @huytransformer
 
 # Model bring-up
-src/maxtext/assets @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande
-src/maxtext/configs/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng
-src/maxtext/checkpoint_conversion @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @hengtaoguo @gagika @shralex @richjames0 @NicoGrande
-src/maxtext/layers @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng @xibinliu
-src/maxtext/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng @xibinliu
+src/maxtext/assets @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @aireenmei
+src/maxtext/configs/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng @aireenmei
+src/maxtext/checkpoint_conversion @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @hengtaoguo @gagika @shralex @richjames0 @NicoGrande @aireenmei
+src/maxtext/layers @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng @xibinliu @huytransformer @aireenmei
+src/maxtext/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng @xibinliu @huytransformer @aireenmei
 
 # Features
 src/maxtext/experimental/rl @A9isha @khatwanimohit @xuefgu @gagika @richjames0 @shralex @NicoGrande @darisoy
@@ -29,10 +29,10 @@ src/maxtext/inference @vipannalla @mitalisi @gpolovets1 @mailvijayasingh @jrplat
 src/dependencies/ @bvandermoon @SurbhiJainUSC @parambole @richjames0 @shralex @darisoy
 
 # Docs
-docs/ @jacoguzo @bvandermoon @SurbhiJainUSC @richjames0 @shralex @gobbleturk @RissyRan @gagika @A9isha @jiangjy1982 @vipannalla @darisoy
+docs/ @jacoguzo @bvandermoon @SurbhiJainUSC @richjames0 @shralex @gobbleturk @RissyRan @gagika @A9isha @jiangjy1982 @vipannalla @darisoy @aireenmei
 
 # Workflow files
-.github/workflows/ @gobbleturk @khatwanimohit @shralex @parambole @bvandermoon @richjames0 @darisoy @xibinliu
+.github/workflows/ @gobbleturk @khatwanimohit @shralex @parambole @bvandermoon @richjames0 @darisoy @xibinliu @huytransformer
 
 # Benchmarking/Recipes
-benchmarks/ @SujeethJinesh @bvandermoon @richjames0 @shralex @vipannalla @mitalisi @RissyRan @shauryagup @NuojCheng @gobbleturk @khatwanimohit @Obliviour @notabee @suexu1025 @darisoy
+benchmarks/ @SujeethJinesh @bvandermoon @richjames0 @shralex @vipannalla @mitalisi @RissyRan @shauryagup @NuojCheng @gobbleturk @khatwanimohit @Obliviour @notabee @suexu1025 @darisoy @aireenmei

.github/workflows/build_and_push_docker_image.yml

@@ -42,6 +42,10 @@ on:
         required: false
         type: boolean
         default: false
+      run_tests:
+        required: false
+        type: boolean
+        default: true
     outputs:
       build_result:
         description: 'The result of build_and_push job.'
@@ -184,7 +188,7 @@ jobs:
   run_ci_tests:
     name: Run Unit and Integration Tests
     needs: [pre_build_check, build_and_push]
-    if: needs.pre_build_check.outputs.should_run == 'true' && inputs.include_test_assets == true
+    if: needs.pre_build_check.outputs.should_run == 'true' && inputs.include_test_assets == true && inputs.run_tests == true
     uses: ./.github/workflows/run_ci_tests.yml
     with:
       image_name: ${{ inputs.image_name }}

.github/workflows/check_docs_build.yml

@@ -17,12 +17,11 @@
 name: Check the Documentation Build
 
 on:
-  pull_request:
-    types: [opened, synchronize]
   workflow_call:
     inputs:
       maxtext_sha:
-        required: true
+        description: 'The specific MaxText commit SHA.'
+        required: false
         type: string
   workflow_dispatch:
 
@@ -38,7 +37,7 @@ jobs:
         with:
           persist-credentials: false
           fetch-depth: 0
-          ref: ${{ inputs.maxtext_sha }}
+          ref: ${{ inputs.maxtext_sha || github.sha }}
 
       - name: Check if only documentation changed
         id: check

.github/workflows/ci_pipeline.yml

@@ -126,12 +126,23 @@ jobs:
             echo "run_notebooks=false" >> $GITHUB_OUTPUT
           fi
 
+  code_quality_check:
+    uses: ./.github/workflows/code_quality.yml
+    with:
+      maxtext_sha: ${{ inputs.maxtext_sha || github.sha }}
+
+  docs_build_check:
+    uses: ./.github/workflows/check_docs_build.yml
+    with:
+      maxtext_sha: ${{ inputs.maxtext_sha || github.sha }}
+
   build_and_upload_maxtext_package:
-    needs: analyze_code_changes
-    # Run if either tests or notebooks need to run
+    needs: [analyze_code_changes, code_quality_check, docs_build_check]
+    # Run if either tests or notebooks need to run; on PRs, gate on code quality + docs passing
     if: |
-      needs.analyze_code_changes.outputs.run_tests == 'true' ||
-      needs.analyze_code_changes.outputs.run_notebooks == 'true'
+      always() &&
+      (needs.analyze_code_changes.outputs.run_tests == 'true' || needs.analyze_code_changes.outputs.run_notebooks == 'true') &&
+      (github.event_name != 'pull_request' || (needs.code_quality_check.result == 'success' && needs.docs_build_check.result == 'success'))
     uses: ./.github/workflows/build_package.yml
     with:
       device_type: tpu
@@ -154,15 +165,60 @@ jobs:
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
+  gate_test_run:
+    name: Gate and Formalize Parameters
+    needs: [analyze_code_changes, build_and_upload_maxtext_package]
+    if: |
+      always() &&
+      needs.analyze_code_changes.outputs.run_tests == 'true' &&
+      needs.build_and_upload_maxtext_package.result == 'success'
+    runs-on: ubuntu-latest
+    outputs:
+      total_workers: ${{ steps.set-params.outputs.total_workers }}
+      worker_groups: ${{ steps.set-params.outputs.worker_groups }}
+    steps:
+      - id: set-params
+        name: Formalize Test Suite Parameters
+        run: |
+          TPU_UNIT_TOTAL_WORKERS=2
+          TPU_UNIT_WORKER_GROUPS='[1, 2]'
+          echo "total_workers=${TPU_UNIT_TOTAL_WORKERS}" >> "$GITHUB_OUTPUT"
+          echo "worker_groups=${TPU_UNIT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT"
+
   tpu-tests:
     name: ${{ matrix.flavor || 'TPU' }} tests
+    needs: [build_and_upload_maxtext_package, gate_test_run]
+    uses: ./.github/workflows/run_tests_coordinator.yml
+    strategy:
+      fail-fast: false
+      matrix:
+        flavor: [tpu-unit, tpu-integration, tpu-post-training-unit]
+    with:
+      flavor: ${{ matrix.flavor }}
+      base_image: maxtext-unit-test-tpu:py312
+      is_scheduled_run: ${{ github.event_name == 'schedule' }}
+      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
+
+  maxtext_cpu_torch_reference_tests:
+    name: cpu-torch-reference tests
     needs: [build_and_upload_maxtext_package]
     if: needs.analyze_code_changes.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_tests_coordinator.yml
+    with:
+      flavor: cpu-torch-reference
+      base_image: maxtext-unit-test-tpu:py312
+      is_scheduled_run: ${{ github.event_name == 'schedule' }}
+      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
+
+  tpu7x-tests:
+    name: TPU7X tests
+    needs: [build_and_upload_maxtext_package, gate_test_run]
+    if: github.ref == 'refs/heads/main' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch')
+    uses: ./.github/workflows/run_tests_coordinator.yml
     strategy:
       fail-fast: false
       matrix:
-        flavor: [tpu-unit, tpu-integration, tpu-post-training-unit]
+        flavor: [tpu7x-unit, tpu7x-integration, tpu7x-post-training-unit]
     with:
       flavor: ${{ matrix.flavor }}
       base_image: maxtext-unit-test-tpu:py312
@@ -171,8 +227,7 @@ jobs:
 
   gpu-tests:
     name: ${{ matrix.flavor || 'GPU' }} tests
-    needs: [build_and_upload_maxtext_package]
-    if: needs.analyze_code_changes.outputs.run_tests == 'true'
+    needs: [build_and_upload_maxtext_package, gate_test_run]
     strategy:
       fail-fast: false
       matrix:
@@ -186,8 +241,7 @@ jobs:
 
   cpu-tests:
     name: ${{ matrix.flavor || 'CPU' }} tests
-    needs: [build_and_upload_maxtext_package]
-    if: needs.analyze_code_changes.outputs.run_tests == 'true'
+    needs: [build_and_upload_maxtext_package, gate_test_run]
     uses: ./.github/workflows/run_tests_coordinator.yml
     strategy:
       fail-fast: false
@@ -199,31 +253,13 @@ jobs:
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
-  setup-pathways-parameters:
-    name: Setup Pathways Parameters
-    runs-on: ubuntu-latest
-    outputs:
-      total_workers: ${{ steps.set-params.outputs.total_workers }}
-      worker_groups: ${{ steps.set-params.outputs.worker_groups }}
-    steps:
-      - id: set-params
-        name: Set Pathways Worker Parameters
-        run: |
-          # TPU worker constants
-          TPU_UNIT_TOTAL_WORKERS=2
-          TPU_UNIT_WORKER_GROUPS='[1, 2]'
-
-          echo "total_workers=${TPU_UNIT_TOTAL_WORKERS}" >> "$GITHUB_OUTPUT"
-          echo "worker_groups=${TPU_UNIT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT"
-
   maxtext_tpu_pathways_unit_tests:
-    needs: [build_and_upload_maxtext_package, setup-pathways-parameters]
-    if: needs.analyze_code_changes.outputs.run_tests == 'true'
+    needs: [build_and_upload_maxtext_package, gate_test_run]
     uses: ./.github/workflows/run_pathways_tests.yml
     strategy:
         fail-fast: false
         matrix:
-          group: ${{ fromJSON(needs.setup-pathways-parameters.outputs.worker_groups) }}
+          group: ${{ fromJSON(needs.gate_test_run.outputs.worker_groups || '[1, 2]') }}
     with:
       device_type: tpu
       device_name: v6e-4
@@ -236,11 +272,11 @@ jobs:
       container_resource_option: "--privileged"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
-      total_workers: ${{ needs.setup-pathways-parameters.outputs.total_workers }}
+      total_workers: ${{ needs.gate_test_run.outputs.total_workers || '2' }}
       worker_group: ${{ matrix.group }}
 
   maxtext_tpu_pathways_integration_tests:
-    needs: build_and_upload_maxtext_package
+    needs: [build_and_upload_maxtext_package, gate_test_run]
     uses: ./.github/workflows/run_pathways_tests.yml
     strategy:
         fail-fast: false
@@ -259,20 +295,20 @@ jobs:
 
   all_tests_passed:
     name: All Required Tests Passed
-    needs: [tpu-tests, gpu-tests, cpu-tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests]
+    needs: [build_and_upload_maxtext_package, gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_cpu_torch_reference_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check]
     if: always()
     runs-on: ubuntu-latest
     steps:
       - name: Check test results
         run: |
-          if [ "${NEEDS_ANALYZE_CHANGES_OUTPUTS_RUN_TESTS}" == "false" ]; then
-            echo "Tests were skipped"
-            exit 0
-          fi
-
-          # Otherwise, check that build and all tests passed or were skipped
+          # Check that build and all tests passed or were skipped
           echo "Build result: ${NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT}"
+          echo "Code Quality result: ${NEEDS_CODE_QUALITY_CHECK_RESULT}"
+          echo "Docs Build result: ${NEEDS_DOCS_BUILD_CHECK_RESULT}"
+          echo "Gate result: ${NEEDS_GATE_TEST_RUN_RESULT}"
           echo "TPU Tests (Matrix) result: ${NEEDS_TPU_TESTS_RESULT}"
+          echo "TPU7X Tests (Matrix) result: ${NEEDS_TPU7X_TESTS_RESULT}"
+          echo "CPU torch reference tests result: ${NEEDS_MAXTEXT_CPU_TORCH_REFERENCE_TESTS_RESULT}"
           echo "GPU Tests (Matrix) result: ${NEEDS_GPU_TESTS_RESULT}"
           echo "CPU Tests (Matrix) result: ${NEEDS_CPU_TESTS_RESULT}"
           echo "Pathways Unit result: ${NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT}"
@@ -286,10 +322,14 @@ jobs:
 
           echo "All required tests passed successfully"
         env:
-          NEEDS_ANALYZE_CODE_CHANGES_OUTPUTS_RUN_TESTS: ${{ needs.analyze_code_changes.outputs.run_tests }}
           NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT: ${{ needs.build_and_upload_maxtext_package.result }}
+          NEEDS_CODE_QUALITY_CHECK_RESULT: ${{ needs.code_quality_check.result }}
+          NEEDS_DOCS_BUILD_CHECK_RESULT: ${{ needs.docs_build_check.result }}
+          NEEDS_GATE_TEST_RUN_RESULT: ${{ needs.gate_test_run.result }}
           NEEDS_CPU_TESTS_RESULT: ${{ needs.cpu-tests.result }}
           NEEDS_TPU_TESTS_RESULT: ${{ needs.tpu-tests.result }}
+          NEEDS_TPU7X_TESTS_RESULT: ${{ needs.tpu7x-tests.result }}
+          NEEDS_MAXTEXT_CPU_TORCH_REFERENCE_TESTS_RESULT: ${{ needs.maxtext_cpu_torch_reference_tests.result }}
           NEEDS_GPU_TESTS_RESULT: ${{ needs.gpu-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 }}
@@ -325,7 +365,7 @@ jobs:
 
   notify_failure:
     name: Notify failed build # creates an issue or modifies last open existing issue for failed build
-    needs: [tpu-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests]
+    needs: [gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_cpu_torch_reference_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check]
     if: ${{ always() }}
     runs-on: ubuntu-latest
     permissions:
@@ -339,7 +379,7 @@ jobs:
 
   investigate_failure:
     name: Investigate failed build # investigates failure of scheduled run and comments on tracking issue
-    needs: [tpu-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, notify_failure]
+    needs: [gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_cpu_torch_reference_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check, notify_failure]
     if: ${{ always() && contains(needs.*.result, 'failure') && github.event_name == 'schedule' }}
     uses: ./.github/workflows/gemini_investigate.yml
     permissions:

.github/workflows/code_quality.yml

@@ -17,24 +17,22 @@
 name: Code Quality
 
 on:
-  pull_request:
-
-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
+  workflow_call:
+    inputs:
+      maxtext_sha:
+        description: 'The specific MaxText commit SHA.'
+        required: false
+        type: string
 
 jobs:
   qa:
     name: "Static code-quality checkers"
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v5
+      with:
+        fetch-depth: 0
+        ref: ${{ inputs.maxtext_sha || github.sha }}
 
     - name: Install uv and set the Python version
       uses: astral-sh/setup-uv@v7
@@ -56,7 +54,12 @@ jobs:
 
     - name: Run pre-commit checks on just the files that have changed
       run: |
-        git fetch origin "$GITHUB_BASE_REF":"$GITHUB_BASE_REF"
-        git branch "$GITHUB_HEAD_REF"
         . "$GITHUB_WORKSPACE"/venv/bin/activate
-        pre-commit run --from-ref "$GITHUB_BASE_REF" --to-ref "$GITHUB_HEAD_REF" --show-diff-on-failure
+        if [ "${{ github.event_name }}" == "pull_request" ]; then
+          git fetch origin "$GITHUB_BASE_REF":"$GITHUB_BASE_REF"
+          git branch "$GITHUB_HEAD_REF"
+          pre-commit run --from-ref "$GITHUB_BASE_REF" --to-ref "$GITHUB_HEAD_REF" --show-diff-on-failure
+        else
+          echo "Not a pull request event (${{ github.event_name }}), checking all files."
+          pre-commit run --all-files --show-diff-on-failure
+        fi