Merge branch 'main' of github.com:AI-Hypercomputer/maxtext into shuningjin-ckpt-structure

Author: shuningjin

.github/CODEOWNERS

@@ -3,7 +3,7 @@
 # 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 @suexu1025 @jesselu-google
-src/MaxText/utils/ckpt_conversion @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @hengtaoguo @gagika @shralex @richjames0 @NicoGrande
+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 @suexu1025 @jesselu-google
 
 # Features

.github/workflows/UploadDockerImages.yml

@@ -15,7 +15,7 @@
 # This workflow builds and pushes MaxText images for both TPU and GPU devices.
 # It runs automatically daily at 12am UTC, on Pull Requests, or manually via Workflow Dispatch.
 
-name: Build and Test Images
+name: Build Images
 
 on:
   schedule:
@@ -32,11 +32,6 @@ on:
           - all
           - tpu
           - gpu
-      for_dev_test:
-        description: 'For development test purpose. All images will be added a -test suffix'
-        required: false
-        type: boolean
-        default: false
 
 permissions:
   contents: read
@@ -47,7 +42,6 @@ jobs:
     outputs:
       maxtext_sha: ${{ steps.vars.outputs.maxtext_sha }}
       image_date: ${{ steps.vars.outputs.image_date }}
-      image_suffix: ${{ steps.vars.outputs.image_suffix }}
     steps:
      - name: Checkout MaxText
        uses: actions/checkout@v5
@@ -61,13 +55,6 @@ jobs:
         # Image date
         echo "image_date=$(date +%Y-%m-%d)" >> $GITHUB_OUTPUT
 
-        # If for_dev_test is true, set suffix to -test, otherwise empty
-        if [[ "${{ github.event.inputs.for_dev_test }}" == "true" ]]; then
-          echo "image_suffix=-test" >> $GITHUB_OUTPUT
-        else
-          echo "image_suffix=" >> $GITHUB_OUTPUT
-        fi
-
   tpu-pre-training:
     name: ${{ matrix.image_name }}
     needs: setup
@@ -85,7 +72,7 @@ jobs:
             dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
-      image_name: ${{ matrix.image_name }}${{ needs.setup.outputs.image_suffix  }}
+      image_name: ${{ matrix.image_name }}
       device: ${{ matrix.device }}
       build_mode: ${{ matrix.build_mode }}
       dockerfile: ${{ matrix.dockerfile }}
@@ -109,13 +96,14 @@ jobs:
             dockerfile: ./dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
-      image_name: ${{ matrix.image_name }}${{ needs.setup.outputs.image_suffix  }}
+      image_name: ${{ matrix.image_name }}
       device: ${{ matrix.device }}
       build_mode: ${{ matrix.build_mode }}
       dockerfile: ${{ matrix.dockerfile }}
       maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
       image_date: ${{ needs.setup.outputs.image_date }}
       base_image: gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:${{ needs.setup.outputs.image_date }}
+      is_post_training: true
 
   gpu-pre-training:
     name: ${{ matrix.image_name }}
@@ -134,48 +122,9 @@ jobs:
             dockerfile: ./dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
-      image_name: ${{ matrix.image_name }}${{ needs.setup.outputs.image_suffix  }}
+      image_name: ${{ matrix.image_name }}
       device: ${{ matrix.device }}
       build_mode: ${{ matrix.build_mode }}
       dockerfile: ${{ matrix.dockerfile }}
       maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
       image_date: ${{ needs.setup.outputs.image_date }}
-
-  # TEST JOBS
-  pre-training-tpu-tests:
-    needs: [setup, tpu-pre-training]
-    strategy:
-      fail-fast: false
-      matrix:
-        image: [maxtext_jax_stable, maxtext_jax_nightly]
-    uses: ./.github/workflows/test_and_tag_docker_image.yml
-    with:
-      image_name: ${{ matrix.image }}${{ needs.setup.outputs.image_suffix  }}
-      image_date: ${{ needs.setup.outputs.image_date }}
-      test_mode: tpu-pre-training
-
-  post-training-tpu-tests:
-    needs: [setup, tpu-post-training]
-    strategy:
-      fail-fast: false
-      matrix:
-        image: [maxtext_post_training_stable, maxtext_post_training_nightly]
-    uses: ./.github/workflows/test_and_tag_docker_image.yml
-    with:
-      image_name: ${{ matrix.image }}${{ needs.setup.outputs.image_suffix  }}
-      image_date: ${{ needs.setup.outputs.image_date }}
-      test_mode: tpu-post-training
-
-
-  pre-training-gpu-tests:
-    needs: [setup, gpu-pre-training]
-    strategy:
-      fail-fast: false
-      matrix:
-        image: [maxtext_gpu_jax_stable, maxtext_gpu_jax_nightly]
-    uses: ./.github/workflows/test_and_tag_docker_image.yml
-    with:
-      image_name: ${{ matrix.image }}${{ needs.setup.outputs.image_suffix  }}
-      image_date: ${{ needs.setup.outputs.image_date }}
-      test_mode: gpu-pre-training
-

