[PATCH] PyTorch Training Images CVE Patch (#5011)

Yadan-Wei · Yadan Wei · web-flow · commit b7c041b74783 · 2025-07-16T10:54:17.000-07:00
* Patch nvjpeg for PT TR

---------

Co-authored-by: Yadan Wei &lt;yadanwei@amazon.com&gt;
diff --git a/pytorch/training/buildspec-2-5-ec2.yml b/pytorch/training/buildspec-2-5-ec2.yml
@@ -5,7 +5,7 @@ framework: &FRAMEWORK pytorch
 version: &VERSION 2.5.1
 short_version: &SHORT_VERSION "2.5"
 arch_type: x86
-autopatch_build: "True"
+# autopatch_build: "True"
 
 repository_info:
   training_repository: &TRAINING_REPOSITORY
@@ -21,6 +21,9 @@ context:
     changehostname:
       source: docker/build_artifacts/changehostname.c
       target: changehostname.c
+    start_cuda_compat:
+      source: docker/build_artifacts/start_cuda_compat.sh
+      target: start_cuda_compat.sh
     start_with_right_hostname:
       source: docker/build_artifacts/start_with_right_hostname.sh
       target: start_with_right_hostname.sh
diff --git a/pytorch/training/buildspec-2-5-sm.yml b/pytorch/training/buildspec-2-5-sm.yml
@@ -5,7 +5,7 @@ framework: &FRAMEWORK pytorch
 version: &VERSION 2.5.1
 short_version: &SHORT_VERSION "2.5"
 arch_type: x86
-autopatch_build: "True"
+# autopatch_build: "True"
 
 repository_info:
   training_repository: &TRAINING_REPOSITORY
diff --git a/pytorch/training/buildspec-2-6-ec2.yml b/pytorch/training/buildspec-2-6-ec2.yml
@@ -5,7 +5,7 @@ framework: &FRAMEWORK pytorch
 version: &VERSION 2.6.0
 short_version: &SHORT_VERSION "2.6"
 arch_type: x86
-autopatch_build: "True"
+# autopatch_build: "True"
 
 repository_info:
   training_repository: &TRAINING_REPOSITORY
diff --git a/pytorch/training/buildspec-2-6-sm.yml b/pytorch/training/buildspec-2-6-sm.yml
@@ -5,7 +5,7 @@ framework: &FRAMEWORK pytorch
 version: &VERSION 2.6.0
 short_version: &SHORT_VERSION "2.6"
 arch_type: x86
-autopatch_build: "True"
+# autopatch_build: "True"
 
 repository_info:
   training_repository: &TRAINING_REPOSITORY
diff --git a/pytorch/training/buildspec-2-7-ec2.yml b/pytorch/training/buildspec-2-7-ec2.yml
@@ -5,7 +5,7 @@ framework: &FRAMEWORK pytorch
 version: &VERSION 2.7.1
 short_version: &SHORT_VERSION "2.7"
 arch_type: x86
-autopatch_build: "True"
+# autopatch_build: "True"
 
 repository_info:
   training_repository: &TRAINING_REPOSITORY
diff --git a/pytorch/training/buildspec-2-7-sm.yml b/pytorch/training/buildspec-2-7-sm.yml
@@ -5,7 +5,7 @@ framework: &FRAMEWORK pytorch
 version: &VERSION 2.7.1
 short_version: &SHORT_VERSION "2.7"
 arch_type: x86
-autopatch_build: "True"
+# autopatch_build: "True"
 
 repository_info:
   training_repository: &TRAINING_REPOSITORY
diff --git a/pytorch/training/docker/2.5/py3/Dockerfile.cpu b/pytorch/training/docker/2.5/py3/Dockerfile.cpu
@@ -166,15 +166,18 @@ RUN /opt/conda/bin/mamba install -y -c conda-forge \
     fsspec \
     "idna>=3.7" \
     "tqdm>=4.66.3" \
-    "requests>=2.32.0" \
-    "setuptools>=70.0.0" \
-    "urllib3<2" \
+    "requests>=2.32.4" \
+    "setuptools>=80.9.0" \
+    "urllib3>=2.5.0" \
     "awscli<2" \
  && /opt/conda/bin/mamba clean -afy \
  && rm -rf /etc/apt/sources.list.d/*
 
 # Install common pip packages (in case of conda package is not available)
-RUN pip install --no-cache-dir opencv-python mpi4py
+RUN pip install --no-cache-dir "opencv-python==4.11.0.86" mpi4py
+
+#address pip cve
+RUN pip install --no-cache-dir --upgrade "pip>=25.1.1"
 
 RUN curl -o /license.txt https://aws-dlc-licenses.s3.amazonaws.com/pytorch-2.5/license.txt
 
@@ -298,7 +301,7 @@ RUN pip install --no-cache-dir -U \
     "sagemaker>=2,<3" \
     "sagemaker-experiments<1" \
     sagemaker-pytorch-training \
-    sagemaker-training
+    "sagemaker-training>=5.0.0"
 
 # Install extra packages
 RUN /opt/conda/bin/mamba install -y -c conda-forge \
diff --git a/pytorch/training/docker/2.5/py3/Dockerfile.ec2.cpu.core_packages.json b/pytorch/training/docker/2.5/py3/Dockerfile.ec2.cpu.core_packages.json
@@ -1,14 +1,14 @@
 {
   "accelerate": {
-    "version_specifier": "==1.1.1",
+    "version_specifier": "==1.8.1",
     "skip": "True"
   },
   "fastai": {
-    "version_specifier": "==2.7.18",
+    "version_specifier": "==2.8.2",
     "skip": "True"
   },
   "s3torchconnector": {
-    "version_specifier": "==1.2.6",
+    "version_specifier": "==1.4.2",
     "skip": "True"
   },
   "torch": {
@@ -43,7 +43,7 @@
     "version_specifier": ">=70.0.0"
   },
   "urllib3": {
-    "version_specifier": "<2"
+    "version_specifier": ">=2.5.0"
   },
   "awscli": {
     "version_specifier": "<2"
diff --git a/pytorch/training/docker/2.5/py3/Dockerfile.sagemaker.cpu.core_packages.json b/pytorch/training/docker/2.5/py3/Dockerfile.sagemaker.cpu.core_packages.json
@@ -1,14 +1,14 @@
 {
   "accelerate": {
-    "version_specifier": "==1.5.1",
+    "version_specifier": "==1.8.1",
     "skip": "True"
   },
   "fastai": {
-    "version_specifier": "==2.7.19",
+    "version_specifier": "==2.8.2",
     "skip": "True"
   },
   "s3torchconnector": {
-    "version_specifier": "==1.3.2",
+    "version_specifier": "==1.4.2",
     "skip": "True"
   },
   "torch": {
@@ -43,7 +43,7 @@
     "version_specifier": ">=70.0.0"
   },
   "urllib3": {
-    "version_specifier": "<2"
+    "version_specifier": ">=2.5.0"
   },
   "awscli": {
     "version_specifier": "<2"
diff --git a/pytorch/training/docker/2.5/py3/cu124/Dockerfile.ec2.gpu.core_packages.json b/pytorch/training/docker/2.5/py3/cu124/Dockerfile.ec2.gpu.core_packages.json
@@ -1,10 +1,10 @@
 {
   "accelerate": {
-    "version_specifier": "==1.1.1",
+    "version_specifier": "==1.8.1",
     "skip": "True"
   },
   "fastai": {
-    "version_specifier": "==2.7.18",
+    "version_specifier": "==2.8.2",
     "skip": "True"
   },
   "flash-attn": {
@@ -16,7 +16,7 @@
     "skip": "True"
   },
   "s3torchconnector": {
-    "version_specifier": "==1.2.6",
+    "version_specifier": "==1.4.2",
     "skip": "True"
   },
   "torch": {
@@ -51,7 +51,7 @@
     "version_specifier": ">=70.0.0"
   },
   "urllib3": {
-    "version_specifier": "<2"
+    "version_specifier": ">=2.5.0"
   },
   "awscli": {
     "version_specifier": "<2"
diff --git a/pytorch/training/docker/2.5/py3/cu124/Dockerfile.gpu b/pytorch/training/docker/2.5/py3/cu124/Dockerfile.gpu
@@ -128,6 +128,20 @@ RUN apt-get update \
  && rm -rf /var/lib/apt/lists/* \
  && apt-get clean
 
+# patch nvjpeg
+RUN mkdir -p /tmp/nvjpeg \
+&& cd /tmp/nvjpeg \
+&& wget https://developer.download.nvidia.com/compute/cuda/redist/libnvjpeg/linux-x86_64/libnvjpeg-linux-x86_64-12.4.0.76-archive.tar.xz \
+&& tar -xvf libnvjpeg-linux-x86_64-12.4.0.76-archive.tar.xz \
+&& rm -rf /usr/local/cuda/targets/x86_64-linux/lib/libnvjpeg* \
+&& rm -rf /usr/local/cuda/targets/x86_64-linux/include/nvjpeg.h \
+&& cp libnvjpeg-linux-x86_64-12.4.0.76-archive/lib/libnvjpeg* /usr/local/cuda/targets/x86_64-linux/lib/ \
+&& cp libnvjpeg-linux-x86_64-12.4.0.76-archive/include/* /usr/local/cuda/targets/x86_64-linux/include/ \
+&& rm -rf /tmp/nvjpeg \
+# patch cuobdump and nvdisasm
+&& rm -rf /usr/local/cuda/bin/cuobjdump* \
+&& rm -rf /usr/local/cuda/bin/nvdisasm*
+
 # For EFA, below flags are needed to install EFA on docker image
 #  -n, --no-verify       Skip EFA device verification and test
 #  -l, --skip-limit-conf Skip EFA limit configuration
@@ -246,17 +260,20 @@ RUN /opt/conda/bin/mamba install -y -c conda-forge \
     fsspec \
     "idna>=3.7" \
     "tqdm>=4.66.3" \
-    "requests>=2.32.0" \
-    "setuptools>=70.0.0" \
-    "urllib3<2" \
+    "requests>=2.32.4" \
+    "setuptools>=80.9.0" \
+    "urllib3>=2.5.0" \
     "awscli<2" \
     libgcc \
     ninja \
  && /opt/conda/bin/mamba clean -afy \
  && rm -rf /etc/apt/sources.list.d/*
 
 # Install common pip packages (in case of conda package is not available)
-RUN pip install --no-cache-dir opencv-python mpi4py
+RUN pip install --no-cache-dir "opencv-python==4.11.0.86" mpi4py
+
+#address pip cve
+RUN pip install --no-cache-dir --upgrade "pip>=25.1.1"
 
 RUN curl -o /license.txt https://aws-dlc-licenses.s3.amazonaws.com/pytorch-2.5/license.txt
 
@@ -457,7 +474,7 @@ RUN pip install --no-cache-dir -U \
     "sagemaker>=2,<3" \
     "sagemaker-experiments<1" \
     sagemaker-pytorch-training \
-    sagemaker-training
+    "sagemaker-training>=5.0.0"
 
 # Install extra packages
 RUN /opt/conda/bin/mamba install -y -c conda-forge \
diff --git a/pytorch/training/docker/2.5/py3/cu124/Dockerfile.sagemaker.gpu.core_packages.json b/pytorch/training/docker/2.5/py3/cu124/Dockerfile.sagemaker.gpu.core_packages.json
@@ -1,10 +1,10 @@
 {
   "accelerate": {
-    "version_specifier": "==1.5.1",
+    "version_specifier": "==1.8.1",
     "skip": "True"
   },
   "fastai": {
-    "version_specifier": "==2.7.19",
+    "version_specifier": "==2.8.2",
     "skip": "True"
   },
   "flash-attn": {
@@ -16,7 +16,7 @@
     "skip": "True"
   },
   "s3torchconnector": {
-    "version_specifier": "==1.3.2",
+    "version_specifier": "==1.4.2",
     "skip": "True"
   },
   "torch": {
@@ -51,7 +51,7 @@
     "version_specifier": ">=70.0.0"
   },
   "urllib3": {
-    "version_specifier": "<2"
+    "version_specifier": ">=2.5.0"
   },
   "awscli": {
     "version_specifier": "<2"
diff --git a/pytorch/training/docker/2.6/py3/Dockerfile.cpu b/pytorch/training/docker/2.6/py3/Dockerfile.cpu
@@ -185,9 +185,9 @@ RUN pip install --no-cache-dir \
     "tqdm>=4.66.3" \
     "requests>=2.32.0" \
     "setuptools>=70.0.0" \
-    "urllib3<2" \
+    "urllib3>=2.5.0" \
     "awscli<2" \
-    opencv-python \
+    opencv-python==4.11.0.86 \
     mpi4py
 
 RUN curl -o /license.txt https://aws-dlc-licenses.s3.amazonaws.com/pytorch-2.6/license.txt
@@ -334,7 +334,8 @@ RUN pip install --no-cache-dir -U \
     "sagemaker>=2,<3" \
     "sagemaker-experiments<1" \
     sagemaker-pytorch-training \
-    "sagemaker-training==4.9.0"
+    # sagemaker-training is upgraded to 5.0.0 by autopatch
+    "sagemaker-training>=5.0.0"
 
 # Install extra packages
 RUN pip install --no-cache-dir -U \
diff --git a/pytorch/training/docker/2.6/py3/Dockerfile.ec2.cpu.core_packages.json b/pytorch/training/docker/2.6/py3/Dockerfile.ec2.cpu.core_packages.json
@@ -4,7 +4,7 @@
     "skip": "True"
   },
   "s3torchconnector": {
-    "version_specifier": "==1.4.1",
+    "version_specifier": "==1.4.2",
     "skip": "True"
   },
   "torch": {
diff --git a/pytorch/training/docker/2.6/py3/Dockerfile.sagemaker.cpu.core_packages.json b/pytorch/training/docker/2.6/py3/Dockerfile.sagemaker.cpu.core_packages.json
@@ -4,7 +4,7 @@
     "skip": "True"
   },
   "s3torchconnector": {
-    "version_specifier": "==1.4.1",
+    "version_specifier": "==1.4.2",
     "skip": "True"
   },
   "torch": {
diff --git a/pytorch/training/docker/2.6/py3/cu126/Dockerfile.ec2.gpu.core_packages.json b/pytorch/training/docker/2.6/py3/cu126/Dockerfile.ec2.gpu.core_packages.json
@@ -12,7 +12,7 @@
     "skip": "True"
   },
   "s3torchconnector": {
-    "version_specifier": "==1.4.1",
+    "version_specifier": "==1.4.2",
     "skip": "True"
   },
   "torch": {
diff --git a/pytorch/training/docker/2.6/py3/cu126/Dockerfile.gpu b/pytorch/training/docker/2.6/py3/cu126/Dockerfile.gpu
@@ -131,6 +131,19 @@ RUN apt-get update \
  && rm -rf /var/lib/apt/lists/* \
  && apt-get clean
 
+RUN mkdir -p /tmp/nvjpeg \
+&& cd /tmp/nvjpeg \
+&& wget https://developer.download.nvidia.com/compute/cuda/redist/libnvjpeg/linux-x86_64/libnvjpeg-linux-x86_64-12.4.0.76-archive.tar.xz \
+&& tar -xvf libnvjpeg-linux-x86_64-12.4.0.76-archive.tar.xz \
+&& rm -rf /usr/local/cuda/targets/x86_64-linux/lib/libnvjpeg* \
+&& rm -rf /usr/local/cuda/targets/x86_64-linux/include/nvjpeg.h \
+&& cp libnvjpeg-linux-x86_64-12.4.0.76-archive/lib/libnvjpeg* /usr/local/cuda/targets/x86_64-linux/lib/ \
+&& cp libnvjpeg-linux-x86_64-12.4.0.76-archive/include/* /usr/local/cuda/targets/x86_64-linux/include/ \
+&& rm -rf /tmp/nvjpeg \
+# patch cuobjdump and nvdisasm
+&& rm -rf /usr/local/cuda/bin/cuobjdump* \
+&& rm -rf /usr/local/cuda/bin/nvdisasm* 
+
 # For EFA, below flags are needed to install EFA on docker image
 #  -n, --no-verify       Skip EFA device verification and test
 #  -l, --skip-limit-conf Skip EFA limit configuration
@@ -230,10 +243,10 @@ RUN pip install --no-cache-dir \
     "tqdm>=4.66.3" \
     "requests>=2.32.0" \
     "setuptools>=70.0.0" \
-    "urllib3<2" \
+    "urllib3>=2.5.0" \
     "awscli<2" \
     ninja \
-    opencv-python \
+    opencv-python==4.11.0.86 \
     mpi4py
 
 RUN curl -o /license.txt https://aws-dlc-licenses.s3.amazonaws.com/pytorch-2.6/license.txt
@@ -448,7 +461,8 @@ RUN pip install --no-cache-dir -U \
     "sagemaker>=2,<3" \
     "sagemaker-experiments<1" \
     sagemaker-pytorch-training \
-    "sagemaker-training==4.9.0"
+    # sagemaker-training is upgraded to 5.0.0 by autopatch
+    "sagemaker-training>=5.0.0"
 
 # Install extra packages
 RUN pip install --no-cache-dir -U \
@@ -460,7 +474,6 @@ RUN pip install --no-cache-dir -U \
     shap \
     scikit-learn \
     seaborn \
-    # pinned for sagemaker==2.233.0
     cloudpickle 
 
 # Copy workaround script for incorrect hostname
diff --git a/pytorch/training/docker/2.6/py3/cu126/Dockerfile.sagemaker.gpu.core_packages.json b/pytorch/training/docker/2.6/py3/cu126/Dockerfile.sagemaker.gpu.core_packages.json
@@ -12,7 +12,7 @@
     "skip": "True"
   },
   "s3torchconnector": {
-    "version_specifier": "==1.4.1",
+    "version_specifier": "==1.4.2",
     "skip": "True"
   },
   "torch": {
diff --git a/pytorch/training/docker/2.7/py3/Dockerfile.cpu b/pytorch/training/docker/2.7/py3/Dockerfile.cpu
@@ -183,9 +183,9 @@ RUN pip install --no-cache-dir \
     "tqdm>=4.66.3" \
     "requests>=2.32.0" \
     "setuptools>=70.0.0" \
-    "urllib3<2" \
+    "urllib3>=2.5.0" \
     "awscli<2" \
-    opencv-python \
+    "opencv-python==4.11.0.86" \
     mpi4py \
     jinja2>=3.1.6 \
     tornado>=6.5.1
diff --git a/pytorch/training/docker/2.7/py3/Dockerfile.ec2.cpu.core_packages.json b/pytorch/training/docker/2.7/py3/Dockerfile.ec2.cpu.core_packages.json
diff --git a/pytorch/training/docker/2.7/py3/Dockerfile.sagemaker.cpu.core_packages.json b/pytorch/training/docker/2.7/py3/Dockerfile.sagemaker.cpu.core_packages.json
diff --git a/pytorch/training/docker/2.7/py3/cu128/Dockerfile.ec2.gpu.core_packages.json b/pytorch/training/docker/2.7/py3/cu128/Dockerfile.ec2.gpu.core_packages.json
diff --git a/pytorch/training/docker/2.7/py3/cu128/Dockerfile.gpu b/pytorch/training/docker/2.7/py3/cu128/Dockerfile.gpu
diff --git a/pytorch/training/docker/2.7/py3/cu128/Dockerfile.sagemaker.gpu.core_packages.json b/pytorch/training/docker/2.7/py3/cu128/Dockerfile.sagemaker.gpu.core_packages.json
