Update Archiver to v0.3.1 / Integ tests

Author: dk19y

src/sagemaker_pytorch_serving_container/torchserve.py

@@ -42,7 +42,6 @@
 )
 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"
 
@@ -117,12 +116,8 @@ def _adapt_to_ts_format(handler_service):
         DEFAULT_TS_MODEL_NAME,
         "--handler",
         handler_service,
-        "--serialized-file",
-        os.path.join(environment.model_dir, DEFAULT_TS_MODEL_SERIALIZED_FILE),
         "--export-path",
         DEFAULT_TS_MODEL_DIRECTORY,
-        "--extra-files",
-        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,6 +1,8 @@
 ARG region
 FROM 763104351884.dkr.ecr.$region.amazonaws.com/pytorch-inference:1.6.0-cpu-py3
 
+RUN git clone https://github.com/pytorch/serve/ && cd serve/model-archiver/ && git checkout patch_release_0_3_1 && pip install . && cd -
+
 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 git clone https://github.com/pytorch/serve/ && cd serve/model-archiver/ && git checkout patch_release_0_3_1 && pip install . && cd -
+
 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

@@ -38,6 +38,8 @@ RUN useradd -m model-server \
 COPY artifacts/ts-entrypoint.py /usr/local/bin/dockerd-entrypoint.py
 COPY artifacts/config.properties /home/model-server
 
+RUN git clone https://github.com/pytorch/serve/ && cd serve/model-archiver/ && git checkout patch_release_0_3_1 && pip install . && cd -
+
 RUN chmod +x /usr/local/bin/dockerd-entrypoint.py
 
 EXPOSE 8080 8081

test/integration/__init__.py

@@ -28,33 +28,33 @@
 model_cpu_dir = os.path.join(mnist_path, cpu_sub_dir)
 mnist_cpu_script = os.path.join(model_cpu_dir, code_sub_dir, 'mnist.py')
 model_cpu_tar = file_utils.make_tarfile(mnist_cpu_script,
-                                        os.path.join(model_cpu_dir, "model.pth"),
+                                        os.path.join(model_cpu_dir, "torch_model.pth"),
                                         model_cpu_dir,
                                         script_path="code")
 
 model_cpu_1d_dir = os.path.join(model_cpu_dir, '1d')
 mnist_1d_script = os.path.join(model_cpu_1d_dir, code_sub_dir, 'mnist_1d.py')
 model_cpu_1d_tar = file_utils.make_tarfile(mnist_1d_script,
-                                           os.path.join(model_cpu_1d_dir, "model.pth"),
+                                           os.path.join(model_cpu_1d_dir, "torch_model.pth"),
                                            model_cpu_1d_dir,
                                            script_path="code")
 
 model_gpu_dir = os.path.join(mnist_path, gpu_sub_dir)
 mnist_gpu_script = os.path.join(model_gpu_dir, code_sub_dir, 'mnist.py')
 model_gpu_tar = file_utils.make_tarfile(mnist_gpu_script,
-                                        os.path.join(model_gpu_dir, "model.pth"),
+                                        os.path.join(model_gpu_dir, "torch_model.pth"),
                                         model_gpu_dir,
                                         script_path="code")
 
 model_eia_dir = os.path.join(mnist_path, eia_sub_dir)
 mnist_eia_script = os.path.join(model_eia_dir, 'mnist.py')
 model_eia_tar = file_utils.make_tarfile(mnist_eia_script,
-                                        os.path.join(model_eia_dir, "model.pth"),
+                                        os.path.join(model_eia_dir, "torch_model.pth"),
                                         model_eia_dir)
 
 call_model_fn_once_script = os.path.join(model_cpu_dir, code_sub_dir, 'call_model_fn_once.py')
 call_model_fn_once_tar = file_utils.make_tarfile(call_model_fn_once_script,
-                                                 os.path.join(model_cpu_dir, "model.pth"),
+                                                 os.path.join(model_cpu_dir, "torch_model.pth"),
                                                  model_cpu_dir,
                                                  "model_call_model_fn_once.tar.gz",
                                                  script_path="code")

test/resources/mnist/model_cpu/1d/code/mnist_1d.py

@@ -42,6 +42,6 @@ def forward(self, x):
 
 def model_fn(model_dir):
     model = torch.nn.DataParallel(Net())
-    with open(os.path.join(model_dir, 'model.pth'), 'rb') as f:
+    with open(os.path.join(model_dir, 'torch_model.pth'), 'rb') as f:
         model.load_state_dict(torch.load(f))
     return model

test/resources/mnist/model_cpu/1d/model.pth renamed to test/resources/mnist/model_cpu/1d/torch_model.pth

@@ -52,6 +52,6 @@ def forward(self, x):
 def model_fn(model_dir):
     logger.info('model_fn')
     model = torch.nn.DataParallel(Net())
-    with open(os.path.join(model_dir, 'model.pth'), 'rb') as f:
+    with open(os.path.join(model_dir, 'torch_model.pth'), 'rb') as f:
         model.load_state_dict(torch.load(f))
     return model

test/resources/mnist/model_cpu/code/mnist.py

@@ -38,11 +38,11 @@ def model_fn(model_dir):
     logger.info('model_fn: Loading model with TorchScript from {}'.format(model_dir))
     # Scripted model is serialized with torch.jit.save().
     # No need to instantiate model definition then load state_dict
-    model = torch.jit.load('model.pth')
+    model = torch.jit.load('torch_model.pth')
     return model
 
 
 def save_model(model, model_dir):
     logger.info("Saving the model to {}.".format(model_dir))
-    path = os.path.join(model_dir, 'model.pth')
+    path = os.path.join(model_dir, 'torch_model.pth')
     torch.jit.save(model, path)

test/resources/mnist/model_cpu/model.pth renamed to test/resources/mnist/model_cpu/torch_model.pth

@@ -52,6 +52,6 @@ def forward(self, x):
 def model_fn(model_dir):
     logger.info('model_fn')
     model = torch.nn.DataParallel(Net())
-    with open(os.path.join(model_dir, 'model.pth'), 'rb') as f:
+    with open(os.path.join(model_dir, 'torch_model.pth'), 'rb') as f:
         model.load_state_dict(torch.load(f))
     return model