Fix PyTorch 2.8 SageMaker Package Regressions

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

@@ -161,6 +161,7 @@ ENV PATH="/usr/local/bin:${PATH}"
 RUN python -m pip install --upgrade pip --trusted-host pypi.org --trusted-host files.pythonhosted.org
 
 # Install common packages
+# Updated boto3 and awscli versions to match production baseline
 RUN pip install --no-cache-dir \
     cython \
     cryptography \
@@ -172,7 +173,7 @@ RUN pip install --no-cache-dir \
     typing \
     charset-normalizer \
     packaging \
-    boto3 \
+    "boto3>=1.41.0" \
     PyYAML \
     numpy \
     scipy \
@@ -188,7 +189,7 @@ RUN pip install --no-cache-dir \
     "requests>=2.32.0" \
     "setuptools>=70.0.0" \
     "urllib3>=2.5.0" \
-    "awscli" \
+    "awscli>=1.43.0" \
     opencv-python==4.11.0.86 \
     mpi4py \
     jinja2>=3.1.6 \
@@ -280,14 +281,20 @@ ENV SAGEMAKER_TRAINING_MODULE=sagemaker_pytorch_container.training:main
 WORKDIR /
 
 # Install SM packages
+# Version pins to prevent package regressions and fix protobuf CVE vulnerabilities
 RUN pip install --no-cache-dir -U \
     smclarify \
-    "sagemaker>=2" \
+    "sagemaker>=2.200.0" \
     sagemaker-experiments \
-    sagemaker-pytorch-training \
-    sagemaker-training
+    "sagemaker-pytorch-training>=3.0.0" \
+    "sagemaker-training>=5.1.1" \
+    "protobuf>=6.31.1" \
+    "rich>=14.2.0" \
+    "sniffio>=1.3.0" \
+    "wrapt>=2.0.1"
 
 # Install extra packages
+# Required packages missing from earlier builds to match production baseline
 RUN pip install --no-cache-dir -U \
     bokeh \
     imageio \
@@ -297,7 +304,12 @@ RUN pip install --no-cache-dir -U \
     scikit-learn \
     seaborn \
     shap \
-    cloudpickle
+    cloudpickle \
+    google-pasta \
+    multiprocess \
+    pathos \
+    pox \
+    ppft
 
 # Copy workaround script for incorrect hostname
 COPY changehostname.c /

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

@@ -247,15 +247,23 @@ WORKDIR /
 # Install SM packages
 RUN pip install --no-cache-dir -U \
     # address package regression caused by smclarify depedency s3fs"
-    "awscli<1.42.50" \
-    "boto3<1.40.50" \
+    # Pin s3fs before smclarify to prevent downgrade to 0.4.2 which lacks async dependencies
+    # Version pins to prevent package regressions and fix protobuf CVE vulnerabilities
+    "s3fs>=2025.10.0" \
     smclarify \
-    "sagemaker>=2" \
+    "sagemaker>=2.200.0" \
     sagemaker-experiments \
-    sagemaker-pytorch-training \
-    sagemaker-training
+    "sagemaker-pytorch-training>=3.0.0" \
+    "sagemaker-training>=5.1.1" \
+    "protobuf>=6.31.1" \
+    "rich>=14.2.0" \
+    "sniffio>=1.3.0" 
+
+# Install AWS CLI and boto3 separately to get compatible versions
+RUN pip install --no-cache-dir -U awscli boto3
 
 # Install extra packages
+# Required packages missing from earlier builds to match production baseline
 RUN pip install --no-cache-dir -U \
     bokeh \
     imageio \
@@ -265,7 +273,12 @@ RUN pip install --no-cache-dir -U \
     shap \
     scikit-learn \
     seaborn \
-    cloudpickle
+    cloudpickle \
+    google-pasta \
+    multiprocess \
+    pathos \
+    pox \
+    ppft
 
 COPY setup_oss_compliance.sh setup_oss_compliance.sh
 RUN bash setup_oss_compliance.sh ${PYTHON} && rm setup_oss_compliance.sh