[PyTorch][Training][EC2][SageMaker]PyTorch 2.8 Currency Release

pytorch/training/docker/2.8/py3/Dockerfile.cpu

@@ -44,6 +44,11 @@ LABEL dlc_major_version="1"
 ARG PYTHON
 ARG PYTHON_VERSION
 ARG PYTHON_SHORT_VERSION
+ARG PYTORCH_VERSION
+ARG TORCHTNT_VERSION
+ARG TORCHDATA_VERSION
+ARG TORCHAUDIO_VERSION
+ARG TORCHVISION_VERSION
 
 ARG OPEN_MPI_VERSION
 
@@ -133,7 +138,6 @@ RUN rm -rf /root/.ssh/ \
  && cp /root/.ssh/id_rsa.pub /root/.ssh/authorized_keys \
  && printf "Host *\n StrictHostKeyChecking no\n" >> /root/.ssh/config
 
-
 # install python
 RUN cd /tmp/ \
 && wget https://www.python.org/ftp/python/${PYTHON_VERSION}/Python-${PYTHON_VERSION}.tgz \
@@ -185,12 +189,30 @@ RUN pip install --no-cache-dir \
     "setuptools>=70.0.0" \
     "urllib3>=2.5.0" \
     "awscli" \
-    # opencv-python 4.12.0.88 reuqires numpy<2.3.0, which is not compatible with previous prod image(2.3.1)
     opencv-python==4.11.0.86 \
     mpi4py \
     jinja2>=3.1.6 \
     tornado>=6.5.1
 
+# Install PyTorch
+RUN pip install --no-cache-dir -U torch==${PYTORCH_VERSION} \
+    torchvision==${TORCHVISION_VERSION} \
+    torchaudio==${TORCHAUDIO_VERSION} \
+    --index-url https://download.pytorch.org/whl/cpu \
+    && pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
+    torchdata==${TORCHDATA_VERSION} \
+    s3torchconnector \
+    fastai \
+    accelerate \
+    # pin numpy requirement for fastai dependency
+    # requires explicit declaration of spacy, thic, blis
+    spacy \
+    # pin thinc due to incompatibility with numpy 1.26.4 (sagemaker doesn't support latest numpy)
+    thinc==8.3.4  \
+    blis \
+    numpy \
+ && pip uninstall -y dataclasses
+
 RUN curl -o /license.txt https://aws-dlc-licenses.s3.amazonaws.com/pytorch-2.8/license.txt
 
 COPY deep_learning_container.py /usr/local/bin/deep_learning_container.py
@@ -221,46 +243,10 @@ RUN rm -rf /root/.cache | true
 
 FROM common AS ec2
 
-ARG PYTHON
-ARG PYTHON_SHORT_VERSION
-ARG PYTORCH_VERSION
-ARG TORCHTNT_VERSION
-ARG TORCHDATA_VERSION
-ARG TORCHAUDIO_VERSION
-ARG TORCHVISION_VERSION
-
 WORKDIR /
 
-# Install PyTorch
-RUN pip install --no-cache-dir -U torch==${PYTORCH_VERSION} \
-    torchvision==${TORCHVISION_VERSION} \
-    torchaudio==${TORCHAUDIO_VERSION} \
-    --index-url https://download.pytorch.org/whl/cpu \
-    && pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
-    torchdata==${TORCHDATA_VERSION} \
-    s3torchconnector \
-    fastai \
-    accelerate \
-    # pin numpy requirement for fastai dependency
-    # requires explicit declaration of spacy, thic, blis
-    spacy \
-    thinc  \
-    blis \
-    numpy \
- && pip uninstall -y dataclasses
-
-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* \
- && rm -rf /tmp/tmp*
-
-# Removing the cache as it is needed for security verification
-RUN rm -rf /root/.cache | true
+COPY setup_oss_compliance.sh setup_oss_compliance.sh
+RUN bash setup_oss_compliance.sh ${PYTHON} && rm setup_oss_compliance.sh
 
 COPY dockerd_entrypoint.sh /usr/local/bin/dockerd_entrypoint.sh
 RUN chmod +x /usr/local/bin/dockerd_entrypoint.sh
