From 0c3dabd68cb8d3a559a68c3fbb7db1de1f5093c8 Mon Sep 17 00:00:00 2001
From: Prisha Jain <prishajain@google.com>
Date: Tue, 15 Sep 2026 10:59:59 +0000
Subject: [PATCH] Fix nightly and stable Docker image build pipeline

- Switch base image from deprecated python:3.12-slim-bullseye to python:3.12-slim-bookworm
- Replace retired google-cloud-sdk apt package with google-cloud-cli
- Add if: always() to nightly build step in UploadDockerImages.yml so a stable build failure does not silently skip nightly publishing
---
 .github/workflows/UploadDockerImages.yml | 7 +++++++
 docker_build_dependency_image.sh         | 2 +-
 maxdiffusion_dependencies.Dockerfile     | 9 +++++----
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml
index 26bf37056..df4f5ecd9 100644
--- a/.github/workflows/UploadDockerImages.yml
+++ b/.github/workflows/UploadDockerImages.yml
@@ -32,8 +32,15 @@ jobs:
     - name: Cleanup old docker images
       run: docker system prune --all --force
     - name: build maxdiffusion stable image
+      id: build_stable
       run: |
         bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable MODE=stable PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_stable
+    # The stable and nightly images are independent artifacts with independent
+    # consumers. Without `if: always()` a failure in the stable build silently
+    # skips the nightly build, which is how the nightly image went stale for
+    # 18 consecutive days while its own build was healthy.
     - name: build maxdiffusion nightly image
+      id: build_nightly
+      if: always()
       run: |
         bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_nightly MODE=nightly PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_nightly
diff --git a/docker_build_dependency_image.sh b/docker_build_dependency_image.sh
index 48685b56d..ab5512139 100644
--- a/docker_build_dependency_image.sh
+++ b/docker_build_dependency_image.sh
@@ -68,7 +68,7 @@ if [[ ${DEVICE} == "gpu" ]]; then
   docker build --network host --build-arg MODE=${MODE} --build-arg JAX_VERSION=$JAX_VERSION --build-arg DEVICE=$DEVICE --build-arg BASEIMAGE=$BASEIMAGE -f ./maxdiffusion_gpu_dependencies.Dockerfile -t ${LOCAL_IMAGE_NAME} .
 else
   # Default to maxdiffusion_dependencies.Dockerfile for non-GPU builds
-  export BASEIMAGE=${BASEIMAGE:-python:3.12-slim-bullseye}
+  export BASEIMAGE=${BASEIMAGE:-python:3.12-slim-bookworm}
   docker build --no-cache \
     --network=host \
     --build-arg MODE=${MODE} \
diff --git a/maxdiffusion_dependencies.Dockerfile b/maxdiffusion_dependencies.Dockerfile
index a87430169..7cd87f8b4 100644
--- a/maxdiffusion_dependencies.Dockerfile
+++ b/maxdiffusion_dependencies.Dockerfile
@@ -1,5 +1,5 @@
-# Use Python 3.12-slim-bullseye as the base image unless overridden
-ARG BASEIMAGE=python:3.12-slim-bullseye
+# Use Python 3.12-slim-bookworm as the base image unless overridden
+ARG BASEIMAGE=python:3.12-slim-bookworm
 FROM $BASEIMAGE
 
 # Environment variable for no-cache-dir and pip root user warning
@@ -23,8 +23,9 @@ RUN apt-get update && apt-get install -y apt-utils git curl gnupg procps iproute
 RUN curl -fsSL https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && \
     echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee /etc/apt/sources.list.d/google-cloud-sdk.list
 
-# Install the Google Cloud SDK
-RUN apt-get update && apt-get install -y google-cloud-sdk && rm -rf /var/lib/apt/lists/*
+# Install the Google Cloud CLI. The `google-cloud-sdk` package was retired
+# upstream and is replaced by `google-cloud-cli`, which provides gcloud/gsutil/bq.
+RUN apt-get update && apt-get install -y google-cloud-cli && rm -rf /var/lib/apt/lists/*
 
 # Install diagnostic and storage dependencies using uv
 RUN python -m uv pip install --system \