Resolve merge conflicts: merge main into PR #4911 with ring_axis support

Merged main branch which simplified dcn_parallelism handling. Kept ring_axis
logic from PR while using the cleaner config.dcn_parallelism.copy() from main.

Co-Authored-By: Jeremy <jeremy@systalyze.com>
Co-Authored-By: Sudarsanan <sudarsanan@systalyze.com>
Co-Authored-By: Armin <armin@systalyze.com>

Author: 4 people

.github/scripts/analyze_code_changes.sh

@@ -112,6 +112,7 @@ EXCLUDED_FILES=(
   '^\.github/scripts/'
   '^tools/'
   '\.md$'
+  'src/maxtext/version.py$'
 )
 
 # Loop through every changed file

.github/workflows/build_package.yml

@@ -43,6 +43,7 @@ jobs:
   build_and_upload:
     name: Build Wheel
     runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
+    timeout-minutes: 15
     container: python:3.12-slim-bookworm
     outputs:
       maxtext_sha: ${{ steps.vars.outputs.maxtext_sha }}

.github/workflows/check_docs_build.yml

@@ -32,6 +32,7 @@ jobs:
   build-docs:
     name: Build Sphinx Docs
     runs-on: ubuntu-latest
+    timeout-minutes: 15
     steps:
       - name: Checkout repository
         uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3  # v6.0.0

.github/workflows/ci_pipeline.yml

@@ -51,6 +51,7 @@ jobs:
   analyze_code_changes:
     name: Analyze Code Changes for Test Orchestration
     runs-on: ubuntu-latest
+    timeout-minutes: 15
     outputs:
       run_tests: ${{ steps.check.outputs.run_tests }}
       run_notebooks: ${{ steps.check.outputs.run_notebooks }}
@@ -128,6 +129,7 @@ jobs:
       needs.analyze_code_changes.outputs.run_tests == 'true' &&
       needs.build_and_upload_maxtext_package.result == 'success'
     runs-on: ubuntu-latest
+    timeout-minutes: 15
     outputs:
       total_workers: ${{ steps.set-params.outputs.total_workers }}
       worker_groups: ${{ steps.set-params.outputs.worker_groups }}
@@ -183,8 +185,7 @@ jobs:
     needs: [gate_test_run]
     if: |
       always() &&
-      needs.gate_test_run.result == 'success' &&
-      github.ref == 'refs/heads/main' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch')
+      needs.gate_test_run.result == 'success' && github.event_name != 'pull_request'
     uses: ./.github/workflows/run_tests_coordinator.yml
     strategy:
       fail-fast: false
@@ -305,6 +306,7 @@ jobs:
     needs: [build_and_upload_maxtext_package, gate_test_run, tpu-pretrain-tests, tpu-posttrain-tests, tpu7x-tests, gpu-tests, cpu-pretrain-tests, cpu-posttrain-tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check]
     if: always()
     runs-on: ubuntu-latest
+    timeout-minutes: 15
     steps:
       - name: Check test results
         run: |
@@ -348,6 +350,7 @@ jobs:
     needs: [analyze_code_changes, build_and_upload_maxtext_package, maxtext_jupyter_notebooks]
     if: always()
     runs-on: ubuntu-latest
+    timeout-minutes: 15
     steps:
       - name: Check notebooks results
         run: |
@@ -377,6 +380,7 @@ jobs:
     needs: [gate_test_run, tpu-pretrain-tests, tpu-posttrain-tests, tpu7x-tests, gpu-tests, cpu-pretrain-tests, cpu-posttrain-tests, maxtext_jupyter_notebooks, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check]
     if: ${{ always() }}
     runs-on: ubuntu-latest
+    timeout-minutes: 15
     permissions:
       issues: write
     steps:

.github/workflows/code_quality.yml

@@ -28,6 +28,7 @@ jobs:
   qa:
     name: "Pre-commit Linters"
     runs-on: ubuntu-latest
+    timeout-minutes: 15
     steps:
     - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5.1.0
       with:

.github/workflows/promote_docker_image.yml

@@ -18,6 +18,7 @@
 # 2. If the DAG succeeded, it promotes the Docker images by tagging the artifact with the latest tag.
 
 name: Promote MaxText Docker Image
+run-name: "Triggered by Airflow ${{ github.event.client_payload.test_type }} Tests DAG"
 
 on:
   repository_dispatch:
@@ -39,11 +40,11 @@ jobs:
           DAG_ID: ${{ github.event.client_payload.dag_id }}
           DAG_RUN_ID: ${{ github.event.client_payload.dag_run_id }}
           SHA: ${{ github.event.client_payload.sha }}
-          GITHUB_RUN_ID: ${{ github.event.client_payload.github_run_id }}
+          CALLER_RUN_ID: ${{ github.event.client_payload.github_run_id }}
           TEST_TYPE: ${{ github.event.client_payload.test_type }}
         run: |
           echo "================================"
