fix(dependencies): reduce container image vulnerabilities in TPU stable images

[khatwanimohit]

Description

Addresses b/532071207 to remediate vulnerabilities in the official TPU stable container image (maxtext_jax_stable).

Summary of Changes:

1. src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile:
   • Added apt-get upgrade -y during Debian system dependency installation to pick up security fixes for base OS libraries (libexpat1, libarchive13).
   • Upgraded core Python packaging tools (pip, setuptools, wheel, uv) and removed the obsolete ensurepip/_bundled wheel cache.
2. src/dependencies/scripts/setup.sh:
   • Added apt-get upgrade -y and apt upgrade -y when installing gcsfuse and system tools to ensure the latest patched gcsfuse Go binary is installed.
   • Added cleanup_unneeded_build_artifacts() post-installation to:
     • Remove flaxlib_src/Cargo.lock left in site-packages from the flax wheel build context.
     • Delete embedded virtualenv seed wheels containing outdated pip and setuptools.

BUGS: b/532071207

Tests

1. Verification Inside the Built Container

• gcsfuse: Upgraded to 3.11.3 (compiled with Go 1.26.7). Eliminates CRITICAL CVE-2026-39821 and 9 other HIGH Go CVEs from b/532071207.
• flaxlib_src / pyo3: Removed unused Cargo.lock (ls /usr/local/lib/python3.12/site-packages/flaxlib_src returns no such file or directory). Eliminates 3 pyo3 CVEs (CVE-2025-46746, CVE-2025-46747, CVE-2025-46748).
• black: Upgraded to 26.5.1 (safe version, eliminates CVE-2026-4444).
• pip & setuptools: Upgraded to pip 26.2.1 and setuptools 84.0.0; embedded .whl files deleted (all 11 bundled CVEs eliminated).
• Debian OS Libraries: libexpat1 upgraded to 2.5.0-1+deb12u3 and libarchive13 upgraded to 3.6.2-1+deb12u5 (eliminates 15 expat CVEs and 2 libarchive CVEs).

2. Functional & Smoke Tests Inside Container

• python3 -c "import jax; import flax": passes cleanly (JAX: 0.11.1, FLAX: 0.12.9).
• python3 -c "import maxtext": imports successfully from /deps/src/maxtext/__init__.py.
• pytest /deps/tests/unit/checkpoint_context_test.py: 18 passed in 13.31s.

3. Google Artifact Registry Vulnerability Scan

Pushed test build to staging Artifact Registry to run official GCP vulnerability scanning:

• Total Findings: 320 -> 277 (-43 CVEs)
• Fixable Findings: 44 -> 1 (-97.7% reduction; all CRITICALs, Python, Rust, and Debian fixable CVEs eliminated)
• CRITICAL Fixable CVEs: 1 -> 0 (-100%)
• HIGH Fixable CVEs: 14 -> 1 (-92.9%)
• MED / LOW Fixable CVEs: 29 -> 0 (-100%)

Checklist

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[gemini-code-assist]

Code Review

This pull request introduces package upgrades and cleanup routines in both the Dockerfile and setup script to remove build artifacts that trigger security scanners. The feedback highlights several improvement opportunities: ensuring non-interactive package installations in both the Dockerfile and setup script (especially when using sudo bash), dynamically locating the ensurepip directory instead of hardcoding the Python version path, using --no-cache-dir with pip, and avoiding Python list comprehensions for side effects in the cleanup script.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!