@@ -289,36 +275,10 @@ FROM common AS sagemaker
 LABEL maintainer="Amazon AI"
 LABEL dlc_major_version="1"
 
-ARG PYTHON
-ARG PYTHON_SHORT_VERSION
-ARG PYTORCH_VERSION
-ARG TORCHTNT_VERSION
-ARG TORCHDATA_VERSION
-ARG TORCHAUDIO_VERSION
-ARG TORCHVISION_VERSION
-
 ENV SAGEMAKER_TRAINING_MODULE=sagemaker_pytorch_container.training:main
 
 WORKDIR /
 
-# Install PyTorch
-RUN pip install --no-cache-dir -U torch==${PYTORCH_VERSION} \
-    torchvision==${TORCHVISION_VERSION} \
-    torchaudio==${TORCHAUDIO_VERSION} \
-    --index-url https://download.pytorch.org/whl/cpu \
-    && pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
-    torchdata==${TORCHDATA_VERSION} \
-    s3torchconnector \
-    fastai \
-    accelerate \
-    # pin numpy requirement for fastai dependency
-    # requires explicit declaration of spacy, thic, blis
-    spacy \
-    thinc==8.3.4  \
-    blis \
-    numpy \
- && pip uninstall -y dataclasses
-
 # Install SM packages
 RUN pip install --no-cache-dir -U \
     smclarify \
@@ -345,18 +305,8 @@ COPY start_with_right_hostname.sh /usr/local/bin/start_with_right_hostname.sh
 
 RUN chmod +x /usr/local/bin/start_with_right_hostname.sh
 
-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* \
- && rm -rf /tmp/tmp*
-
-# Removing the cache as it is needed for security verification
-RUN rm -rf /root/.cache | true
+COPY setup_oss_compliance.sh setup_oss_compliance.sh
+RUN bash setup_oss_compliance.sh ${PYTHON} && rm setup_oss_compliance.sh
 
 ENTRYPOINT ["bash", "-m", "start_with_right_hostname.sh"]
 CMD ["/bin/bash"]

pytorch/training/docker/2.8/py3/cu129/Dockerfile.gpu

@@ -7,21 +7,9 @@ ARG TORCHAUDIO_VERSION=2.8.0
 ARG TORCHVISION_VERSION=0.23.0
 ARG TORCHDATA_VERSION=0.11.0
 
-ARG GDRCOPY_VERSION=2.5
+ARG GDRCOPY_VERSION=2.5.1
 ARG TE_VERSION=2.5