-          echo "Github Run ID: ${GITHUB_RUN_ID}"
+          echo "Github Run ID: ${CALLER_RUN_ID}"
           echo "DAG ID:        ${DAG_ID}"
           echo "DAG Run ID:    ${DAG_RUN_ID}"
           echo "Commit SHA:    ${SHA}"
@@ -57,16 +58,16 @@ jobs:
             *) echo "DAG ended with unexpected state: ${STATE}"; exit 1 ;;
           esac
 
-          if [ -n "$GITHUB_RUN_ID" ]; then
-            echo "Checking nightly build status for run: $GITHUB_RUN_ID"
-            gh run watch "$GITHUB_RUN_ID" --exit-status --interval 60 --repo "$GITHUB_REPOSITORY"
+          if [ -n "$CALLER_RUN_ID" ]; then
+            echo "Checking nightly build status for run: $CALLER_RUN_ID"
+            gh run watch "$CALLER_RUN_ID" --exit-status --interval 60 --repo "$GITHUB_REPOSITORY"
           else
             echo "Error: No Github Run ID provided. Cannot verify nightly build status."
             exit 1
           fi
 
   tag_docker_image:
-    name: Promote ${{ matrix.image_name }} Docker Image
+    name: Promote ${{ matrix.test_type }} Docker Image
     needs: handle_result
     if: needs.handle_result.result == 'success'
     runs-on: linux-x86-n2-16-buildkit
@@ -76,13 +77,11 @@ jobs:
       matrix:
         include:
           - test_type: pre_training
-            image_name: maxtext_jax_nightly
+            stable_image_name: maxtext_jax_stable
+            nightly_image_name: maxtext_jax_nightly
           - test_type: post_training
-            image_name: maxtext_post_training_nightly
-          - test_type: pre_training
-            image_name: maxtext_jax_stable
-          - test_type: post_training
-            image_name: maxtext_post_training_stable
+            stable_image_name: maxtext_post_training_stable
+            nightly_image_name: maxtext_post_training_nightly
     steps:
       - name: Configure Docker
         if: ${{ github.event.client_payload.test_type == '' || github.event.client_payload.test_type == matrix.test_type }}
@@ -91,15 +90,33 @@ jobs:
         if: ${{ github.event.client_payload.test_type == '' || github.event.client_payload.test_type == matrix.test_type }}
         shell: bash
         env:
-          GITHUB_RUN_ID: ${{ github.event.client_payload.github_run_id }}
+          CALLER_RUN_ID: ${{ github.event.client_payload.github_run_id }}
           PROJECT_NAME: ${{ vars.PROJECT_NAME }}
-          IMAGE_NAME: ${{ matrix.image_name }}
+          STABLE_IMAGE_NAME: ${{ matrix.stable_image_name }}
+          NIGHTLY_IMAGE_NAME: ${{ matrix.nightly_image_name }}
         run: |
-          SOURCE_IMAGE="gcr.io/${PROJECT_NAME}/${IMAGE_NAME}"
+          tag_image() {
+            local img="$1"
+            local tag="$2"
+            echo "Tagging ${img}:${CALLER_RUN_ID} as ${tag}..."
+            gcloud container images add-tag "${img}:${CALLER_RUN_ID}" "${img}:${tag}" --quiet
+          }
 
-          # Add the traceability tag to confirm it passed validation suite
-          gcloud container images add-tag "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" \
-            "${SOURCE_IMAGE}:verified-${GITHUB_RUN_ID}" --quiet
+          # 1. Check if stable image exists and tag it
+          STABLE_IMAGE="gcr.io/${PROJECT_NAME}/${STABLE_IMAGE_NAME}"
+          if gcloud container images describe "${STABLE_IMAGE}:${CALLER_RUN_ID}" >/dev/null 2>&1; then
+            echo "Stable image ${STABLE_IMAGE}:${CALLER_RUN_ID} exists. Tagging stable image..."
+            tag_image "${STABLE_IMAGE}" "verified-e2e-${CALLER_RUN_ID}"
+            tag_image "${STABLE_IMAGE}" "latest"
+          else
+            echo "Stable image ${STABLE_IMAGE}:${CALLER_RUN_ID} does not exist. Skipping stable image tagging."
+          fi
 
-          # Add "latest" tag
-          gcloud container images add-tag "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" "${SOURCE_IMAGE}:latest" --quiet
+          # 2. Check if nightly image exists and tag it
+          NIGHTLY_IMAGE="gcr.io/${PROJECT_NAME}/${NIGHTLY_IMAGE_NAME}"
+          if gcloud container images describe "${NIGHTLY_IMAGE}:${CALLER_RUN_ID}" >/dev/null 2>&1; then
+            echo "Nightly image ${NIGHTLY_IMAGE}:${CALLER_RUN_ID} exists. Tagging nightly image..."
+            tag_image "${NIGHTLY_IMAGE}" "latest"
+          else
+            echo "Nightly image ${NIGHTLY_IMAGE}:${CALLER_RUN_ID} does not exist. Skipping nightly image tagging."
+          fi

