Update Dockerfile for Model Archiver 0.3.1

Author: dk19y

test/container/1.6.0/Dockerfile.dlc.cpu

@@ -1,6 +1,8 @@
 ARG region
 FROM 763104351884.dkr.ecr.$region.amazonaws.com/pytorch-inference:1.6.0-cpu-py3
 
+RUN pip install --upgrade torch-model-archiver==0.3.1
+
 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

test/container/1.6.0/Dockerfile.dlc.gpu

@@ -1,6 +1,8 @@
 ARG region
 FROM 763104351884.dkr.ecr.$region.amazonaws.com/pytorch-inference:1.6.0-gpu-py3
 
+RUN pip install --upgrade torch-model-archiver==0.3.1
+
 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

test/container/1.6.0/Dockerfile.pytorch

@@ -3,7 +3,8 @@ FROM pytorch/pytorch:1.6.0-cuda10.1-cudnn7-runtime
 LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
 LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
 
-ARG TS_VERSION=0.1.1
+ARG TS_VERSION=0.3.1
+ARG TS_ARCHIVER_VERSION=0.3.1
 
 ENV SAGEMAKER_SERVING_MODULE sagemaker_pytorch_serving_container.serving:main
 ENV TEMP=/home/model-server/tmp
@@ -25,7 +26,7 @@ RUN conda install -c conda-forge opencv==4.0.1 \
     && ln -s /opt/conda/bin/pip /usr/local/bin/pip3
 
 RUN pip install torchserve==$TS_VERSION \
- && pip install torch-model-archiver==$TS_VERSION
+ && pip install torch-model-archiver==$TS_ARCHIVER_VERSION
 
 COPY dist/sagemaker_pytorch_inference-*.tar.gz /sagemaker_pytorch_inference.tar.gz
 RUN pip install --no-cache-dir /sagemaker_pytorch_inference.tar.gz && \