Tag RC Docker images during build and run CI, E2E, and performance tests in parallel

.github/workflows/build_and_push_docker_image.yml

@@ -37,6 +37,9 @@ on:
       maxtext_sha:
         required: true
         type: string
+      image_date:
+        required: true
+        type: string
       workflow:
         required: false
         type: string
@@ -45,6 +48,11 @@ on:
         required: false
         type: boolean
         default: false
+      is_rc:
+        description: 'Whether the image is a Release Candidate'
+        required: false
+        type: boolean
+        default: false
       fast_rebuild:
         description: 'Use code-overlay rebuild instead of full build'
         required: false
@@ -238,7 +246,7 @@ jobs:
           fi
           echo "BUILD_REPO=${BUILD_REPO}" >> "$GITHUB_ENV"
           echo "build_repo=${BUILD_REPO}" >> "$GITHUB_OUTPUT"
-          echo "maxtext_short_sha=$(git rev-parse --short HEAD)" >> "$GITHUB_OUTPUT"
+          echo "maxtext_short_sha=$(git rev-parse --short=7 HEAD)" >> "$GITHUB_OUTPUT"
 
       - name: Resolve source image digest
         if: ${{ inputs.fast_rebuild }}
@@ -407,7 +415,6 @@ jobs:
           gcloud container images add-tag "${SOURCE_IMAGE}@${IMAGE_DIGEST}" "${TEMP_IMG}" --quiet
 
           # Add date tag
-          IMAGE_DATE="$(date +%Y-%m-%d)"
           gcloud container images add-tag "${SOURCE_IMAGE}@${IMAGE_DIGEST}" "$SOURCE_IMAGE:$IMAGE_DATE" --quiet
 
           # Convert date to YYYYMMDD format
@@ -416,14 +423,30 @@ jobs:
           # Add MaxText tag
           gcloud container images add-tag "${SOURCE_IMAGE}@${IMAGE_DIGEST}" "${SOURCE_IMAGE}:maxtext_${MAXTEXT_SHA}_${clean_date}" --quiet
 
+          # Add Release Candidate tags
+          if [[ "${IS_RC}" == "true" ]]; then
+            gcloud container images add-tag "${SOURCE_IMAGE}@${IMAGE_DIGEST}" "${SOURCE_IMAGE}:${MAXTEXT_SHA}" --quiet
+            if [[ "${INPUTS_REPOSITORY_NAME}" == "maxtext-images" ]]; then
+              gcloud container images add-tag "${SOURCE_IMAGE}@${IMAGE_DIGEST}" "${SOURCE_IMAGE}:latest-rc" --quiet
+            else
+              echo "Skipping :latest-rc tag (functional and performance tests disabled for this run)."
+            fi
+          fi
+
           # Also keep staging tags in sync when promoting
           if [[ "${BUILD_REPO}" != "${INPUTS_REPOSITORY_NAME}" ]]; then
             gcloud container images add-tag "${STAGING_IMAGE}@${IMAGE_DIGEST}" "${STAGING_IMAGE}:${IMAGE_DATE}" --quiet
             gcloud container images add-tag "${STAGING_IMAGE}@${IMAGE_DIGEST}" "${STAGING_IMAGE}:maxtext_${MAXTEXT_SHA}_${clean_date}" --quiet
+            if [[ "${IS_RC}" == "true" ]]; then
+              gcloud container images add-tag "${STAGING_IMAGE}@${IMAGE_DIGEST}" "${STAGING_IMAGE}:${MAXTEXT_SHA}" --quiet
+              gcloud container images add-tag "${STAGING_IMAGE}@${IMAGE_DIGEST}" "${STAGING_IMAGE}:latest-rc" --quiet
+            fi
           fi
         env:
           BUILD_REPO: ${{ needs.build_and_push.outputs.build_repo }}
           IMAGE_DIGEST: ${{ needs.build_and_push.outputs.image_digest }}
