Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Sep 4, 2026
+27 −7
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,10 @@ FROM $BASEIMAGE
# Install system dependencies including C++20 compiler for vLLM
RUN if [ -f /etc/os-release ] && grep -q "bullseye" /etc/os-release; then \
echo "deb http://deb.debian.org/debian bookworm main" > /etc/apt/sources.list.d/bookworm.list && \
apt-get update && apt-get install -y --no-install-recommends -t bookworm gcc-12 g++-12 build-essential cmake ninja-build curl gnupg && \
apt-get update && DEBIAN_FRONTEND=noninteractive apt-get upgrade -y && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends -t bookworm gcc-12 g++-12 build-essential cmake ninja-build curl gnupg && \
rm -rf /var/lib/apt/lists/*; \
else \
apt-get update && apt-get install -y --no-install-recommends gcc-12 g++-12 build-essential cmake ninja-build curl gnupg && \
apt-get update && DEBIAN_FRONTEND=noninteractive apt-get upgrade -y && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends gcc-12 g++-12 build-essential cmake ninja-build curl gnupg && \
rm -rf /var/lib/apt/lists/*; \
fi
Comment thread
khatwanimohit marked this conversation as resolved.

Expand All @@ -23,6 +23,10 @@ RUN apt-get update && apt-get install -y google-cloud-cli && rm -rf /var/lib/apt
# Set the default Python version to 3.12
RUN update-alternatives --install /usr/bin/python3 python3 /usr/local/bin/python3.12 1

# Upgrade pip, setuptools, wheel, uv and clean up ensurepip bundled cache
RUN python3 -m pip install --upgrade --no-cache-dir pip setuptools wheel uv && \
python3 -c 'import ensurepip, os, shutil; shutil.rmtree(os.path.join(os.path.dirname(ensurepip.__file__), "_bundled"), ignore_errors=True)'

# Set environment variables for Google Cloud SDK, Python 3.12, and GCC 12
ENV PATH="/usr/local/google-cloud-sdk/bin:/usr/local/bin/python3.12:${PATH}"
ENV CC=gcc-12
Expand Down
26 changes: 21 additions & 5 deletions src/dependencies/scripts/setup.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,18 +125,24 @@ fi
echo "Python version check passed. Continuing with script."
echo "--------------------------------------------------"

apt-get update && apt-get install -y sudo
apt-get update && apt-get upgrade -y && apt-get install -y sudo
(sudo bash || bash) <<'EOF'
apt update && \
apt install -y numactl lsb-release gnupg curl net-tools iproute2 procps lsof git ethtool && \
# sudo strips the exports set by the parent shell, so re-export them here to
# keep apt-get non-interactive.
export DEBIAN_FRONTEND=noninteractive
export NEEDRESTART_SUSPEND=1
export NEEDRESTART_MODE=l
apt-get update && \
apt-get upgrade -y && \
apt-get install -y numactl lsb-release gnupg curl net-tools iproute2 procps lsof git ethtool && \
export GCSFUSE_REPO=gcsfuse-`lsb_release -c -s`
echo "deb https://packages.cloud.google.com/apt $GCSFUSE_REPO main" | tee /etc/apt/sources.list.d/gcsfuse.list
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
apt update -y && apt -y install gcsfuse
apt-get update -y && apt-get -y install gcsfuse
rm -rf /var/lib/apt/lists/*
EOF

python3 -m pip install -U setuptools wheel uv
python3 -m pip install --no-cache-dir -U pip setuptools wheel uv

# Set environment variables
for ARGUMENT in "$@"; do
Expand Down Expand Up @@ -203,6 +209,14 @@ install_maxtext_package_without_deps() {
fi
}

cleanup_unneeded_build_artifacts() {
echo "Cleaning up build artifacts that trigger security scanners..."
# Clean up flaxlib_src Cargo.lock / unneeded source files
rm -rf /usr/local/lib/python3*/site-packages/flaxlib_src
# Clean up embedded virtualenv seed wheels if virtualenv was installed
find /usr/local /root -name "*.whl" -path "*/virtualenv/seed/wheels/embed/*" -delete 2>/dev/null || true
}

install_maxtext_with_deps() {
if [[ "$DEVICE" != "tpu" && "$DEVICE" != "gpu" ]]; then
echo -e "\n\nError: DEVICE must be either 'tpu' or 'gpu'.\n\n"
Expand All @@ -219,6 +233,7 @@ install_maxtext_with_deps() {
python3 -m src.dependencies.scripts.install_pre_train_extra_deps

install_maxtext_package_without_deps
cleanup_unneeded_build_artifacts
}

install_post_training_deps() {
Expand All @@ -243,6 +258,7 @@ install_post_training_deps() {
echo "Installing requirements from $dep_name"
UV_TORCH_BACKEND=cpu python3 -m uv pip install --resolution=lowest -r "$dep_name"
python3 -m src.dependencies.scripts.install_post_train_extra_deps
cleanup_unneeded_build_artifacts
}

# ---------- Post-Training workflow installation ----------
Expand Down
Loading