Squashes docker layers from the build.

Author: nSircombe

ML-Frameworks/pytorch-aarch64/CHANGELOG.md

@@ -10,6 +10,7 @@ where `YY` is the year, and `MM` the month of the increment.
 ### Added
 
 ### Changed
+ - Updated Dockerfile to minimise layers in the final image.
 
 ### Removed
 

ML-Frameworks/pytorch-aarch64/Dockerfile

@@ -18,7 +18,9 @@
 
 # Specify DOCKER_IMAGE_MIRROR if you want to use a mirror of hub.docker.com
 ARG DOCKER_IMAGE_MIRROR=""
-FROM ${DOCKER_IMAGE_MIRROR}ubuntu:22.04
+FROM ${DOCKER_IMAGE_MIRROR}ubuntu:22.04 AS workshop
+
+ARG USERNAME
 
 ARG TORCH_WHEEL
 ENV TORCH_WHEEL=$TORCH_WHEEL
@@ -44,7 +46,7 @@ RUN apt-get update && apt-get install -y \
     && rm -rf /var/lib/apt/lists/*
 
 # DOCKER_USER for the Docker user
-ENV DOCKER_USER=ubuntu
+ENV DOCKER_USER=${USERNAME}
 
 # Setup default user
 RUN useradd --create-home -s /bin/bash -m $DOCKER_USER && echo "$DOCKER_USER:Portland" | chpasswd && adduser $DOCKER_USER sudo
@@ -112,4 +114,16 @@ RUN sudo mkdir -p /home/ubuntu/gen_ai_utils/ && \
     sudo git config --global --add safe.directory /home/ubuntu/gen_ai_utils/torchchat && \
     sudo git checkout 90749d280bbc116fcc121a1eda1b60f1dba5b675
 
+# Move build into final image as a single layer.
+FROM ${DOCKER_IMAGE_MIRROR}ubuntu:22.04
+
+ARG USERNAME
+
+ENV DOCKER_USER=${USERNAME}
+
+COPY --from=workshop / /
+RUN chown $DOCKER_USER:$DOCKER_USER /home/$DOCKER_USER
+
+USER $DOCKER_USER
+WORKDIR /home/$DOCKER_USER
 CMD ["bash", "-l"]

ML-Frameworks/pytorch-aarch64/build.sh

@@ -50,8 +50,9 @@ torch_wheel_name=$(grep -o "torch-.*.whl" $build_log | head -n -1 | tail -n 1)
 # Use the second to last match, otherwise grep finds itself
 torch_ao_wheel_name=$(grep -o "torchao-.*.whl" $build_log | head -n -1 | tail -n 1)
 
-docker build -t toolsolutions-pytorch:latest \
+DOCKER_BUILDKIT=1 docker build -t toolsolutions-pytorch:latest \
     --build-arg TORCH_WHEEL=results/$torch_wheel_name \
     --build-arg DOCKER_IMAGE_MIRROR \
     --build-arg TORCH_AO_WHEEL=ao/dist/$torch_ao_wheel_name \
+    --build-arg USERNAME=ubuntu \
     .

ML-Frameworks/tensorflow-aarch64/CHANGELOG.md

@@ -10,6 +10,7 @@ where `YY` is the year, and `MM` the month of the increment.
 ### Added
 
 ### Changed
+ - Updated Dockerfile to minimise layers in the final image.
 
 ### Removed
 

ML-Frameworks/tensorflow-aarch64/Dockerfile

@@ -17,7 +17,9 @@
 
 # Specify DOCKER_IMAGE_MIRROR if you want to use a mirror of hub.docker.com
 ARG DOCKER_IMAGE_MIRROR=""
-FROM ${DOCKER_IMAGE_MIRROR}ubuntu:22.04
+FROM ${DOCKER_IMAGE_MIRROR}ubuntu:22.04 AS workshop
+
+ARG USERNAME
 
 ARG TENSORFLOW_WHEEL
 ENV TENSORFLOW_WHEEL=$TENSORFLOW_WHEEL
@@ -76,4 +78,16 @@ RUN pip install -r requirements.txt --no-deps
 # Setup Examples
 COPY --chown=$DOCKER_USER examples/ /home/$DOCKER_USER/
 
+# Move build into final image as a single layer.
+FROM ${DOCKER_IMAGE_MIRROR}ubuntu:22.04
+
+ARG USERNAME
+
+ENV DOCKER_USER=${USERNAME}
+
+COPY --from=workshop / /
+RUN chown $DOCKER_USER:$DOCKER_USER /home/$DOCKER_USER
+
+USER $DOCKER_USER
+WORKDIR /home/$DOCKER_USER
 CMD ["bash", "-l"]