diff --git a/.github/workflows/build_and_push_docker_image.yml b/.github/workflows/build_and_push_docker_image.yml index 612c68ad07..f88a18e4da 100644 --- a/.github/workflows/build_and_push_docker_image.yml +++ b/.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 }} diff --git a/.github/workflows/build_package.yml b/.github/workflows/build_package.yml index 8eb0c74c8c..6cf4fb1b40 100644 --- a/.github/workflows/build_package.yml +++ b/.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 diff --git a/.github/workflows/gpu_docker_images_pipeline.yml b/.github/workflows/gpu_docker_images_pipeline.yml index 1dc505ce7a..d1272b0f9d 100644 --- a/.github/workflows/gpu_docker_images_pipeline.yml +++ b/.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 }} diff --git a/.github/workflows/rc_release.yml b/.github/workflows/rc_release.yml index c0ca85c8e9..c6c849d3cd 100644 --- a/.github/workflows/rc_release.yml +++ b/.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: diff --git a/.github/workflows/release_pipeline.yml b/.github/workflows/release_pipeline.yml index d5a2723a0a..4e04cec065 100644 --- a/.github/workflows/release_pipeline.yml +++ b/.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 }} diff --git a/.github/workflows/tpu_docker_images_pipeline.yml b/.github/workflows/tpu_docker_images_pipeline.yml index 37cd0adf47..47936e477a 100644 --- a/.github/workflows/tpu_docker_images_pipeline.yml +++ b/.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 }}