Fix PyTorch SDK Failure

Author: dk19y

buildspec.yml

@@ -2,7 +2,7 @@ version: 0.2
 
 env:
   variables:
-    FRAMEWORK_VERSION: '1.5.0'
+    FRAMEWORK_VERSION: '1.6.0'
     EIA_FRAMEWORK_VERSION: '1.3.1'
     CPU_INSTANCE_TYPE: 'ml.c4.xlarge'
     GPU_INSTANCE_TYPE: 'ml.p2.8xlarge'

src/sagemaker_pytorch_serving_container/torchserve.py

@@ -43,6 +43,7 @@
 DEFAULT_TS_MODEL_DIRECTORY = os.path.join(os.getcwd(), ".sagemaker", "ts", "models")
 DEFAULT_TS_MODEL_NAME = "model"
 DEFAULT_TS_MODEL_SERIALIZED_FILE = "model.pth"
+DEFAULT_TS_CODE_DIR = "code"
 DEFAULT_HANDLER_SERVICE = "sagemaker_pytorch_serving_container.handler_service"
 
 ENABLE_MULTI_MODEL = os.getenv("SAGEMAKER_MULTI_MODEL", "false") == "true"
@@ -121,7 +122,7 @@ def _adapt_to_ts_format(handler_service):
         "--export-path",
         DEFAULT_TS_MODEL_DIRECTORY,
         "--extra-files",
-        os.path.join(environment.model_dir, environment.Environment().module_name + ".py"),
+        os.path.join(environment.model_dir, DEFAULT_TS_CODE_DIR, environment.Environment().module_name + ".py"),
         "--version",
         "1",
     ]

test/container/1.6.0/Dockerfile.dlc.cpu

@@ -1,23 +1,6 @@
 ARG region
-FROM 763104351884.dkr.ecr.$region.amazonaws.com/pytorch-inference:1.5.0-cpu-py3
-
-ARG TS_VERSION=0.1.1
-RUN apt-get update \
-    && apt-get install -y --no-install-recommends software-properties-common \
-    && add-apt-repository ppa:openjdk-r/ppa \
-    && apt-get update \
-    && apt-get install -y --no-install-recommends openjdk-11-jdk
-
-RUN pip install torchserve==$TS_VERSION \
-    && pip install torch-model-archiver==$TS_VERSION
-
-RUN pip uninstall torch \
-    && pip uninstall torchvision \
-    && pip install torch=1.6.0 \
-    && pip install torchvision=0.7.0
+FROM 763104351884.dkr.ecr.$region.amazonaws.com/pytorch-inference:1.6.0-cpu-py3
 
 COPY dist/sagemaker_pytorch_inference-*.tar.gz /sagemaker_pytorch_inference.tar.gz
 RUN pip install --upgrade --no-cache-dir /sagemaker_pytorch_inference.tar.gz && \
     rm /sagemaker_pytorch_inference.tar.gz
-
-CMD ["torchserve", "--start", "--ts-config", "/home/model-server/config.properties", "--model-store", "/home/model-server/"]

test/container/1.6.0/Dockerfile.dlc.gpu

@@ -1,23 +1,6 @@
 ARG region
-FROM 763104351884.dkr.ecr.$region.amazonaws.com/pytorch-inference:1.5.0-cpu-py3
-
-ARG TS_VERSION=0.1.1
-RUN apt-get update \
-    && apt-get install -y --no-install-recommends software-properties-common \
-    && add-apt-repository ppa:openjdk-r/ppa \
-    && apt-get update \
-    && apt-get install -y --no-install-recommends openjdk-11-jdk
-
-RUN pip install torchserve==$TS_VERSION \
-    && pip install torch-model-archiver==$TS_VERSION
-
-RUN pip uninstall torch \
-    && pip uninstall torchvision \
-    && pip install torch=1.6.0 \
-    && pip install torchvision=0.7.0
+FROM 763104351884.dkr.ecr.$region.amazonaws.com/pytorch-inference:1.6.0-gpu-py3
 
 COPY dist/sagemaker_pytorch_inference-*.tar.gz /sagemaker_pytorch_inference.tar.gz
 RUN pip install --upgrade --no-cache-dir /sagemaker_pytorch_inference.tar.gz && \
     rm /sagemaker_pytorch_inference.tar.gz
-
-CMD ["torchserve", "--start", "--ts-config", "/home/model-server/config.properties", "--model-store", "/home/model-server/"]