Add integration tests for basic training failure cases (#33)

Author: laurenyu

tests/data/iris/failure_script.py

@@ -0,0 +1,3 @@
+def estimator_fn(run_config, params):
+    """For use with integration tests expecting failures."""
+    raise Exception('This failure is expected.')

tests/data/mxnet_mnist/failure_script.py

@@ -0,0 +1,3 @@
+def train(**kwargs):
+    """For use with integration tests expecting failures."""
+    raise Exception('This failure is expected.')

tests/integ/test_mxnet_train.py

@@ -70,3 +70,20 @@ def test_deploy_model(mxnet_training_job, sagemaker_session):
 
         data = numpy.zeros(shape=(1, 1, 28, 28))
         predictor.predict(data)
+
+
+def test_failed_training_job(sagemaker_session):
+    with timeout(minutes=15):
+        script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'failure_script.py')
+        data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
+
+        mx = MXNet(entry_point=script_path, role='SageMakerRole',
+                   train_instance_count=1, train_instance_type='ml.c4.xlarge',
+                   sagemaker_session=sagemaker_session)
+
+        train_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'train'),
+                                                       key_prefix='integ-test-data/mxnet_mnist/train-failure')
+
+        with pytest.raises(ValueError) as e:
+            mx.fit(train_input)
+        assert 'This failure is expected' in str(e.value)

tests/integ/test_tf.py

@@ -19,6 +19,8 @@
 from tests.integ import DATA_DIR, REGION
 from tests.integ.timeout import timeout_and_delete_endpoint, timeout
 
+DATA_PATH = os.path.join(DATA_DIR, 'iris', 'data')
+
 
 @pytest.fixture(scope='module')
 def sagemaker_session():
@@ -28,7 +30,6 @@ def sagemaker_session():
 def test_tf(sagemaker_session):
     with timeout(minutes=15):
         script_path = os.path.join(DATA_DIR, 'iris', 'iris-dnn-classifier.py')
-        data_path = os.path.join(DATA_DIR, 'iris', 'data')
 
         estimator = TensorFlow(entry_point=script_path,
                                role='SageMakerRole',
@@ -40,7 +41,7 @@ def test_tf(sagemaker_session):
                                sagemaker_session=sagemaker_session,
                                base_job_name='test-tf')
 
-        inputs = estimator.sagemaker_session.upload_data(path=data_path, key_prefix='integ-test-data/tf_iris')
+        inputs = estimator.sagemaker_session.upload_data(path=DATA_PATH, key_prefix='integ-test-data/tf_iris')
         estimator.fit(inputs)
         print('job succeeded: {}'.format(estimator.latest_training_job.name))
 
@@ -49,3 +50,22 @@ def test_tf(sagemaker_session):
 
         result = json_predictor.predict([6.4, 3.2, 4.5, 1.5])
         print('predict result: {}'.format(result))
+
+
+def test_failed_tf_training(sagemaker_session):
+    with timeout(minutes=15):
+        script_path = os.path.join(DATA_DIR, 'iris', 'failure_script.py')
+        estimator = TensorFlow(entry_point=script_path,
+                               role='SageMakerRole',
+                               training_steps=1,
+                               evaluation_steps=1,
+                               hyperparameters={'input_tensor_name': 'inputs'},
+                               train_instance_count=1,
+                               train_instance_type='ml.c4.xlarge',
+                               sagemaker_session=sagemaker_session)
+
+        inputs = estimator.sagemaker_session.upload_data(path=DATA_PATH, key_prefix='integ-test-data/tf-failure')
+
+        with pytest.raises(ValueError) as e:
+            estimator.fit(inputs)
+        assert 'This failure is expected' in str(e.value)