fix(docker): Fix CPU/TPU nightly Docker build dependencies

[Toshi-31]

Summary

Fixes the CPU/TPU Docker image build pipeline so that Google Cloud Build can resume publishing daily maxdiffusion_jax_stable_stack_nightly images.

Root Causes

1. Base Image EOL: python:3.12-slim-bullseye (Debian 11) package mirrors moved, breaking apt-get update.
2. Retired APT Package: google-cloud-sdk was retired upstream by Google Cloud apt repositories in favor of google-cloud-cli, causing apt-get install -y google-cloud-sdk to exit with code 100.

Changes

1. Upgrade BASEIMAGE default from python:3.12-slim-bullseye to python:3.12-slim-bookworm (Debian 12) in maxdiffusion_dependencies.Dockerfile and docker_build_dependency_image.sh.
2. Replace google-cloud-sdk with google-cloud-cli in maxdiffusion_dependencies.Dockerfile, modernizing the GPG key dearmoring step.

How This Was Identified

1. Automated Pipeline Outage: The scheduled nightly image build workflow (.github/workflows/build_runner.yml) was failing during the Docker build stage, halting the publication of nightly runner images (maxdiffusion_jax_stable_stack_nightly) to GCR.
2. Build Log Analysis:
   • apt-get install -y google-cloud-sdk exited with code 100 (E: Unable to locate package google-cloud-sdk) because Google Cloud's apt repositories deprecated the package name google-cloud-sdk in favor of google-cloud-cli.
   • The Debian 11 python:3.12-slim-bullseye base image was hitting repository mirror archive warnings and outdated keyring handling for Google Cloud apt sources.

How This Fix Has Been Verified

1. GitHub Actions CI Workflow:
   • The automated build-image workflow on this PR ran end-to-end and successfully built both CPU and TPU runner images: Workflow Run 36115942148 / Job 108009968725.
2. Container Build Verification:
   • Verified clean local/container builds for both CPU and TPU targets (bash docker_build_dependency_image.sh MODE=stable DEVICE=tpu) confirming google-cloud-cli installs without apt errors and dependencies resolve on Debian 12 (bookworm).
3. Cluster Runtime Verification:
   • Deployed the resulting container image on the TPU benchmark cluster via XPK, confirming that the container initializes cleanly, detects TPU devices, and launches training workloads without missing library or runtime environment errors.

Some pylink formatting changes have also been done because unit tests were failing.

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

[gemini-code-assist]

Code Review

This pull request updates the base Docker images from Debian Bullseye to Bookworm (python:3.12-slim-bookworm) and replaces the deprecated google-cloud-sdk package with google-cloud-cli. It also modernizes the GPG key addition process in the GPU Dockerfile to avoid using the deprecated apt-key utility, alongside some minor formatting cleanups in Python files. Feedback on the GPU Dockerfile highlights a potential build failure if curl and gnupg are not pre-installed in the base image, suggesting they be explicitly installed before fetching the GPG key.

[prishajain1]

Can you add how this was identified and how this fix has been verified to work

[Toshi-31]

Can you add how this was identified and how this fix has been verified to work

Done

[mbohlool]

Thanks for fixing the nightly build pipeline! A couple of minor suggestions and a nit for consideration / follow-up:

💡 Suggestions

1. GPU Dockerfile parity (maxdiffusion_gpu_dependencies.Dockerfile):
   maxdiffusion_gpu_dependencies.Dockerfile still references the retired google-cloud-sdk package and uses deprecated apt-key. While this PR is scoped to CPU/TPU, updating the GPU Dockerfile similarly (ensuring curl and gnupg are installed prior to dearmoring) would prevent GPU build breaks down the road.
2. APT layer consolidation (maxdiffusion_dependencies.Dockerfile):
   Lines 20–27 currently run two separate apt-get update && apt-get install passes. Configuring the Google Cloud repository before the initial install could collapse them into a single RUN layer to shave off build time and reduce image size.

🟢 Nit

• Stale comment (maxdiffusion_dependencies.Dockerfile:L9): The comment reads # Set environment variables for Google Cloud SDK and Python 3.10, but line 10 sets PYTHON_VERSION=3.12.

[Toshi-31]

Thanks for fixing the nightly build pipeline! A couple of minor suggestions and a nit for consideration / follow-up:

💡 Suggestions

1. GPU Dockerfile parity (maxdiffusion_gpu_dependencies.Dockerfile):
   maxdiffusion_gpu_dependencies.Dockerfile still references the retired google-cloud-sdk package and uses deprecated apt-key. While this PR is scoped to CPU/TPU, updating the GPU Dockerfile similarly (ensuring curl and gnupg are installed prior to dearmoring) would prevent GPU build breaks down the road.
2. APT layer consolidation (maxdiffusion_dependencies.Dockerfile):
   Lines 20–27 currently run two separate apt-get update && apt-get install passes. Configuring the Google Cloud repository before the initial install could collapse them into a single RUN layer to shave off build time and reduce image size.

🟢 Nit

• Stale comment (maxdiffusion_dependencies.Dockerfile:L9): The comment reads # Set environment variables for Google Cloud SDK and Python 3.10, but line 10 sets PYTHON_VERSION=3.12.

Done