+          IMAGE_DATE: ${{ inputs.image_date }}
+          IS_RC: ${{ inputs.is_rc }}
           MAXTEXT_SHA: ${{ needs.build_and_push.outputs.maxtext_short_sha }}
           GITHUB_RUN_ID: ${{ github.run_id }}
           INPUTS_REPOSITORY_NAME: ${{ inputs.repository_name }}

.github/workflows/build_package.yml

@@ -36,6 +36,9 @@ on:
       maxtext_sha:
         description: "MaxText short SHA used for the build"
         value: ${{ jobs.build_and_upload.outputs.maxtext_sha }}
+      image_date:
+        description: "Date (YYYY-MM-DD) used for tagging Docker images"
+        value: ${{ jobs.build_and_upload.outputs.image_date }}
 
 permissions:
   contents: read
@@ -47,6 +50,7 @@ jobs:
     container: python:3.12-slim-trixie
     outputs:
       maxtext_sha: ${{ steps.vars.outputs.maxtext_sha }}
+      image_date: ${{ steps.vars.outputs.image_date }}
     steps:
       - name: Checkout MaxText
         uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5.1.0
@@ -62,6 +66,7 @@ jobs:
           # MaxText SHA used to build the package
           MAXTEXT_SHA="$MAXTEXT_SHA_INPUT"
           echo "maxtext_sha=${MAXTEXT_SHA}" >> "$GITHUB_OUTPUT"
+          echo "image_date=$(date -u +%Y-%m-%d)" >> "$GITHUB_OUTPUT"
       - name: Install build tools
         run: |
           apt-get update -y && apt-get install -y --no-install-recommends curl ca-certificates

.github/workflows/gpu_docker_images_pipeline.yml

@@ -113,6 +113,7 @@ jobs:
       workflow: ${{ matrix.workflow }}
       dockerfile: maxtext_gpu_dependencies.Dockerfile
       maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }}
+      image_date: ${{ needs.build_maxtext_package.outputs.image_date }}
       include_test_assets: true
       fast_rebuild: ${{ inputs.mode == 'fast-rebuild' }}
       source_run_id: ${{ needs.resolve_source.outputs.source_run_id || inputs.source_run_id }}

.github/workflows/rc_release.yml

@@ -47,7 +47,6 @@ jobs:
     outputs:
       maxtext_sha: ${{ steps.vars.outputs.maxtext_sha }}
       short_sha: ${{ steps.vars.outputs.short_sha }}
-      rc_date: ${{ steps.vars.outputs.rc_date }}
     steps:
       - name: Checkout MaxText
         uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5.1.0
@@ -65,7 +64,6 @@ jobs:
           {
             echo "maxtext_sha=${MAXTEXT_SHA}"
             echo "short_sha=${SHORT_SHA}"
-            echo "rc_date=$(date -u +%Y%m%d)"
           } >> "$GITHUB_OUTPUT"
 
   build_maxtext_package:
@@ -111,13 +109,15 @@ jobs:
       workflow: ${{ matrix.workflow }}
       dockerfile: ${{ matrix.dockerfile }}
       maxtext_sha: ${{ needs.get_maxtext_sha.outputs.maxtext_sha }}
+      image_date: ${{ needs.build_maxtext_package.outputs.image_date }}
       include_test_assets: true
+      is_rc: true
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
   run_ci_tests:
     name: Run ${{ matrix.name }} CI Tests
-    needs: build_release_candidate_images
+    needs: [build_release_candidate_images, get_maxtext_sha]
     if: github.event_name == 'schedule' || inputs.repository_name == 'maxtext-images'
     strategy:
       fail-fast: false
@@ -135,69 +135,16 @@ jobs:
     with:
       image_name: ${{ matrix.image_name }}
       repository_name: 'maxtext-images'
-      image_tag: ${{ github.run_id }}
+      image_tag: ${{ needs.get_maxtext_sha.outputs.short_sha }}
       device: ${{ matrix.device }}
       workflow: ${{ matrix.workflow }}
 
