add local integration tests (#10)

Author: Dan

test/__init__.py

@@ -0,0 +1,117 @@
+#  Copyright 2019 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://www.apache.org/licenses/LICENSE-2.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 boto3
+import pytest
+from sagemaker import LocalSession, Session
+from sagemaker.mxnet import MXNet
+
+logger = logging.getLogger(__name__)
+logging.getLogger('boto').setLevel(logging.INFO)
+logging.getLogger('botocore').setLevel(logging.INFO)
+logging.getLogger('factory.py').setLevel(logging.INFO)
+logging.getLogger('auth.py').setLevel(logging.INFO)
+logging.getLogger('connectionpool.py').setLevel(logging.INFO)
+
+SCRIPT_PATH = os.path.dirname(os.path.realpath(__file__))
+
+
+def pytest_addoption(parser):
+    parser.addoption('--docker-base-name', default='preprod-mxnet-serving')
+    parser.addoption('--region', default='us-west-2')
+    parser.addoption('--framework-version', default=MXNet.LATEST_VERSION)
+    parser.addoption('--py-version', default='3', choices=['2', '3'])
+    parser.addoption('--processor', default='cpu', choices=['gpu', 'cpu'])
+    parser.addoption('--aws-id', default=None)
+    parser.addoption('--instance-type', default=None)
+    parser.addoption('--accelerator-type', default=None)
+    # If not specified, will default to {framework-version}-{processor}-py{py-version}
+    parser.addoption('--tag', default=None)
+
+
+@pytest.fixture(scope='session')
+def docker_base_name(request):
+    return request.config.getoption('--docker-base-name')
+
+
+@pytest.fixture(scope='session')
+def region(request):
+    return request.config.getoption('--region')
+
+
+@pytest.fixture(scope='session')
+def framework_version(request):
+    return request.config.getoption('--framework-version')
+
+
+@pytest.fixture(scope='session')
+def py_version(request):
+    return int(request.config.getoption('--py-version'))
+
+
+@pytest.fixture(scope='session')
+def processor(request):
+    return request.config.getoption('--processor')
+
+
+@pytest.fixture(scope='session')
+def aws_id(request):
+    return request.config.getoption('--aws-id')
+
+
+@pytest.fixture(scope='session')
+def tag(request, framework_version, processor, py_version):
+    provided_tag = request.config.getoption('--tag')
+    default_tag = '{}-{}-py{}'.format(framework_version, processor, py_version)
+    return provided_tag if provided_tag is not None else default_tag
+
+
+@pytest.fixture(scope='session')
+def instance_type(request, processor):
+    provided_instance_type = request.config.getoption('--instance-type')
+    default_instance_type = 'ml.c4.xlarge' if processor == 'cpu' else 'ml.p2.xlarge'
+    return provided_instance_type if provided_instance_type is not None else default_instance_type
+
+
+@pytest.fixture(scope='session')
+def accelerator_type(request):
+    return request.config.getoption('--accelerator-type')
+
+
+@pytest.fixture(scope='session')
+def docker_image(docker_base_name, tag):
+    return '{}:{}'.format(docker_base_name, tag)
+
+
+@pytest.fixture(scope='session')
+def ecr_image(aws_id, docker_base_name, tag, region):
+    return '{}.dkr.ecr.{}.amazonaws.com/{}:{}'.format(aws_id, region, docker_base_name, tag)
+
+
+@pytest.fixture(scope='session')
+def sagemaker_session(region):
+    return Session(boto_session=boto3.Session(region_name=region))
+
+
+@pytest.fixture(scope='session')
+def sagemaker_local_session(region):
+    return LocalSession(boto_session=boto3.Session(region_name=region))
+
+
+@pytest.fixture(scope='session')
+def local_instance_type(processor):
+    return 'local' if processor == 'cpu' else 'local_gpu'

test/conftest.py

@@ -0,0 +1,17 @@
+# Copyright 2019 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 os
+
+RESOURCE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'resources'))

test/integration/__init__.py

@@ -0,0 +1,45 @@
+# Copyright 2019 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
+
+from contextlib import contextmanager
+import fcntl
+import os
+import tarfile
+import time
+
+from test.integration import RESOURCE_PATH
+
+LOCK_PATH = os.path.join(RESOURCE_PATH, 'local_mode_lock')
+
+
+@contextmanager
+def lock():
+    # Since Local Mode uses the same port for serving, we need a lock in order
+    # to allow concurrent test execution.
+    local_mode_lock_fd = open(LOCK_PATH, 'w')
+    local_mode_lock = local_mode_lock_fd.fileno()
+
+    fcntl.lockf(local_mode_lock, fcntl.LOCK_EX)
+
+    try:
+        yield
+    finally:
+        time.sleep(5)
+        fcntl.lockf(local_mode_lock, fcntl.LOCK_UN)
+
+
+def assert_output_files_exist(output_path, directory, files):
+    with tarfile.open(os.path.join(output_path, '{}.tar.gz'.format(directory))) as tar:
+        for f in files:
+            tar.getmember(f)

test/integration/local/local_mode_utils.py

