SAI Vulnerability Fixes (aws#3594)

* SAI Vulnerabilities

* Undo Dev Config

Author: amzn-choeric

release_images_inference.yml

@@ -42,4 +42,18 @@ release_images:
       disable_sm_tag: True # [Default: False] Set to True to prevent SageMaker Abbreviated Tags from being attached
       # to images being published.
       force_release: False  # [Default: False] Set to True to force images to be published even if the same image
-      # has already been published. Re-released image will have minor version incremented by 1.
+      # has already been published. Re-released image will have minor version incremented by 1.
+  4:
+    framework: "stabilityai_pytorch"
+    version: "2.0.1"
+    arch_type: "x86"
+    customer_type: "sagemaker"
+    inference:
+      device_types: ["gpu"]
+      python_versions: ["py310"]
+      os_version: "ubuntu20.04"
+      cuda_version: "cu118"
+      sgm_version: "0.1.0"
+      example: False        
+      disable_sm_tag: False
+      force_release: False

stabilityai/pytorch/inference/buildspec.yml

@@ -23,15 +23,15 @@ images:
   BuildStabilityaiPytorchGpuPy310InferenceDockerImage:
     <<: *INFERENCE_REPOSITORY
     build: &STABILITYAI_PYTORCH_GPU_INFERENCE_PY3 false
-    image_size_baseline: 14000
+    image_size_baseline: 17000
     device_type: &DEVICE_TYPE gpu
     python_version: &DOCKER_PYTHON_VERSION py3
     tag_python_version: &TAG_PYTHON_VERSION py310
     cuda_version: &CUDA_VERSION cu118
     os_version: &OS_VERSION ubuntu20.04
     sgm_version: &SGM_VERSION 0.1.0
     tag: !join [ *VERSION, "-", 'sgm',*SGM_VERSION, '-', *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *CUDA_VERSION, "-", *OS_VERSION, "-sagemaker" ]
-    docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *CUDA_VERSION, /Dockerfile.,
+    docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *CUDA_VERSION, /Dockerfile., sagemaker.,
                          *DEVICE_TYPE ]
     context:
           <<: *INFERENCE_CONTEXT                         

stabilityai/pytorch/inference/docker/2.0/py3/cu118/Dockerfile.gpu renamed to stabilityai/pytorch/inference/docker/2.0/py3/cu118/Dockerfile.sagemaker.gpu

@@ -5,10 +5,12 @@ ARG PYTHON=python3
 ARG XFORMERS_VERSION=0.0.20
 
 # Resolve CVE
-RUN apt-get update
+RUN apt-get update \
+  && apt-mark hold libcudnn8 \
+  && apt-get upgrade -y \
+  && apt-get clean \
+  && rm -rf /var/lib/apt/lists/* 
 RUN pip --no-cache-dir install -U pip
-RUN apt-get install -y unattended-upgrades
-RUN unattended-upgrade
 
 # xformers must be installed from source due to the older version of python in the DLC
 RUN pip install ninja \
@@ -36,8 +38,10 @@ RUN pip install --no-cache-dir -U \
 # Resolve CVE
 RUN pip install --no-cache-dir -U \
     werkzeug \
-    urllib3 \
-    opencv-python
+    # Upper limit restriction by boto3
+    urllib3==2.0.7 \
+    opencv-python \
+    pyarrow
 
 # Configure Torchserve for large model loading
 ENV TS_DEFAULT_RESPONSE_TIMEOUT=1000

test/sagemaker_tests/pytorch/inference/integration/local/test_serving.py

@@ -158,11 +158,13 @@ def _predictor(
     )
 
     with local_mode_utils.lock():
+        predictor = None
         try:
             predictor = model.deploy(1, instance_type)
             yield predictor
         finally:
-            predictor.delete_endpoint()
+            if predictor:
+                predictor.delete_endpoint()
 
 
 def _assert_prediction_npy_json(predictor, test_loader, content_type, accept):

test/sagemaker_tests/pytorch/inference/requirements.txt

@@ -1,5 +1,7 @@
 boto3
 coverage
+# Docker v7.0.0 breaks compatibility with Docker Compose v1 (SageMaker Local)
+docker<=6.1.3
 docker-compose
 Flask==1.1.1
 fabric

test/test_utils/__init__.py

@@ -287,6 +287,8 @@ def get_dockerfile_path_for_image(image_uri, python_version=None):
         framework_path = framework.replace("_", os.path.sep)
     elif "habana" in image_uri:
         framework_path = os.path.join("habana", framework)
+    elif "stabilityai" in framework:
+        framework_path = framework.replace("_", os.path.sep)
     else:
         framework_path = framework
 
@@ -317,7 +319,6 @@ def get_dockerfile_path_for_image(image_uri, python_version=None):
     neuron_sdk_version = get_neuron_sdk_version_from_tag(image_uri)
 
     dockerfile_name = get_expected_dockerfile_filename(device_type, image_uri)
-
     dockerfiles_list = [
         path
         for path in glob(os.path.join(python_version_path, "**", dockerfile_name), recursive=True)