Trigger unit tests for docker images upload workflow

xibinliu · xibinliu · commit f8d52410d386 · 2026-02-09T17:31:10.000Z
- images will only be tagged to "latest" when units pass
diff --git a/.github/workflows/UploadDockerImages.yml b/.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 Images
+name: Build and Test Images
 
 on:
   schedule:
@@ -32,6 +32,11 @@ 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
@@ -42,6 +47,7 @@ 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
@@ -55,6 +61,13 @@ 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
@@ -72,7 +85,7 @@ jobs:
             dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
-      image_name: ${{ matrix.image_name }}
+      image_name: ${{ matrix.image_name }}${{ needs.setup.outputs.image_suffix  }}
       device: ${{ matrix.device }}
       build_mode: ${{ matrix.build_mode }}
       dockerfile: ${{ matrix.dockerfile }}
@@ -96,14 +109,13 @@ 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 }}
+      image_name: ${{ matrix.image_name }}${{ needs.setup.outputs.image_suffix  }}
       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 }}
@@ -122,9 +134,48 @@ jobs:
             dockerfile: ./dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
-      image_name: ${{ matrix.image_name }}
+      image_name: ${{ matrix.image_name }}${{ needs.setup.outputs.image_suffix  }}
       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
+
diff --git a/.github/workflows/test_and_tag_docker_image.yml b/.github/workflows/test_and_tag_docker_image.yml
@@ -0,0 +1,93 @@
+# 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 test and tag MaxText Docker image to GCR.
+name: Test and Tag MaxText Docker Images
+
+on:
+  workflow_call:
+    inputs:
+      image_name:
+        required: true
+        type: string
+      test_mode:
+        description: "Test mode (tpu-pre-training, tpu-post-training, gpu-pre-training)"
+        required: true
+        type: string
+      image_date:
+        required: true
+        type: string
+
+permissions:
+  contents: read
+
+jobs:
+  test:
+    strategy:
+      fail-fast: false
+      matrix:
+        flavor: >-
+            ${{ fromJSON('{
+              "gpu-pre-training": ["gpu-unit", "gpu-integration"],
+              "tpu-post-training": ["post-training-tpu-unit", "post-training-tpu-integration"],
+              "tpu-pre-training": ["tpu-unit", "tpu-integration", "cpu-unit"]
+            }')[inputs.test_mode] }}
+    uses: ./.github/workflows/run_tests_coordinator.yml
+    with:
+      flavor: ${{ matrix.flavor }}
+      base_image: ${{ inputs.image_name }}:${{ inputs.image_date }}-build-${{ github.run_id }}
+      is_scheduled_run: true
+      maxtext_installed: true
+
+  notebook-test:
+    if: inputs.test_mode == 'tpu-post-training'
+    uses: ./.github/workflows/run_jupyter_notebooks.yml
+    with:
+      device_type: tpu
+      device_name: v6e-4
+      base_image: ${{ inputs.image_name }}:${{ inputs.image_date }}-build-${{ github.run_id }}
+      cloud_runner: linux-x86-ct6e-180-4tpu
+      maxtext_installed: true
+    secrets:
+      HF_TOKEN: ${{ secrets.HF_TOKEN }}
+
+  tagging:
+    needs: [test, notebook-test]
+    if: |
+      always() &&
+      needs.test.result == 'success' &&
+      (needs.notebook-test.result == 'success' || needs.notebook-test.result == 'skipped')
+    runs-on: ubuntu-latest
+    container: google/cloud-sdk:524.0.0
+    steps:
+      - name: Configure Docker
+        run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q
+
+      - name: Create Production Tags
+        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 }}"
+
+          # Validate existence first
+          gcloud container images describe "$TEMP_IMG" > /dev/null
+
+          # 1. Date Tag
+          gcloud container images add-tag "$TEMP_IMG" "$SOURCE_IMAGE:${{ inputs.image_date }}" --quiet
+
+          # 2. Latest Tag
+          gcloud container images add-tag "$TEMP_IMG" "$SOURCE_IMAGE:latest" --quiet
+
+          # 3. Clean up Temporary Tag
+          gcloud container images untag "$TEMP_IMG" --quiet
