build_and_push_docker_image.yml

# 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 build and push MaxText Docker image to GCR.

name: Build and Push MaxText Docker Images

on:
  workflow_call:
    inputs:
      image_name:
        required: true
        type: string
      device:
        required: true
        type: string
      build_mode:
        required: true
        type: string
      dockerfile:
        required: true
        type: string
      maxtext_sha:
        required: true
        type: string
      image_date:
        required: true
        type: string
      base_image:
        required: false
        type: string
        default: ''
      is_post_training:
        required: false
        type: boolean
        default: false

permissions:
  contents: read

jobs:
  build_and_push:
    runs-on: linux-x86-n2-16-buildkit
    container: google/cloud-sdk:524.0.0
    if: >
      github.event_name == 'schedule' ||
      github.event_name == 'pull_request' ||
      github.event_name == 'workflow_dispatch' && (
        github.event.inputs.target_device == 'all' ||
        github.event.inputs.target_device == 'tpu' ||
        github.event.inputs.target_device == 'gpu'
      )
    steps:
      - name: Check if build should run
        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
            echo "should_run=false" >> $GITHUB_OUTPUT
            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."
          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
        if: steps.check.outputs.should_run == 'true'
        with:
          # This ensures that every job clones the exact same commit as "setup" job
          ref: ${{ inputs.maxtext_sha }}

      - name: Checkout post-training dependencies
        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
          git clone https://github.com/vllm-project/tpu-inference.git ./tpu-inference

      - name: Mark git repositories as safe
        run: git config --global --add safe.directory '*'
        if: steps.check.outputs.should_run == 'true'

      - name: Configure Docker
        run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q
        if: steps.check.outputs.should_run == 'true'

      - name: Set up Docker BuildX
        uses: docker/setup-buildx-action@v3.11.1
        if: steps.check.outputs.should_run == 'true'
        with:
          driver: remote
          endpoint: tcp://localhost:1234

      - name: Build and push Docker image
        uses: docker/build-push-action@v6
        if: steps.check.outputs.should_run == 'true'
        with:
          push: true
          context: .
          file: ${{ inputs.dockerfile }}
          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 }}
            MODE=${{ inputs.build_mode }}
            JAX_VERSION=NONE
            LIBTPU_VERSION=NONE
            INCLUDE_TEST_ASSETS=true
            ${{ inputs.base_image != '' && format('BASEIMAGE={0}', inputs.base_image) || '' }}

      - name: Add tags to Docker image
        if: steps.check.outputs.should_run == 'true'
        shell: bash
        run: |
          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)

          # Add MaxText tag
          maxtext_hash=$(git rev-parse --short HEAD)
          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet


          # Add post-training dependencies tags
          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 }}