Added Jax Stable Stack build to daily TPU build workflow

.github/workflows/UploadDockerImages.yml

@@ -0,0 +1,35 @@
+# Copyright 2024 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:
+  push:
+    branches: [ "pb-jax-ss" ]
+
+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 stack image
+      run: |
+        bash docker_maxdiffusion_image_upload.sh PROJECT_ID=tpu-prod-env-multipod BASEIMAGE=us-docker.pkg.dev/tpu-prod-env-multipod/jax-ss/tpu:jax0.4.30-rev1 CLOUD_IMAGE_NAME=maxdiffusion-jax-ss IMAGE_TAG=jax0.4.30-rev1 MAXDIFFUSION_REQUIREMENTS_FILE=requirements_with_jax_ss.txt BUILD_CONTEXT=/dev/shm/docker-context DELETE_LOCAL_IMAGE=true

docker_maxdiffusion_image_upload.sh

@@ -25,6 +25,8 @@
 
 set -e
 
+export LOCAL_IMAGE_NAME=maxdiffusion_base_image
+
 # Set environment variables
 for ARGUMENT in "$@"; do
     IFS='=' read -r KEY VALUE <<< "$ARGUMENT"
@@ -57,18 +59,36 @@ if [[ ! -v MAXDIFFUSION_REQUIREMENTS_FILE ]]; then
   exit 1
 fi
 
+# BUILD_CONTEXT, default to current directory if not set
+BUILD_CONTEXT=${BUILD_CONTEXT:-.}
+
+# Default: Don't delete local image
+DELETE_LOCAL_IMAGE="false"
+
+gcloud auth configure-docker --quiet
+gcloud auth configure-docker us-docker.pkg.dev --quiet
+
 COMMIT_HASH=$(git rev-parse --short HEAD)
 
 echo "Building JAX SS MaxDiffusion at commit hash ${COMMIT_HASH} . . ."  
 
+IMAGE_DATE=$(date +%Y-%m-%d)
+
+FULL_IMAGE_TAG=${IMAGE_TAG}-${IMAGE_DATE}
+
 docker build \
   --build-arg JAX_SS_BASEIMAGE=${BASEIMAGE} \
   --build-arg COMMIT_HASH=${COMMIT_HASH} \
   --build-arg MAXDIFFUSION_REQUIREMENTS_FILE=${MAXDIFFUSION_REQUIREMENTS_FILE} \
   --network=host \
-  -t gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${IMAGE_TAG} \
-  -f ./maxdiffusion_jax_ss_tpu.Dockerfile .
+  -t gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${FULL_IMAGE_TAG} \
+  -f ${BUILD_CONTEXT}/maxdiffusion_jax_ss_tpu.Dockerfile ${BUILD_CONTEXT}
+
+docker push gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${FULL_IMAGE_TAG}
 
-docker push gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${IMAGE_TAG}
+echo "All done, check out your artifacts at: gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${FULL_IMAGE_TAG}"
 
-echo "All done, check out your artifacts at: gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${IMAGE_TAG}"
+if [ "$DELETE_LOCAL_IMAGE" == "true" ]; then
+  docker rmi gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}/tpu:${FULL_IMAGE_TAG}
+  echo "Local image deleted."
+fi