Add tests

Author: saimidu

buildspec.yml

@@ -3,7 +3,7 @@ version: 0.2
 env:
   variables:
     FRAMEWORK_VERSION: '1.6.0'
-    EIA_FRAMEWORK_VERSION: '1.3.1'
+    EIA_FRAMEWORK_VERSION: '1.5.1'
     CPU_INSTANCE_TYPE: 'ml.c4.xlarge'
     GPU_INSTANCE_TYPE: 'ml.p2.8xlarge'
     EIA_ACCELERATOR_TYPE: 'ml.eia2.medium'

test/integration/__init__.py

@@ -24,6 +24,8 @@
 gpu_sub_dir = 'model_gpu'
 eia_sub_dir = 'model_eia'
 code_sub_dir = 'code'
+default_sub_dir = 'default_model'
+default_sub_eia_dir = 'default_model_eia'
 
 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')
@@ -59,6 +61,18 @@
                                                  "model_call_model_fn_once.tar.gz",
                                                  script_path="code")
 
+default_model_dir = os.path.join(mnist_path, default_sub_dir)
+default_model_script = os.path.join(default_model_dir, "mnist.py")
+default_model_tar = file_utils.make_tarfile(
+    default_model_script, os.path.join(default_model_dir, "model.pt"), default_model_dir, script_path="code"
+)
+
+default_model_eia_dir = os.path.join(mnist_path, default_sub_dir)
+default_model_eia_script = os.path.join(default_model_eia_dir, "mnist.py")
+default_model_eia_tar = file_utils.make_tarfile(
+    default_model_eia_script, os.path.join(default_model_eia_dir, "model.pt"), default_model_eia_dir
+)
+
 ROLE = 'dummy/unused-role'
 DEFAULT_TIMEOUT = 20
 PYTHON3 = 'py3'

test/integration/sagemaker/test_mnist_default_inference.py

@@ -0,0 +1,101 @@
+# Copyright 2019-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import numpy as np
+import pytest
+import sagemaker
+from sagemaker.pytorch import PyTorchModel
+
+from integration import (
+    model_cpu_tar,
+    model_gpu_tar,
+    mnist_cpu_script,
+    mnist_gpu_script,
+    model_eia_tar,
+    mnist_eia_script,
+    default_model_script,
+    default_model_tar,
+    default_model_eia_script,
+    default_model_eia_tar,
+)
+from integration.sagemaker.timeout import timeout_and_delete_endpoint
+
+
+@pytest.mark.cpu_test
+def test_default_inference_cpu(sagemaker_session, image_uri, instance_type):
+    instance_type = instance_type or "ml.c4.xlarge"
+    # Scripted model is serialized with torch.jit.save().
+    # Default inference test doesn't need to instantiate model definition
+    _test_default_inference(sagemaker_session, image_uri, instance_type, default_model_tar, default_model_script)
+
+
+@pytest.mark.gpu_test
+def test_default_inference_gpu(sagemaker_session, image_uri, instance_type):
+    instance_type = instance_type or "ml.p2.xlarge"
+    # Scripted model is serialized with torch.jit.save().
+    # Default inference test doesn't need to instantiate model definition
+    _test_default_inference(sagemaker_session, image_uri, instance_type, default_model_tar, default_model_script)
+
+
+@pytest.mark.eia_test
+def test_default_inference_eia(sagemaker_session, image_uri, instance_type, accelerator_type):
+    instance_type = instance_type or "ml.c4.xlarge"
+    # Scripted model is serialized with torch.jit.save().
+    # Default inference test doesn't need to instantiate model definition
+    _test_default_inference(
+        sagemaker_session,
+        image_uri,
+        instance_type,
+        default_model_eia_tar,
+        default_model_eia_script,
+        accelerator_type=accelerator_type,
+    )
+
+
+def _test_default_inference(
+    sagemaker_session, image_uri, instance_type, model_tar, mnist_script, accelerator_type=None
+):
+    endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-pytorch-serving")
+
+    model_data = sagemaker_session.upload_data(
+        path=model_tar,
+        key_prefix="sagemaker-pytorch-serving/models",
+    )
+
+    pytorch = PyTorchModel(
+        model_data=model_data,
+        role="SageMakerRole",
+        entry_point=mnist_script,
+        image=image_uri,
+        sagemaker_session=sagemaker_session,
+    )
+    with timeout_and_delete_endpoint(endpoint_name, sagemaker_session, minutes=30):
+        # Use accelerator type to differentiate EI vs. CPU and GPU. Don't use processor value
+        if accelerator_type is not None:
+            predictor = pytorch.deploy(
+                initial_instance_count=1,
+                instance_type=instance_type,
+                accelerator_type=accelerator_type,
+                endpoint_name=endpoint_name,
+            )
+        else:
+            predictor = pytorch.deploy(
+                initial_instance_count=1, instance_type=instance_type, endpoint_name=endpoint_name
+            )
+
+        batch_size = 100
+        data = np.random.rand(batch_size, 1, 28, 28).astype(np.float32)
+        output = predictor.predict(data)
+
+        assert output.shape == (batch_size, 10)

test/resources/mnist/default_model/code/mnist.py

@@ -0,0 +1,35 @@
+# Copyright 2019-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+import logging
+import os
+import sys
+
+import torch
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+def predict_fn(input_data, model):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = model.to(device)
+    input_data = input_data.to(device)
+    return model(input_data)
+
+
+def save_model(model, model_dir):
+    logger.info("Saving the model to {}.".format(model_dir))
+    path = os.path.join(model_dir, 'model.pt')
+    torch.jit.save(model, path)

test/resources/mnist/default_model_eia/code/mnist.py

@@ -0,0 +1,35 @@
+# Copyright 2019-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+import logging
+import os
+import sys
+
+import torch
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+def predict_fn(input_data, model):
+    logger.info('Performing EIA inference with Torch JIT context with input of size {}'.format(input_data.shape))
+    # With EI, client instance should be CPU for cost-efficiency.
+    # Sub-graphs with unsupported arguments run locally. Server runs with CUDA
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = model.to(device)
+    input_data = input_data.to(device)
+    with torch.no_grad():
+        # Set the target device to the accelerator ordinal
+        with torch.jit.optimized_execution(True, {'target_device': 'eia:0'}):
+            return model(input_data)

test/resources/mnist/default_model_eia/model.pt

@@ -55,3 +55,9 @@ def model_fn(model_dir):
     with open(os.path.join(model_dir, 'torch_model.pth'), 'rb') as f:
         model.load_state_dict(torch.load(f))
     return model
+
+
+def save_model(model, model_dir):
+    logger.info("Saving the model to {}.".format(model_dir))
+    path = os.path.join(model_dir, 'model.pt')
+    torch.jit.save(model, path)

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

@@ -58,8 +58,14 @@ def eia_inference_handler():
 
 
 def test_default_model_fn(inference_handler):
-    with pytest.raises(NotImplementedError):
-        inference_handler.default_model_fn("model_dir")
+    with mock.patch("sagemaker_pytorch_serving_container.default_pytorch_inference_handler.os") as mock_os:
+        mock_os.getenv.return_value = "true"
+        mock_os.path.join.return_value = "model_dir"
+        mock_os.path.exists.return_value = True
+        with mock.patch("torch.jit.load") as mock_torch:
+            mock_torch.return_value = DummyModel()
+            model = inference_handler.default_model_fn("model_dir")
+    assert model is not None
 
 
 def test_default_input_fn_json(inference_handler, tensor):