diff --git a/setup.py b/setup.py index ce6d778f..f37eb73c 100644 --- a/setup.py +++ b/setup.py @@ -54,7 +54,7 @@ def read(fname): # different packages for different variants install_requires=['numpy', 'retrying', 'sagemaker-inference>=1.3.1'], extras_require={ - 'test': ['boto3==1.10.32', 'coverage==4.5.3', 'docker-compose==1.23.2', 'flake8==3.7.7', 'Flask==1.1.1', + 'test': ['boto3>=1.10.44', 'coverage==4.5.3', 'docker-compose==1.23.2', 'flake8==3.7.7', 'Flask==1.1.1', 'mock==2.0.0', 'pytest==4.4.0', 'pytest-cov==2.7.1', 'pytest-xdist==1.28.0', 'PyYAML==3.10', 'sagemaker==1.56.3', 'sagemaker-containers>=2.5.4', 'six==1.12.0', 'requests==2.20.0', 'requests_mock==1.6.0', 'torch==1.6.0', 'torchvision==0.7.0', 'tox==3.7.0'] diff --git a/src/sagemaker_pytorch_serving_container/torchserve.py b/src/sagemaker_pytorch_serving_container/torchserve.py index 95362352..d90e0000 100644 --- a/src/sagemaker_pytorch_serving_container/torchserve.py +++ b/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,14 +116,12 @@ 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", + "--extra-files", + os.path.join(environment.model_dir) ] logger.info(model_archiver_cmd) diff --git a/test/container/1.6.0/Dockerfile.dlc.cpu b/test/container/1.6.0/Dockerfile.dlc.cpu index 44667c02..152a5753 100644 --- a/test/container/1.6.0/Dockerfile.dlc.cpu +++ b/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 diff --git a/test/container/1.6.0/Dockerfile.dlc.gpu b/test/container/1.6.0/Dockerfile.dlc.gpu index e48fc985..c10fc836 100644 --- a/test/container/1.6.0/Dockerfile.dlc.gpu +++ b/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 diff --git a/test/container/1.6.0/Dockerfile.pytorch b/test/container/1.6.0/Dockerfile.pytorch index debf50ff..ca9042f2 100644 --- a/test/container/1.6.0/Dockerfile.pytorch +++ b/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 && \ diff --git a/test/integration/__init__.py b/test/integration/__init__.py index 0f10294d..0442aa25 100644 --- a/test/integration/__init__.py +++ b/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") diff --git a/test/resources/mnist/model_cpu/1d/code/mnist_1d.py b/test/resources/mnist/model_cpu/1d/code/mnist_1d.py index f027e700..457141e9 100644 --- a/test/resources/mnist/model_cpu/1d/code/mnist_1d.py +++ b/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 diff --git a/test/resources/mnist/model_cpu/1d/model.pth b/test/resources/mnist/model_cpu/1d/torch_model.pth similarity index 100% rename from test/resources/mnist/model_cpu/1d/model.pth rename to test/resources/mnist/model_cpu/1d/torch_model.pth diff --git a/test/resources/mnist/model_cpu/code/mnist.py b/test/resources/mnist/model_cpu/code/mnist.py index e0a1dfbe..3c773c90 100644 --- a/test/resources/mnist/model_cpu/code/mnist.py +++ b/test/resources/mnist/model_cpu/code/mnist.py @@ -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 diff --git a/test/resources/mnist/model_cpu/model.pth b/test/resources/mnist/model_cpu/torch_model.pth similarity index 100% rename from test/resources/mnist/model_cpu/model.pth rename to test/resources/mnist/model_cpu/torch_model.pth diff --git a/test/resources/mnist/model_eia/mnist.py b/test/resources/mnist/model_eia/mnist.py index ebc0bff0..53f0be8d 100644 --- a/test/resources/mnist/model_eia/mnist.py +++ b/test/resources/mnist/model_eia/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) diff --git a/test/resources/mnist/model_eia/model.pth b/test/resources/mnist/model_eia/torch_model.pth similarity index 100% rename from test/resources/mnist/model_eia/model.pth rename to test/resources/mnist/model_eia/torch_model.pth diff --git a/test/resources/mnist/model_gpu/code/mnist.py b/test/resources/mnist/model_gpu/code/mnist.py index e0a1dfbe..3c773c90 100644 --- a/test/resources/mnist/model_gpu/code/mnist.py +++ b/test/resources/mnist/model_gpu/code/mnist.py @@ -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 diff --git a/test/resources/mnist/model_gpu/model.pth b/test/resources/mnist/model_gpu/torch_model.pth similarity index 100% rename from test/resources/mnist/model_gpu/model.pth rename to test/resources/mnist/model_gpu/torch_model.pth diff --git a/test/unit/test_model_server.py b/test/unit/test_model_server.py index 108cd3c3..aeaec28e 100644 --- a/test/unit/test_model_server.py +++ b/test/unit/test_model_server.py @@ -145,16 +145,12 @@ def test_adapt_to_ts_format(path_exists, make_dir, subprocess_check_call, set_py torchserve.DEFAULT_TS_MODEL_NAME, "--handler", handler_service, - "--serialized-file", - os.path.join(environment.model_dir, torchserve.DEFAULT_TS_MODEL_SERIALIZED_FILE), "--export-path", torchserve.DEFAULT_TS_MODEL_DIRECTORY, - "--extra-files", - os.path.join(environment.model_dir, - torchserve.DEFAULT_TS_CODE_DIR, - environment.Environment().module_name + ".py"), "--version", "1", + "--extra-files", + environment.model_dir ] subprocess_check_call.assert_called_once_with(model_archiver_cmd)