diff --git a/pytorch/training/docker/2.8/py3/Dockerfile.cpu b/pytorch/training/docker/2.8/py3/Dockerfile.cpu
index 9e305de1b54a..c97508902096 100644
--- a/pytorch/training/docker/2.8/py3/Dockerfile.cpu
+++ b/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 /
diff --git a/pytorch/training/docker/2.8/py3/cu129/Dockerfile.gpu b/pytorch/training/docker/2.8/py3/cu129/Dockerfile.gpu
index 0b90890db2fb..5e24300d78d3 100644
--- a/pytorch/training/docker/2.8/py3/cu129/Dockerfile.gpu
+++ b/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