[chiajunglien]

Description

Resolves:

• b/573710066: Update run_e2e_tests.yml, promote_docker_image.yml, and maxtext_e2e_tests Airflow DAG to accept configurable Docker images
• b/574141449: [Milestone 1] Update promote_docker_image.yml to tag RC images

Companion Airflow PR: GoogleCloudPlatform/ml-auto-solutions#1414. Since neither side keeps a fallback, the two PRs should land together.

Changes

1. .github/workflows/rc_release.yml (b/573710066):

   • Pass mode: release_candidate to run_e2e_tests.yml, resolving TODO(b/573710066).

2. .github/workflows/run_e2e_tests.yml (b/573710066):

   • Resolve the Docker images from mode: tpu_pre_training:<short_sha> / tpu_post_training:<short_sha> for
     release_candidate, and maxtext_jax_<mode>:<run_id> / maxtext_post_training_<mode>:<run_id> otherwise.
   • Pass them to the Airflow DAG as a single docker_images dict with tpu_pre_training / tpu_post_training keys.

3. .github/workflows/promote_docker_image.yml (b/573710066, b/574141449):

   • Promote the image reported by Airflow (client_payload.docker_image), or the required docker_image input (full image URI) for manual runs. Removed the job matrix, hardcoded image names, Checkout MaxText, and promote_from_staging().
   • Split the full image URI into the image and IMAGE_TAG, so Airflow and manual runs go through the same path.
   • On E2E pass: attach :verified-e2e-${IMAGE_TAG} and remove any stale :failed-e2e-${IMAGE_TAG}. Daily stable images are also promoted to :latest, together with the matching nightly image, if :verified-ci-${IMAGE_TAG} exists (replacing gh run watch). RC images skip :latest since they already get :latest-rc at build time.
   • On E2E failure, attach :failed-e2e-${IMAGE_TAG}. On an unexpected state, attach :unknown-e2e-${IMAGE_TAG}. The workflow no longer fails in either case.

Tests

• Manually triggered this workflow on this branch with both RC images (tpu_pre_training:1ed249c and tpu_post_training:1ed249c): both got :verified-e2e-1ed249c and no :latest, including tpu_post_training, which has :verified-ci-1ed249c. The test tags were removed afterwards.
  • https://github.com/AI-Hypercomputer/maxtext/actions/runs/39694102317
  • https://github.com/AI-Hypercomputer/maxtext/actions/runs/39694650729
• Validated the tagging logic locally with a mocked gcloud: pass / fail / unknown states, :latest for stable images only (with nightly promotion), missing nightly image, and missing image.
• Checked the end-to-end payloads against feat(multipod): Support configurable Docker images and scope-based conversion image in maxtext_e2e_tests  GoogleCloudPlatform/ml-auto-solutions#1414: docker_images sent to Airflow, and the callback state / docker_image read back.
• codespell, yamllint, and actionlint passed.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Note

Gemini is unable to generate a review for this pull request due to the file types involved not being currently supported.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[chiajunglien]

Also addressed the points we discussed offline:

• Nightly images: when a stable image gets :latest, the matching nightly image is promoted to :latest too.
• docker_image is now a required input for manual runs, together with image_tag.

Other changes that follow from the above:

• Removed the test_type input, since docker_image is now required and test_type was only used by the fallback.
• Updated run-name and the manual trigger step to use docker_image / image_tag.
• Updated the header comment and dropped actions: read from permissions, since gh run watch is gone.