.github/workflows/build_and_push_docker_image.yml

@@ -41,6 +41,10 @@ on:
         required: false
         type: string
         default: ''
+      is_post_training:
+        required: false
+        type: boolean
+        default: false
 
 permissions:
   contents: read
@@ -62,13 +66,18 @@ jobs:
         id: check
         shell: bash
         run: |
-          if [[ "${{ github.event_name }}" == "workflow_dispatch" && "${{ github.event.inputs.target_device }}" != "all" && "${{ github.event.inputs.target_device }}" != "${{ inputs.device }}" ]]; then
+          if [[ "${{ github.event_name }}" == "workflow_dispatch" && "${GITHUB_EVENT_INPUTS_TARGET_DEVICE}" != "all" && "${GITHUB_EVENT_INPUTS_TARGET_DEVICE}" != "${INPUTS_DEVICE}" ]]; then
             echo "should_run=false" >> $GITHUB_OUTPUT
-            echo "Skipping ${{ inputs.image_name }} build for device: ${{ inputs.device }} in ${{ inputs.build_mode }} mode."
+            echo "Skipping ${INPUTS_IMAGE_NAME} build for device: ${INPUTS_DEVICE} in ${INPUTS_BUILD_MODE} mode."
           else
             echo "should_run=true" >> $GITHUB_OUTPUT
-            echo "Building ${{ inputs.image_name }} for device: ${{ inputs.device }} in ${{ inputs.build_mode }} mode."
+            echo "Building ${INPUTS_IMAGE_NAME} for device: ${INPUTS_DEVICE} in ${INPUTS_BUILD_MODE} mode."
           fi
+        env:
+          GITHUB_EVENT_INPUTS_TARGET_DEVICE: ${{ github.event.inputs.target_device }}
+          INPUTS_DEVICE: ${{ inputs.device }}
+          INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
+          INPUTS_BUILD_MODE: ${{ inputs.build_mode }}
 
       - name: Checkout MaxText
         uses: actions/checkout@v5
@@ -78,9 +87,7 @@ jobs:
           ref: ${{ inputs.maxtext_sha }}
 
       - name: Checkout post-training dependencies
-        if: |
-          steps.check.outputs.should_run == 'true' &&
-          contains(inputs.image_name, 'post_training_nightly')
+        if: steps.check.outputs.should_run == 'true' && inputs.image_name == 'maxtext_post_training_nightly'
         run: |
           git clone https://github.com/google/tunix.git ./tunix
           git clone https://github.com/vllm-project/vllm.git ./vllm
@@ -108,7 +115,8 @@ jobs:
           push: true
           context: .
           file: ${{ inputs.dockerfile }}
-          tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:${{ inputs.image_date }}-build-${{ github.run_id }}
+          tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:latest
+          cache-from: type=gha
           outputs: type=image,compression=zstd,force-compression=true
           build-args: |
             DEVICE=${{ inputs.device }}
@@ -122,21 +130,28 @@ jobs:
         if: steps.check.outputs.should_run == 'true'
         shell: bash
         run: |
-          SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}"
-          TEMP_IMG="$SOURCE_IMAGE:${{ inputs.image_date }}-build-${{ github.run_id }}"
+          SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${INPUTS_IMAGE_NAME}"
+
+          # Add date tag
+          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
 
           # Convert date to YYYYMMDD format
-          clean_date=$(echo "${{ inputs.image_date }}" | sed 's/[-:]//g' | cut -c1-8)
+          clean_date=$(echo "${INPUTS_IMAGE_DATE}" | sed 's/[-:]//g' | cut -c1-8)
 
           # Add MaxText tag
           maxtext_hash=$(git rev-parse --short HEAD)
-          gcloud container images add-tag "$TEMP_IMG" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
-          gcloud container images add-tag "$TEMP_IMG" "$SOURCE_IMAGE:${{ inputs.image_date }}" --quiet
+          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
+
 
           # Add post-training dependencies tags
-          for dir in tunix vllm tpu-inference; do
-            if [ -d "./$dir" ]; then
-              dir_hash=$(git -C "$dir" rev-parse --short HEAD)
-              gcloud container images add-tag "$TEMP_IMG" "$SOURCE_IMAGE:${dir}_${dir_hash}_${clean_date}" --quiet
-            fi
-          done
+          if [ "${{ inputs.is_post_training }}" == "true" ]; then
+            for dir in tunix vllm tpu-inference; do
+              if [ -d "./$dir" ]; then
+                dir_hash=$(git -C "$dir" rev-parse --short HEAD)
+                gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${dir}_${dir_hash}_${clean_date}" --quiet
+              fi
+            done
+          fi
+        env:
+          INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
+          INPUTS_IMAGE_DATE: ${{ inputs.image_date }}