Add integration tests to run training jobs with sagemaker

test/conftest.py

@@ -36,6 +36,8 @@ def pytest_addoption(parser):
     parser.addoption('--framework-version', default='1.10.0')
     parser.addoption('--processor', default='cpu', choices=['gpu', 'cpu'])
     parser.addoption('--py-version', default='3', choices=['2', '3'])
+    parser.addoption('--account-id', default='142577830533')
+    parser.addoption('--instance-type', default=None)
 
 
 @pytest.fixture(scope='session')
@@ -80,6 +82,17 @@ def sagemaker_local_session(region):
     return LocalSession(boto_session=boto3.Session(region_name=region))
 
 
+@pytest.fixture(scope='session')
+def account_id(request):
+    return request.config.getoption('--account-id')
+
+
+@pytest.fixture(scope='session')
+def instance_type(request, processor):
+    return request.config.getoption('--instance-type') or \
+        'ml.c4.xlarge' if processor == 'cpu' else 'ml.p2.xlarge'
+
+
 @pytest.fixture(autouse=True)
 def skip_by_device_type(request, processor):
     is_gpu = (processor == 'gpu')
@@ -91,3 +104,9 @@ def skip_by_device_type(request, processor):
 @pytest.fixture(scope='session')
 def docker_image(docker_base_name, tag):
     return '{}:{}'.format(docker_base_name, tag)
+
+
+@pytest.fixture(scope='session')
+def ecr_image(account_id, docker_base_name, tag, region):
+    return '{}.dkr.ecr.{}.amazonaws.com/{}:{}'.format(
+        account_id, region, docker_base_name, tag)

test/integration/sagemaker/test_mnist.py

@@ -0,0 +1,35 @@
+# Copyright 2017-2018 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
+
+from sagemaker.tensorflow import TensorFlow
+
+
+def test_mnist(sagemaker_session, ecr_image, instance_type):
+    resource_path = os.path.join(os.path.dirname(__file__), '../..', 'resources')
+    script = os.path.join(resource_path, 'mnist', 'mnist.py')
+    estimator = TensorFlow(entry_point=script,
+                           role='SageMakerRole',
+                           training_steps=1,
+                           evaluation_steps=1,
+                           train_instance_count=1,
+                           train_instance_type=instance_type,
+                           sagemaker_session=sagemaker_session,
+                           image_name=ecr_image,
+                           base_job_name='test-sagemaker-mnist')
+    inputs = estimator.sagemaker_session.upload_data(
+        path=os.path.join(resource_path, 'mnist', 'data'),
+        key_prefix='scriptmode/mnist')
+    estimator.fit(inputs)