-ARG FLASH_ATTN_VERSION=2.8.2
-
-FROM public.ecr.aws/deep-learning-containers/base:12.9.1-gpu-py312-ubuntu22.04-ec2 AS base_image
-# base has EFA, PYTHON and CUDA 12.9
-
-# This arg required to stop docker build waiting for region configuration while installing tz data from ubuntu 20
-ENV DEBIAN_FRONTEND=noninteractive
-
-RUN apt-get update \
- && apt-get upgrade -y \
- && apt-get autoremove -y \
- && apt-get clean \
- && rm -rf /var/lib/apt/lists/*
+ARG FLASH_ATTN_VERSION=2.8.3
 
 #################################################################
 #   ____
@@ -37,7 +25,8 @@ RUN apt-get update \
 #                      |___/                        |_|
 #################################################################
 
-FROM base_image AS common
+FROM public.ecr.aws/deep-learning-containers/base:12.9.1-gpu-py312-ubuntu22.04-ec2 AS common
+# base has EFA, PYTHON and CUDA 12.9
 
 LABEL maintainer="Amazon AI"
 LABEL dlc_major_version="1"
@@ -138,7 +127,7 @@ RUN pip install --no-cache-dir -U torch==${PYTORCH_VERSION} \
     # pin numpy requirement for fastai dependency
     # requires explicit declaration of spacy, thic, blis
     spacy \
-    #thinc 8.3.6 is not compatible with numpy 1.26.4 (sagemaker doesn't support latest numpy)
+    # pin thinc due to incompatibility with numpy 1.26.4 (sagemaker doesn't support latest numpy)
     thinc==8.3.4 \
     blis \
     numpy \
@@ -155,9 +144,6 @@ RUN pip install --no-cache-dir git+https://github.com/NVIDIA/TransformerEngine.g
 
 RUN curl -o /license.txt https://aws-dlc-licenses.s3.amazonaws.com/pytorch-2.8/license.txt
 
-COPY deep_learning_container.py /usr/local/bin/deep_learning_container.py
-RUN chmod +x /usr/local/bin/deep_learning_container.py
-
 COPY start_cuda_compat.sh /usr/local/bin/start_cuda_compat.sh
 RUN chmod +x /usr/local/bin/start_cuda_compat.sh
 
@@ -171,10 +157,23 @@ RUN echo 'source /usr/local/bin/bash_telemetry.sh' >> /etc/bash.bashrc
 RUN cd /tmp \
  && git clone https://github.com/NVIDIA/gdrcopy.git -b v${GDRCOPY_VERSION} \
  && cd gdrcopy \
- && sed -ie '12s@$@ -L $(CUDA)/lib64/stubs@' tests/Makefile \
+ && sed -ie '13s@$@ -L $(CUDA)/lib64/stubs@' tests/Makefile \
  && CUDA=${CUDA_HOME} make install \
  && rm -rf /tmp/gdrcopy
 
+# Install common packages used by both EC2 and SageMaker
+RUN apt-get update \
+ && apt-get install -y --no-install-recommends \
+    curl \
+    wget \
+    git \
+    jq \
+    emacs \
+    vim \
+    unzip \
+ && rm -rf /var/lib/apt/lists/* \
+ && apt-get clean
+
 # Removing the cache as it is needed for security verification
 RUN rm -rf /root/.cache | true
 
@@ -199,18 +198,6 @@ ARG PYTHON
 
 WORKDIR /
 
-RUN apt-get update \
- && apt-get install -y --no-install-recommends \
-    curl \
-    wget \
-    git \
-    jq \
-    emacs \
-    vim \
-    unzip \
- && rm -rf /var/lib/apt/lists/* \
- && apt-get clean
-
 COPY dockerd_entrypoint.sh /usr/local/bin/dockerd_entrypoint.sh
 RUN chmod +x /usr/local/bin/dockerd_entrypoint.sh
 
@@ -220,6 +207,14 @@ RUN bash setup_oss_compliance.sh ${PYTHON} && rm setup_oss_compliance.sh
 ENTRYPOINT ["bash", "-m", "dockerd_entrypoint.sh"]
 CMD ["/bin/bash"]
 
+# This arg required to stop docker build waiting for region configuration while installing tz data from ubuntu 20
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update \
+ && apt-get upgrade -y \
+ && apt-get autoremove -y \
+ && apt-get clean \
+ && rm -rf /var/lib/apt/lists/*
 
 #################################################################
 #  ____                   __  __       _
@@ -247,18 +242,6 @@ ARG PYTHON
 
 WORKDIR /
 
-RUN apt-get update \
- && apt-get install -y --no-install-recommends \
-    curl \
-    wget \
-    git \
-    jq \
-    emacs \
-    vim \
-    unzip \
- && rm -rf /var/lib/apt/lists/* \
- && apt-get clean
-
 # Install SM packages
 RUN pip install --no-cache-dir -U \
     smclarify \
@@ -288,4 +271,13 @@ COPY start_with_right_hostname.sh /usr/local/bin/start_with_right_hostname.sh
 RUN chmod +x /usr/local/bin/start_with_right_hostname.sh
 
 ENTRYPOINT ["bash", "-m", "start_with_right_hostname.sh"]
-CMD ["/bin/bash"]
+CMD ["/bin/bash"]
+
+# This arg required to stop docker build waiting for region configuration while installing tz data from ubuntu 20
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update \
+ && apt-get upgrade -y \
+ && apt-get autoremove -y \
+ && apt-get clean \
+ && rm -rf /var/lib/apt/lists/*