add elastic inference test

test/integration/__init__.py

@@ -15,3 +15,7 @@
 import os
 
 RESOURCE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'resources'))
+
+# EI is currently only supported in the following regions
+# regions were derived from https://aws.amazon.com/machine-learning/elastic-inference/pricing/
+EI_SUPPORTED_REGIONS = ['us-east-1', 'us-east-2', 'us-west-2', 'eu-west-1', 'ap-northeast-1', 'ap-northeast-2']

test/integration/sagemaker/test_elastic_inference.py

@@ -0,0 +1,69 @@
+#  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 as np
+import pytest
+from sagemaker.mxnet import MXNetModel
+from sagemaker.utils import sagemaker_timestamp
+
+from test.integration import EI_SUPPORTED_REGIONS, RESOURCE_PATH
+from test.integration.sagemaker.timeout import timeout_and_delete_endpoint_by_name
+
+DEFAULT_HANDLER_PATH = os.path.join(RESOURCE_PATH, 'default_handlers')
+SCRIPT_PATH = os.path.join(DEFAULT_HANDLER_PATH, 'model', 'code', 'empty_module.py')
+
+
+@pytest.fixture(autouse=True)
+def skip_if_no_accelerator(accelerator_type):
+    if accelerator_type is None:
+        pytest.skip('Skipping because accelerator type was not provided')
+
+
+@pytest.fixture(autouse=True)
+def skip_if_non_supported_ei_region(region):
+    if region not in EI_SUPPORTED_REGIONS:
+        pytest.skip('EI is not supported in {}'.format(region))
+
+
+@pytest.fixture
+def pretrained_model_data(region):
+    return 's3://sagemaker-sample-data-{}/mxnet/model/resnet/resnet_50.tar.gz'.format(region)
+
+
+@pytest.mark.skip_if_non_supported_ei_region
+@pytest.mark.skip_if_no_accelerator
+def test_deploy_elastic_inference(pretrained_model_data, ecr_image, sagemaker_session,
+                                  instance_type, accelerator_type, framework_version):
+    endpoint_name = 'test-mxnet-ei-deploy-model-{}'.format(sagemaker_timestamp())
+
+    with timeout_and_delete_endpoint_by_name(endpoint_name=endpoint_name, sagemaker_session=sagemaker_session,
+                                             minutes=20):
+        model = MXNetModel(model_data=pretrained_model_data,
+                           entry_point=SCRIPT_PATH,
+                           role='SageMakerRole',
+                           image=ecr_image,
+                           framework_version=framework_version,
+                           sagemaker_session=sagemaker_session)
+
+        predictor = model.deploy(initial_instance_count=1,
+                                 instance_type=instance_type,
+                                 accelerator_type=accelerator_type,
+                                 endpoint_name=endpoint_name)
+
+        random_input = np.zeros(shape=(1, 3, 224, 224))
+
+        predict_response = predictor.predict(random_input.tolist())
+        assert predict_response

test/integration/sagemaker/timeout.py

@@ -0,0 +1,87 @@
+#  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
+
+from contextlib import contextmanager
+import logging
+import signal
+
+from botocore.exceptions import ClientError
+
+LOGGER = logging.getLogger('timeout')
+
+
+class TimeoutError(Exception):
+    pass
+
+
+@contextmanager
+def timeout(seconds=0, minutes=0, hours=0):
+    """Add a signal-based timeout to any block of code.
+    If multiple time units are specified, they will be added together to determine time limit.
+
+    Usage:
+
+    with timeout(seconds=5):
+        my_slow_function(...)
+
+
+    Args:
+        - seconds: The time limit, in seconds.
+        - minutes: The time limit, in minutes.
+        - hours: The time limit, in hours.
+    """
+
+    limit = seconds + 60 * minutes + 3600 * hours
+
+    def handler(signum, frame):
+        raise TimeoutError('timed out after {} seconds'.format(limit))
+
+    try:
+        signal.signal(signal.SIGALRM, handler)
+        signal.alarm(limit)
+
+        yield
+    finally:
+        signal.alarm(0)
+
+
+@contextmanager
+def timeout_and_delete_endpoint(estimator, seconds=0, minutes=0, hours=0):
+    with timeout(seconds=seconds, minutes=minutes, hours=hours) as t:
+        try:
+            yield [t]
+        finally:
+            try:
+                estimator.delete_endpoint()
+                LOGGER.info('deleted endpoint')
+            except ClientError as ce:
+                if ce.response['Error']['Code'] == 'ValidationException':
+                    # avoids the inner exception to be overwritten
+                    pass
+
+
+@contextmanager
+def timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session, seconds=0, minutes=0,
+                                        hours=0):
+    with timeout(seconds=seconds, minutes=minutes, hours=hours) as t:
+        try:
+            yield [t]
+        finally:
+            try:
+                sagemaker_session.delete_endpoint(endpoint_name)
+                LOGGER.info('deleted endpoint {}'.format(endpoint_name))
+            except ClientError as ce:
+                if ce.response['Error']['Code'] == 'ValidationException':
+                    # avoids the inner exception to be overwritten
+                    pass