UploadDockerImages.yml

# Copyright 2023 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 install Python dependencies, run tests and lint with a variety of Python versions
# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python

name: Build Images

on:
  schedule:
    # Run the job daily at 7PM PST (3AM UTC)
    - cron:  '0 3 * * *'

jobs:
  tpu:
    strategy:
      fail-fast: false
      matrix:
        device-type: ["v4-8"]
    runs-on: ["self-hosted", "tpu", "${{ matrix.device-type }}"]
    steps:
    - uses: actions/checkout@v3
    - name: build jax stable image
      run : |
        project=tpu-prod-env-multipod
        local_image_name=maxtext_local_jax_stable
        cloud_image_name=maxtext_jax_stable
        bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$local_image_name MODE=stable
        docker build --build-arg BASEIMAGE=${local_image_name} -f ./maxtext_runner.Dockerfile -t ${local_image_name}_runner .
        gcloud auth configure-docker --quiet
        docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:latest
        docker push gcr.io/$project/${cloud_image_name}:latest
        image_date=$(date +%Y-%m-%d)
        docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:${image_date}
        docker push gcr.io/$project/${cloud_image_name}:${image_date}
    - name: build jax nightly image
      run : |
        project=tpu-prod-env-multipod
        local_image_name=maxtext_local_jax_nightly
        cloud_image_name=maxtext_jax_nightly
        bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$local_image_name MODE=nightly
        docker build --build-arg BASEIMAGE=${local_image_name} -f ./maxtext_runner.Dockerfile -t ${local_image_name}_runner .
        gcloud auth configure-docker --quiet
        docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:latest
        docker push gcr.io/$project/${cloud_image_name}:latest
        image_date=$(date +%Y-%m-%d)
        docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:${image_date}
        docker push gcr.io/$project/${cloud_image_name}:${image_date}

  gpu:
    strategy:
      fail-fast: false
      matrix:
        device-type: ["a100-40gb-4"]
    runs-on: ["self-hosted", "gpu", "${{ matrix.device-type }}"]
    steps:
    - uses: actions/checkout@v3
    - name: build jax stable image
      run : |
        project=tpu-prod-env-multipod
        local_image_name=maxtext_gpu_local_jax_stable
        cloud_image_name=maxtext_gpu_jax_stable
        bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$local_image_name MODE=stable DEVICE=gpu
        docker build --build-arg BASEIMAGE=${local_image_name} -f ./maxtext_runner.Dockerfile -t ${local_image_name}_runner .
        gcloud auth configure-docker --quiet
        docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:latest
        docker push gcr.io/$project/${cloud_image_name}:latest
        image_date=$(date +%Y-%m-%d)
        docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:${image_date}
        docker push gcr.io/$project/${cloud_image_name}:${image_date}
    - name: build jax nightly image
      run : |
        project=tpu-prod-env-multipod
        local_image_name=maxtext_gpu_local_jax_nightly
        cloud_image_name=maxtext_gpu_jax_nightly
        bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$local_image_name MODE=nightly DEVICE=gpu
        docker build --build-arg BASEIMAGE=${local_image_name} -f ./maxtext_runner.Dockerfile -t ${local_image_name}_runner .
        gcloud auth configure-docker --quiet
        docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:latest
        docker push gcr.io/$project/${cloud_image_name}:latest
        image_date=$(date +%Y-%m-%d)
        docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:${image_date}
        docker push gcr.io/$project/${cloud_image_name}:${image_date}