Fix PyTorch 2.8 SageMaker Package Regressions (#5520)

* run tests

* rebuild and run tests

* update sagemaker pin and rebuild image

* rebuild image with sm sdk

* pin versions and rebuild image

* Fix version pins and rebuild image

* pins s3fs to prevent the downgrade and rebuild

* resolve incompatibilty issue

* retry tests

* install newer version of s3fs

* revert config file

* revert version pin changes and pin sm version

* build image

* put back awscli and boto3 pins

* Empty commit to trigger rebuild

* fix typo & rerun

* fix SM version pin & rerun

* add sniffio explicitly & rerun

* revert toml file

---------

Co-authored-by: DevakiBolleneni <devakib@amazon.com>

Author: DevakiBolleneni

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

@@ -188,7 +188,7 @@ RUN pip install --no-cache-dir \
     "requests>=2.32.0" \
     "setuptools>=70.0.0" \
     "urllib3>=2.5.0" \
-    "awscli" \
+    awscli \
     opencv-python==4.11.0.86 \
     mpi4py \
     jinja2>=3.1.6 \
@@ -282,10 +282,12 @@ WORKDIR /
 # Install SM packages
 RUN pip install --no-cache-dir -U \
     smclarify \
-    "sagemaker>=2" \
+    "sagemaker==2.254.1" \
     sagemaker-experiments \
     sagemaker-pytorch-training \
-    sagemaker-training
+    sagemaker-training \
+    # Add sniffio explicitly as it's not included in sagemaker==2.254.1 dependencies
+    sniffio  
 
 # Install extra packages
 RUN pip install --no-cache-dir -U \
@@ -297,7 +299,7 @@ RUN pip install --no-cache-dir -U \
     scikit-learn \
     seaborn \
     shap \
-    cloudpickle
+    cloudpickle 
 
 # Copy workaround script for incorrect hostname
 COPY changehostname.c /

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

@@ -250,10 +250,12 @@ RUN pip install --no-cache-dir -U \
     "awscli<1.42.50" \
     "boto3<1.40.50" \
     smclarify \
-    "sagemaker>=2" \
+    "sagemaker==2.254.1" \
     sagemaker-experiments \
     sagemaker-pytorch-training \
-    sagemaker-training
+    sagemaker-training \
+    # Add sniffio explicitly as it's not included in sagemaker==2.254.1 dependencies
+    sniffio   
 
 # Install extra packages
 RUN pip install --no-cache-dir -U \
@@ -265,7 +267,7 @@ RUN pip install --no-cache-dir -U \
     shap \
     scikit-learn \
     seaborn \
-    cloudpickle
+    cloudpickle 
 
 COPY setup_oss_compliance.sh setup_oss_compliance.sh
 RUN bash setup_oss_compliance.sh ${PYTHON} && rm setup_oss_compliance.sh