feat: multistage build for the base container

base/jobs/docker/1.0/py3/Dockerfile.cpu

@@ -1,29 +1,53 @@
-# This is based loosely on :
-# https://github.com/aws/amazon-sagemaker-examples/blob/master/advanced_functionality/custom-training-containers/script-mode-container/docker/Dockerfile
-
-FROM public.ecr.aws/lts/ubuntu:22.04
-
+FROM public.ecr.aws/lts/ubuntu:22.04 AS builder
 LABEL maintainer="Amazon Braket"
 LABEL major_version="1"
 
-# prevent stopping by user interaction
-ENV DEBIAN_FRONTEND noninteractive
-ENV DEBCONF_NONINTERACTIVE_SEEN true
+# Set environment variables
+ENV DEBIAN_FRONTEND=noninteractive \
+    PYTHONUNBUFFERED=1 \
+    LANG=C.UTF-8 \
+    LC_ALL=C.UTF-8
 
-ENV PYTHONDONTWRITEBYTECODE=1
-ENV PYTHONUNBUFFERED=1
-ENV PYTHONIOENCODING=UTF-8
-ENV LANG=C.UTF-8
-ENV LC_ALL=C.UTF-8
+# Install Python and pip
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3-pip \
+    python3-dev \
+    build-essential \
+    curl \
+    unzip \
+    && rm -rf /var/lib/apt/lists/*
 
-ARG PYTHON=python3.10
-ARG PYTHON_PIP=python3-pip
-ARG PIP=pip3
-ARG PYTHON_VERSION=3.10.11
+# Copy the dependencies file
+COPY requirements.txt /
 
-RUN apt-get update \
- && apt-get install -y --no-install-recommends \
-    build-essential \
+# Install Python dependencies
+RUN pip3 install --no-cache-dir --upgrade pip setuptools \
+    && pip3 install --no-cache-dir -r requirements.txt
+
+# Install OSS compliance tools
+RUN HOME_DIR=/root \
+    && curl -o ${HOME_DIR}/oss_compliance.zip https://aws-dlinfra-utilities.s3.amazonaws.com/oss_compliance.zip \
+    && unzip ${HOME_DIR}/oss_compliance.zip -d ${HOME_DIR}/ \
+    && cp ${HOME_DIR}/oss_compliance/test/testOSSCompliance /usr/local/bin/testOSSCompliance \
+    && chmod +x /usr/local/bin/testOSSCompliance \
+    && chmod +x ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh \
+    && ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh ${HOME_DIR} python3 \
+    && rm -rf ${HOME_DIR}/oss_compliance*
+
+
+# Stage 1: Final Image
+FROM public.ecr.aws/lts/ubuntu:22.04
+LABEL maintainer="Amazon Braket"
+LABEL major_version="1"
+
+# Set environment variables
+ENV DEBIAN_FRONTEND=noninteractive \
+    PYTHONUNBUFFERED=1 \
+    LANG=C.UTF-8 \
+    LC_ALL=C.UTF-8
+
+# Install runtime dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
     openssh-client \
     openssh-server \
     ca-certificates \
@@ -33,93 +57,41 @@ RUN apt-get update \
     libssl3 \
     openssl \
     unzip \
-    wget \
-    zlib1g-dev \
-    # Install dependent library for OpenCV
     libgtk2.0-dev \
- && rm -rf /var/lib/apt/lists/*
-
-# Install Open MPI
-RUN mkdir /tmp/openmpi \
- && cd /tmp/openmpi \
- && curl -fSsL -O https://download.open-mpi.org/release/open-mpi/v4.0/openmpi-4.0.4.tar.gz \
- && tar zxf openmpi-4.0.4.tar.gz \
- && cd openmpi-4.0.4 \
- && ./configure --enable-orterun-prefix-by-default \
- && make -j $(nproc) all \
- && make install \
- && ldconfig \
- && rm -rf /tmp/openmpi
-
-# Create a wrapper for OpenMPI to allow running as root by default
-RUN mv /usr/local/bin/mpirun /usr/local/bin/mpirun.real \
- && echo '#!/bin/bash' > /usr/local/bin/mpirun \
- && echo 'mpirun.real --allow-run-as-root "$@"' >> /usr/local/bin/mpirun \
- && chmod a+x /usr/local/bin/mpirun
-
-RUN echo "hwloc_base_binding_policy = none" >> /usr/local/etc/openmpi-mca-params.conf \
- && echo "rmaps_base_mapping_policy = slot" >> /usr/local/etc/openmpi-mca-params.conf
-
-ENV LD_LIBRARY_PATH=/usr/local/openmpi/lib:/usr/local/lib:$LD_LIBRARY_PATH
-ENV PATH /usr/local/openmpi/bin/:$PATH
-
-# SSH login fix. Otherwise user is kicked off after login
-RUN sed 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' -i /etc/pam.d/sshd
-
-# Create SSH key.
-RUN mkdir -p /root/.ssh/ \
- && mkdir -p /var/run/sshd \
- && ssh-keygen -q -t rsa -N '' -f /root/.ssh/id_rsa \
- && cp /root/.ssh/id_rsa.pub /root/.ssh/authorized_keys \
- && printf "Host *\n  StrictHostKeyChecking no\n" >> /root/.ssh/config
-
-RUN apt-get update \
- && apt-get install -y --no-install-recommends \
-    libbz2-dev \
-    libc6-dev \
-    libffi-dev \
-    libgdbm-dev \
-    liblzma-dev \
-    libncursesw5-dev \
-    libsqlite3-dev \
-    libssl-dev \
-    tk-dev \
- && rm -rf /var/lib/apt/lists/* \
- && apt-get clean
-
-# Installing Python3
-RUN wget https://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VERSION.tgz \
- && tar -xvf Python-$PYTHON_VERSION.tgz \
- && cd Python-$PYTHON_VERSION \
- && ./configure && make && make install \
- && rm -rf ../Python-$PYTHON_VERSION*
-
-# Upgrading pip and creating symbolic link for python3
-RUN ${PIP} --no-cache-dir install --upgrade pip setuptools
-RUN ln -s $(which ${PYTHON}) /usr/local/bin/python \
- && ln -s $(which ${PIP}) /usr/bin/pip
-
-WORKDIR /
-
-RUN apt-get update && apt-get -y install cmake protobuf-compiler
-
-# Copy the dependencies file
-COPY requirements.txt /
-# Installing our custom python libraries
-RUN ${PIP} install --no-cache --upgrade \
-        -r requirements.txt
-
-RUN HOME_DIR=/root \
- && curl -o ${HOME_DIR}/oss_compliance.zip https://aws-dlinfra-utilities.s3.amazonaws.com/oss_compliance.zip \
- && unzip ${HOME_DIR}/oss_compliance.zip -d ${HOME_DIR}/ \
- && cp ${HOME_DIR}/oss_compliance/test/testOSSCompliance /usr/local/bin/testOSSCompliance \
- && chmod +x /usr/local/bin/testOSSCompliance \
- && chmod +x ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh \
- && ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh ${HOME_DIR} ${PYTHON} \
- && rm -rf ${HOME_DIR}/oss_compliance*
-
-# Setup our entry point
-
+    python3-pip \
+    libopenmpi-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# SSH login fix and key generation
+RUN sed -i 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' /etc/pam.d/sshd \
+    && mkdir -p /root/.ssh/ /var/run/sshd \
+    && ssh-keygen -q -t rsa -N '' -f /root/.ssh/id_rsa \
+    && cp /root/.ssh/id_rsa.pub /root/.ssh/authorized_keys \
+    && printf "Host *\n  StrictHostKeyChecking no\n" >> /root/.ssh/config
+
+# Copy Python dependencies from the builder stage
+COPY --from=builder /usr/local/lib/python3.10 /usr/local/lib/python3.10
+COPY --from=builder /usr/local/bin/pip3 /usr/local/bin/pip3
+COPY --from=builder /usr/local/bin/testOSSCompliance /usr/local/bin/testOSSCompliance
+
+# Ensure Python links are correct
+RUN ln -s /usr/bin/python3 /usr/local/bin/python \
+    && ln -s /usr/bin/python3 /usr/local/bin/python3
+
+# Install Open MPI wrapper to allow running as root
+RUN mv /usr/bin/mpirun /usr/bin/mpirun.real \
+    && echo '#!/bin/bash' > /usr/bin/mpirun \
+    && echo 'mpirun.real --allow-run-as-root "$@"' >> /usr/bin/mpirun \
+    && chmod a+x /usr/bin/mpirun
+
+# Set Open MPI configurations
+RUN echo "hwloc_base_binding_policy = none" >> /etc/openmpi/openmpi-mca-params.conf \
+    && echo "rmaps_base_mapping_policy = slot" >> /etc/openmpi/openmpi-mca-params.conf
+
+# Set up environment variables
+ENV LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu/openmpi/lib:$LD_LIBRARY_PATH \
+    PATH=/usr/lib/x86_64-linux-gnu/openmpi/bin:$PATH
+
+# Copy the braket_container.py
 COPY braket_container.py /opt/ml/code/braket_container.py
-
-ENV SAGEMAKER_PROGRAM braket_container.py
+ENV SAGEMAKER_PROGRAM=braket_container.py