-  tag_release_candidate_images:
-    name: Tag ${{ matrix.name }} Docker Image
-    needs: [build_release_candidate_images, get_maxtext_sha, run_ci_tests]
-    if: |
-      always() &&
-      needs.get_maxtext_sha.result == 'success' &&
-      needs.build_release_candidate_images.result == 'success' &&
-      (needs.run_ci_tests.result == 'success' || needs.run_ci_tests.result == 'skipped')
-    runs-on: linux-x86-n2-16-buildkit
-    container: google/cloud-sdk:524.0.0
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - name: 'TPU Pre-Training'
-            image_name: tpu_pre_training
-          - name: 'TPU Post-Training'
-            image_name: tpu_post_training
-    steps:
-      - name: Configure Docker
-        run: gcloud auth configure-docker us-docker.pkg.dev -q
-      - name: Add tags to Docker image
-        shell: bash
-        env:
-          PROJECT_NAME: ${{ vars.PROJECT_NAME }}
-          REPOSITORY_NAME: ${{ inputs.repository_name || 'maxtext-images' }}
-          DOCKER_IMAGE_NAME: ${{ matrix.image_name }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          SHA_RC: ${{ needs.get_maxtext_sha.outputs.short_sha }}
-          RC_DATE: ${{ needs.get_maxtext_sha.outputs.rc_date }}
-          TAG_LATEST_RC: ${{ needs.run_ci_tests.result == 'success' && (github.event_name == 'schedule' || inputs.repository_name == 'maxtext-images') }}
-        run: |
-          SOURCE_IMAGE="us-docker.pkg.dev/${PROJECT_NAME}/${REPOSITORY_NAME}/${DOCKER_IMAGE_NAME}"
-
-          gcloud container images add-tag \
-            "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" \
-            "${SOURCE_IMAGE}:${SHA_RC}" \
-            --quiet
-
-          gcloud container images add-tag \
-            "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" \
-            "${SOURCE_IMAGE}:maxtext_${SHA_RC}_${RC_DATE}" \
-            --quiet
-
-          if [ "${TAG_LATEST_RC}" = "true" ]; then
-            gcloud container images add-tag \
-              "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" \
-              "${SOURCE_IMAGE}:latest-rc" \
-              --quiet
-          else
-            echo "Skipping :latest-rc tag (functional and performance tests disabled for this run)."
-          fi
-
   # TODO(b/573710066): Pass custom Docker image names (tpu_pre_training and
   # tpu_post_training) once run_e2e_tests.yml, promote_docker_image.yml, and
   # the maxtext_e2e_tests Airflow DAG support configurable Docker images.
   run_e2e_tests:
     name: Run E2E tests
-    needs: [get_maxtext_sha, tag_release_candidate_images]
+    needs: [get_maxtext_sha, build_release_candidate_images]
     if: github.event_name == 'schedule' || inputs.repository_name == 'maxtext-images'
     uses: ./.github/workflows/run_e2e_tests.yml
     with:

.github/workflows/release_pipeline.yml

@@ -128,6 +128,7 @@ jobs:
       workflow: ${{ matrix.workflow }}
       dockerfile: ${{ matrix.dockerfile }}
       maxtext_sha: ${{ needs.get_maxtext_sha.outputs.maxtext_sha }}
+      image_date: ${{ needs.build_maxtext_package.outputs.image_date }}
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
@@ -156,6 +157,7 @@ jobs:
       workflow: ${{ matrix.workflow }}
       dockerfile: ${{ matrix.dockerfile }}
       maxtext_sha: ${{ needs.get_maxtext_sha.outputs.maxtext_sha }}
+      image_date: ${{ needs.build_maxtext_package.outputs.image_date }}
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 

.github/workflows/tpu_docker_images_pipeline.yml

@@ -131,6 +131,7 @@ jobs:
       workflow: ${{ matrix.workflow }}
       dockerfile: maxtext_tpu_dependencies.Dockerfile
       maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }}
+      image_date: ${{ needs.build_maxtext_package.outputs.image_date }}
       include_test_assets: true
       fast_rebuild: ${{ inputs.mode == 'fast-rebuild' }}
       source_run_id: ${{ needs.resolve_source.outputs.source_run_id || inputs.source_run_id }}