.github/workflows/pypi_release.yml

@@ -16,10 +16,23 @@
 
 name: Publish MaxText to PyPI
 
-# Triggered by Airflow when E2E DAG completes or can be manually triggered via workflow_dispatch.
+# Triggered by release_pipeline.yml via workflow_call or manually via workflow_dispatch.
 on:
-  repository_dispatch:
-    types: [airflow-dag-complete]
+  workflow_call:
+    inputs:
+      run_id:
+        description: 'GitHub workflow run_id passed to the Airflow DAG run.'
+        required: true
+        type: string
+      maxtext_sha:
+        description: 'MaxText SHA to release to PyPI.'
+        required: true
+        type: string
+      publish:
+        description: 'Set to true to publish to PyPI.'
+        required: false
+        default: true
+        type: boolean
   workflow_dispatch:
     inputs:
       run_id:
@@ -40,77 +53,33 @@ permissions:
   contents: read
 
 jobs:
-  handle_result:
-    name: Handle Airflow DAG Result
-    runs-on: ubuntu-latest
-    # Only run this job if the event that triggered the workflow is 'repository_dispatch'
-    if: github.event_name == 'repository_dispatch'
-    steps:
-      - name: Report DAG result
-        env:
-          STATE: ${{ github.event.client_payload.state }}
-          DAG_ID: ${{ github.event.client_payload.dag_id }}
-          DAG_RUN_ID: ${{ github.event.client_payload.dag_run_id }}
-          SHA: ${{ github.event.client_payload.sha }}
-        run: |
-
-          echo "================================"
-          echo "DAG ID:      ${DAG_ID}"
-          echo "DAG Run ID:  ${DAG_RUN_ID}"
-          echo "Commit SHA:  ${SHA}"
-          echo "State:       ${STATE}"
-          echo "================================"
-
-          # Evaluate state
-          case "${STATE}" in
-            success)
-              echo "DAG '${DAG_ID}' completed successfully."
-              exit 0
-              ;;
-            failed|upstream_failed)
-              echo "DAG '${DAG_ID}' failed with state: ${STATE}."
-              exit 1
-              ;;
-            *)
-              echo "DAG '${DAG_ID}' ended with unexpected state: ${STATE}."
-              exit 1
-              ;;
-          esac
-
   release_approval:
     name: Approve Release
     runs-on: ubuntu-latest
-    needs: handle_result
-    if: |
-      always() &&
-      (needs.handle_result.result == 'success' || needs.handle_result.result == 'skipped') &&
-      (github.event_name == 'repository_dispatch' || inputs.publish)
+    if: inputs.publish == true
     # "release" environment is configured in MaxText repository settings.
     # This environment requires manual approval before proceeding to the next job.
     environment: release
     steps:
-    - name: Acknowledge Approval
-      run: echo "Release approved, proceeding to publishing MaxText package."
+      - name: Acknowledge Approval
+        run: echo "Release approved, proceeding to publishing MaxText package."
 
   build_maxtext_package:
     needs: release_approval
-    if: |
-      always() &&
-      (needs.release_approval.result == 'success' || needs.release_approval.result == 'skipped')
     uses: ./.github/workflows/build_package.yml
     with:
       device_type: tpu
       device_name: v4-8
       cloud_runner: linux-x86-n2-16-buildkit
       # Build package from a particular commit to ensure we are publishing the same code that was tested
-      maxtext_sha: ${{ github.event.client_payload.sha || inputs.maxtext_sha }}
+      maxtext_sha: ${{ inputs.maxtext_sha }}
 
   publish_maxtext_to_pypi:
     name: Publish MaxText to PyPI
     needs: [release_approval, build_maxtext_package]
     runs-on: ubuntu-latest
     environment: release
-    if: needs.build_maxtext_package.result == 'success' && (github.event_name == 'repository_dispatch' || inputs.publish)
+    if: needs.build_maxtext_package.result == 'success' && inputs.publish
     permissions:
       id-token: write # required for PyPI Trusted Publishing (OIDC)
       contents: read
@@ -164,11 +133,10 @@ jobs:
         env:
           PROJECT_NAME: ${{ vars.PROJECT_NAME }}
           IMAGE_NAME: ${{ matrix.image_name }}
-          RUN_ID_INPUT: ${{ github.event_name == 'workflow_dispatch' && inputs.run_id || github.event.client_payload.github_run_id }}
+          GITHUB_RUN_ID: ${{ inputs.run_id }}
           PYPI_VERSION: ${{ needs.get_latest_maxtext_pypi_version.outputs.latest_pypi_version }}
         run: |
           SOURCE_IMAGE="gcr.io/${PROJECT_NAME}/${IMAGE_NAME}"
-          GITHUB_RUN_ID="${RUN_ID_INPUT}"
           gcloud container images add-tag \
             "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" \
             "${SOURCE_IMAGE}:${PYPI_VERSION}" \