@@ -0,0 +1,53 @@
+#  Copyright 2019 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://www.apache.org/licenses/LICENSE-2.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 os
+
+import pytest
+import requests
+from sagemaker.mxnet.model import MXNetModel
+
+import local_mode_utils
+from test.integration import RESOURCE_PATH
+
+DEFAULT_HANDLER_PATH = os.path.join(RESOURCE_PATH, 'default_handlers')
+MODEL_PATH = os.path.join(DEFAULT_HANDLER_PATH, 'model')
+SCRIPT_PATH = os.path.join(MODEL_PATH, 'code', 'empty_module.py')
+
+
+@pytest.fixture(scope='module')
+def predictor(docker_image, sagemaker_local_session, local_instance_type):
+    model = MXNetModel('file://{}'.format(MODEL_PATH),
+                       'SageMakerRole',
+                       SCRIPT_PATH,
+                       image=docker_image,
+                       sagemaker_session=sagemaker_local_session)
+
+    with local_mode_utils.lock():
+        try:
+            predictor = model.deploy(1, local_instance_type)
+            yield predictor
+        finally:
+            sagemaker_local_session.delete_endpoint(model.endpoint_name)
+
+
+def test_default_model_fn(predictor):
+    input = [[1, 2]]
+    output = predictor.predict(input)
+    assert [[4.9999918937683105]] == output
+
+
+def test_default_model_fn_via_requests(predictor):
+    r = requests.post('http://localhost:8080/invocations', json=[[1, 2]])
+    assert [[4.9999918937683105]] == r.json()

test/integration/local/test_default_model_fn.py

@@ -0,0 +1,45 @@
+#  Copyright 2019 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://www.apache.org/licenses/LICENSE-2.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 json
+import os
+
+from sagemaker.mxnet.model import MXNetModel
+
+import local_mode_utils
+from test.integration import RESOURCE_PATH
+
+GLUON_PATH = os.path.join(RESOURCE_PATH, 'gluon_hosting')
+MODEL_PATH = os.path.join(GLUON_PATH, 'model')
+SCRIPT_PATH = os.path.join(MODEL_PATH, 'code', 'gluon.py')
+
+
+# The image should support serving Gluon-created models.
+def test_gluon_hosting(docker_image, sagemaker_local_session, local_instance_type):
+    model = MXNetModel('file://{}'.format(MODEL_PATH),
+                       'SageMakerRole',
+                       SCRIPT_PATH,
+                       image=docker_image,
+                       sagemaker_session=sagemaker_local_session)
+
+    with open(os.path.join(RESOURCE_PATH, 'mnist_images', '04.json'), 'r') as f:
+        input = json.load(f)
+
+    with local_mode_utils.lock():
+        try:
+            predictor = model.deploy(1, local_instance_type)
+            output = predictor.predict(input)
+            assert [4.0] == output
+        finally:
+            sagemaker_local_session.delete_endpoint(model.endpoint_name)

test/integration/local/test_gluon_hosting.py

@@ -0,0 +1,45 @@
+#  Copyright 2019 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://www.apache.org/licenses/LICENSE-2.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 json
+import os
+
+from sagemaker.mxnet.model import MXNetModel
+
+import local_mode_utils
+from test.integration import RESOURCE_PATH
+
+HOSTING_RESOURCE_PATH = os.path.join(RESOURCE_PATH, 'dummy_hosting')
+MODEL_PATH = os.path.join(HOSTING_RESOURCE_PATH, 'code')
+SCRIPT_PATH = os.path.join(HOSTING_RESOURCE_PATH, 'code', 'dummy_hosting_module.py')
+
+
+# The image should use the model_fn and transform_fn defined
+# in the user-provided script when serving.
+def test_hosting(docker_image, sagemaker_local_session, local_instance_type):
+    model = MXNetModel('file://{}'.format(MODEL_PATH),
+                       'SageMakerRole',
+                       SCRIPT_PATH,
+                       image=docker_image,
+                       sagemaker_session=sagemaker_local_session)
+
+    input = json.dumps({'some': 'json'})
+
+    with local_mode_utils.lock():
+        try:
+            predictor = model.deploy(1, local_instance_type)
+            output = predictor.predict(input)
+            assert input == output
+        finally:
+            sagemaker_local_session.delete_endpoint(model.endpoint_name)

test/integration/local/test_hosting.py

@@ -0,0 +1,45 @@
+#  Copyright 2019 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://www.apache.org/licenses/LICENSE-2.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 os
+
+import numpy
+from sagemaker.mxnet import MXNetModel
+
+import local_mode_utils
+from test.integration import RESOURCE_PATH
+
+ONNX_PATH = os.path.join(RESOURCE_PATH, 'onnx')
+MODEL_PATH = os.path.join(ONNX_PATH, 'onnx_model')
+SCRIPT_PATH = os.path.join(MODEL_PATH, 'code', 'onnx_import.py')
+
+
+def test_onnx_import(docker_image, sagemaker_local_session, local_instance_type):
+    model = MXNetModel('file://{}'.format(MODEL_PATH),
+                       'SageMakerRole',
+                       SCRIPT_PATH,
+                       image=docker_image,
+                       sagemaker_session=sagemaker_local_session)
+
+    input = numpy.zeros(shape=(1, 1, 28, 28))
+
+    with local_mode_utils.lock():
+        try:
+            predictor = model.deploy(1, local_instance_type)
+            output = predictor.predict(input)
+        finally:
+            sagemaker_local_session.delete_endpoint(model.endpoint_name)
+
+    # Check that there is a probability for each possible class in the prediction
+    assert len(output[0]) == 10

test/integration/local/test_onnx.py

@@ -0,0 +1,14 @@
+#  Copyright 2019 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://www.apache.org/licenses/LICENSE-2.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.
+
+# nothing here... we are testing default model loading and handlers

test/resources/default_handlers/model/code/empty_module.py

@@ -0,0 +1 @@
+[{"shape": [1, 2], "name": "data"}]