Merge pull request #428 from puneetmatharu/beta-test-debian-pytorch-image

Switch to slim Python/Debian image for PyTorch runner env

Author: puneetmatharu

ML-Frameworks/pytorch-aarch64/Dockerfile

@@ -4,131 +4,90 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # Specify DOCKER_IMAGE_MIRROR if you want to use a mirror of hub.docker.com
+ARG PYTHON_IMAGE=python:3.12-slim
 ARG DOCKER_IMAGE_MIRROR=""
-FROM ${DOCKER_IMAGE_MIRROR}ubuntu:24.04 AS workshop
 
-ARG USERNAME
+# ============================
+# Workshop
+# ============================
+FROM ${DOCKER_IMAGE_MIRROR}${PYTHON_IMAGE} AS workshop
 
+ARG USERNAME
 ARG TORCH_WHEEL
-ENV TORCH_WHEEL=$TORCH_WHEEL
-
 ARG TORCH_AO_WHEEL
-ENV TORCH_AO_WHEEL=$TORCH_AO_WHEEL
-
-RUN if ! [ "$(arch)" = "aarch64" ] ; then exit 1; fi
-
-RUN apt-get update && apt-get install -y \
-    # We need pip to install things, this will also bring in a minimal python3
-    python3-pip \
-    # So that we can create a virtual environment
-    python3-venv \
-    # So that we can call python instead of python3
-    python-is-python3 \
-    # To allow users to install new things if they want
-    sudo \
-    && rm -rf /var/lib/apt/lists/*
-
-# DOCKER_USER for the Docker user
+
+ENV DEBIAN_FRONTEND=noninteractive
 ENV DOCKER_USER=${USERNAME}
 
-# Create user only if it doesn't already exist
-RUN id "$DOCKER_USER" >/dev/null 2>&1 || useradd --create-home -s /bin/bash -m "$DOCKER_USER"
-
-# Set password and add to sudo group
-RUN echo "$DOCKER_USER:ToolSolutionsPyTorch" | chpasswd && adduser "$DOCKER_USER" sudo || true
-
-RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
-RUN echo 'debconf debconf/frontend select Noninteractive' | debconf-set-selections
-
-# Import profile for bash
-COPY bash_profile /home/$DOCKER_USER/.bash_profile
-RUN chown $DOCKER_USER:$DOCKER_USER /home/$DOCKER_USER/.bash_profile
-
-# Add welcome message to warn about dev quality
-COPY welcome.txt /home/$DOCKER_USER/
-RUN echo '[ ! -z "$TERM" -a -r /home/$DOCKER_USER/welcome.txt ] && cat /home/$DOCKER_USER/welcome.txt' >> /etc/bash.bashrc
-RUN echo 'export PATH="$HOME/.local/bin:$PATH"' >>  /etc/bash.bashrc
-
-# Grab the SECURITY.md from the root directory
-COPY --from=rootdir SECURITY.md /home/$DOCKER_USER/
-
-# Remove system Python stuff. Should be safe to wipe after the line above, because
-# python3 -m pip now uses the /usr/local install
-RUN apt-get update && apt-get purge -y \
-      python3-pip \
-      python3-setuptools \
-      python3-pkg-resources \
-      python3-wheel \
-      python3-distutils \
-      python3-lib2to3 \
-      python3-dev \
-      python3.12-dev \
-    && apt-get autoremove -y \
-    && rm -rf /var/lib/apt/lists/*
-
-# Move to userland
-WORKDIR /home/$DOCKER_USER
+RUN test "$(arch)" = "aarch64"
+
+# Install OS dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends curl && rm -rf /var/lib/apt/lists/*
+
+# Create user
+RUN set -eux && id "$DOCKER_USER" >/dev/null 2>&1 || useradd --create-home -s /bin/bash "$DOCKER_USER"
+
+# Copy bash profile and welcome text into user home
+COPY --chown=$DOCKER_USER:$DOCKER_USER bash_profile /home/$DOCKER_USER/.bash_profile
+COPY --chown=$DOCKER_USER:$DOCKER_USER welcome.txt /home/$DOCKER_USER/welcome.txt
+
+# Switch to userland
 USER $DOCKER_USER
+WORKDIR /home/$DOCKER_USER
 
-# Create a per-user virtualenv and use that for everything Python
+# Create virtual environment
 RUN python -m venv /home/$DOCKER_USER/.venv
+ENV PATH="/home/$DOCKER_USER/.venv/bin:${PATH}"
 
-# Make the venv python/pip first on PATH for all subsequent layers and at runtime
-ENV PATH="/home/$DOCKER_USER/.venv/bin:$PATH"
+# Install uv for quicker package installations
+RUN pip install uv
 
 # Update to newer pip/setuptools/wheel (setuptools >= 70.0.0 due to CVE-2024-6345
 # and CVE-2025-47273, wheel >= 0.38.0 due to CVE-2022-40898) and delete old system
 # version (we essentially use apt:python3-pip to bootstrap pip)
-RUN pip install --upgrade pip~=25.2 setuptools~=78.1.1 wheel~=0.45.1
-
-# Base requirements for examples, excluding torch and torch*
-COPY requirements.txt ./
-RUN pip install -r requirements.txt
-
-# Check TORCH_WHEEL was set and copy
-RUN test -n "$TORCH_WHEEL"
-COPY $TORCH_WHEEL /home/$DOCKER_USER/
+RUN uv pip install --upgrade pip~=25.2 setuptools~=78.1.1 wheel~=0.45.1
 
-# Check TORCH_AO_WHEEL was set and copy
-RUN test -n "$TORCH_AO_WHEEL"
-COPY $TORCH_AO_WHEEL /home/$DOCKER_USER/
+# Install non-torch requirements
+COPY --chown=$DOCKER_USER:$DOCKER_USER requirements.txt .
+RUN uv pip install -r requirements.txt
 
-# Install torch* packages, these should be the latest stable (but pinned to
-# minor). We need --no-deps here because the torch* packages depend on the
-# corresponding version of torch. Note: if you add something to this list, you
-# will need to manually add their dependencies. We don't use the nightly
-# versions which corresponding to our torch build because they can disappear,
-# and we usually don't need features from the nightlies.
-# Note: torchvision is pinned to a nightly build, this can be updated
-# at the next vision release, and the `--extra-index-url` removed.
-RUN pip install --pre torchvision==0.25.0.dev20260130 --index-url https://download.pytorch.org/whl/nightly/cpu --no-deps
+# Bring wheels into image
+RUN test -n "${TORCH_WHEEL}" && test -n "${TORCH_AO_WHEEL}"
+COPY --chown=$DOCKER_USER:$DOCKER_USER ${TORCH_WHEEL} /home/$DOCKER_USER/
+COPY --chown=$DOCKER_USER:$DOCKER_USER ${TORCH_AO_WHEEL} /home/$DOCKER_USER/
 
-# We need --no-deps because the torch version won't match the versions on torch*
-RUN pip install "$(basename "$TORCH_WHEEL")" --no-deps \
-    && rm "$(basename "$TORCH_WHEEL")"
+# Install wheels
+RUN set -eux && uv pip install --no-deps "$(basename "$TORCH_WHEEL")" && rm "$(basename "$TORCH_WHEEL")"
+RUN set -eux && uv pip install --no-deps "$(basename "$TORCH_AO_WHEEL")" && rm "$(basename "$TORCH_AO_WHEEL")"
+RUN uv pip install --pre torchvision==0.25.0.dev20260130 --index-url https://download.pytorch.org/whl/nightly/cpu --no-deps
 
-# We need --no-deps because this won't match the torch version
-RUN pip install "$(basename "$TORCH_AO_WHEEL")" --no-deps \
-    && rm "$(basename "$TORCH_AO_WHEEL")"
+# Copy examples/tests into image
+COPY --chown=$DOCKER_USER:$DOCKER_USER examples/ /home/$DOCKER_USER/
+COPY --chown=$DOCKER_USER:$DOCKER_USER pytorch/test /home/$DOCKER_USER/pytorch/test
 
-# Setup Examples and tests
-COPY examples/ /home/$DOCKER_USER/
-COPY pytorch/test /home/$DOCKER_USER/pytorch/test
-
-# Move build into final image as a single layer.
-FROM ${DOCKER_IMAGE_MIRROR}ubuntu:24.04
+# ============================
+# Final flat image
+# ============================
+FROM ${DOCKER_IMAGE_MIRROR}${PYTHON_IMAGE}
 
 ARG USERNAME
-
+ENV DEBIAN_FRONTEND=noninteractive
 ENV DOCKER_USER=${USERNAME}
 
-COPY --from=workshop / /
-RUN chown $DOCKER_USER:$DOCKER_USER /home/$DOCKER_USER
+# Runtime OS bits + UI
+RUN set -eux && \
+    apt-get update && \
+    rm -rf /var/lib/apt/lists/* && \
+    if ! id "$DOCKER_USER" >/dev/null 2>&1; then useradd --create-home -s /bin/bash "$DOCKER_USER"; fi && \
+    echo '[ -n "$TERM" -a -r "$HOME/welcome.txt" ] && cat "$HOME/welcome.txt"' >> /etc/bash.bashrc && \
+    echo 'export PATH="$HOME/.local/bin:$HOME/.venv/bin:$PATH"' >> /etc/bash.bashrc
+
+# Bring in prepped env + code
+COPY --from=workshop --chown=$DOCKER_USER:$DOCKER_USER /home/$DOCKER_USER /home/$DOCKER_USER
 
 USER $DOCKER_USER
 WORKDIR /home/$DOCKER_USER
 
-# Ensure the venv is on PATH in the final image as well
-ENV PATH="/home/$DOCKER_USER/.venv/bin:$PATH"
+ENV PATH="/home/$DOCKER_USER/.venv/bin:${PATH}"
 
 CMD ["bash", "-l"]

ML-Frameworks/pytorch-aarch64/dockerize.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-# SPDX-FileCopyrightText: Copyright 2024, 2025 Arm Limited and affiliates.
+# SPDX-FileCopyrightText: Copyright 2024-2026 Arm Limited and affiliates.
 #
 # SPDX-License-Identifier: Apache-2.0
 
@@ -29,7 +29,7 @@ docker buildx \
     --build-arg DOCKER_IMAGE_MIRROR \
     --build-arg TORCH_WHEEL=$1 \
     --build-arg TORCH_AO_WHEEL=$2 \
-    --build-arg USERNAME=ubuntu \
+    --build-arg USERNAME=debian \
     .
 
 [[ $* == *--build-only* ]] && exit 0