change: make Local Mode export artifacts even after failure (#746)

Author: laurenyu

src/sagemaker/local/image.py

@@ -137,14 +137,15 @@ def train(self, input_data_config, output_data_config, hyperparameters, job_name
             # which contains the exit code and append the command line to it.
             msg = "Failed to run: %s, %s" % (compose_command, str(e))
             raise RuntimeError(msg)
+        finally:
+            artifacts = self.retrieve_artifacts(compose_data, output_data_config, job_name)
 
-        artifacts = self.retrieve_artifacts(compose_data, output_data_config, job_name)
+            # free up the training data directory as it may contain
+            # lots of data downloaded from S3. This doesn't delete any local
+            # data that was just mounted to the container.
+            dirs_to_delete = [data_dir, shared_dir]
+            self._cleanup(dirs_to_delete)
 
-        # free up the training data directory as it may contain
-        # lots of data downloaded from S3. This doesn't delete any local
-        # data that was just mounted to the container.
-        dirs_to_delete = [data_dir, shared_dir]
-        self._cleanup(dirs_to_delete)
         # Print our Job Complete line to have a similar experience to training on SageMaker where you
         # see this line at the end.
         print('===== Job Complete =====')

tests/integ/test_local_mode.py

@@ -1,4 +1,4 @@
-# Copyright 2017-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2017-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
@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 
 import os
+import tarfile
 
 import boto3
 import numpy
@@ -318,6 +319,26 @@ def test_mxnet_local_data_local_script(mxnet_full_version):
             mx.delete_endpoint()
 
 
+@pytest.mark.local_mode
+def test_mxnet_training_failure(sagemaker_local_session, mxnet_full_version, tmpdir):
+    script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'failure_script.py')
+
+    mx = MXNet(entry_point=script_path,
+               role='SageMakerRole',
+               framework_version=mxnet_full_version,
+               py_version=PYTHON_VERSION,
+               train_instance_count=1,
+               train_instance_type='local',
+               sagemaker_session=sagemaker_local_session,
+               output_path='file://{}'.format(tmpdir))
+
+    with pytest.raises(RuntimeError):
+        mx.fit()
+
+    with tarfile.open(os.path.join(str(tmpdir), 'output.tar.gz')) as tar:
+        tar.getmember('failure')
+
+
 @pytest.mark.local_mode
 def test_local_transform_mxnet(sagemaker_local_session, tmpdir, mxnet_full_version):
     data_path = os.path.join(DATA_DIR, 'mxnet_mnist')

tests/unit/test_image.py

@@ -1,4 +1,4 @@
-# Copyright 2017-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2017-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
@@ -303,10 +303,11 @@ def test_check_output():
 
 @patch('sagemaker.local.local_session.LocalSession', Mock())
 @patch('sagemaker.local.image._stream_output', Mock())
-@patch('sagemaker.local.image._SageMakerContainer._cleanup', Mock())
+@patch('sagemaker.local.image._SageMakerContainer._cleanup')
+@patch('sagemaker.local.image._SageMakerContainer.retrieve_artifacts')
 @patch('sagemaker.local.data.get_data_source_instance')
 @patch('subprocess.Popen')
-def test_train(popen, get_data_source_instance, tmpdir, sagemaker_session):
+def test_train(popen, get_data_source_instance, retrieve_artifacts, cleanup, tmpdir, sagemaker_session):
     data_source = Mock()
     data_source.get_root_dir.return_value = 'foo'
     get_data_source_instance.return_value = data_source
@@ -342,6 +343,9 @@ def test_train(popen, get_data_source_instance, tmpdir, sagemaker_session):
         assert os.path.exists(os.path.join(sagemaker_container.container_root, 'output'))
         assert os.path.exists(os.path.join(sagemaker_container.container_root, 'output/data'))
 
+    retrieve_artifacts.assert_called_once()
+    cleanup.assert_called_once()
+
 
 @patch('sagemaker.local.local_session.LocalSession', Mock())
 @patch('sagemaker.local.image._stream_output', Mock())
@@ -371,10 +375,11 @@ def test_train_with_hyperparameters_without_job_name(get_data_source_instance, t
 
 @patch('sagemaker.local.local_session.LocalSession', Mock())
 @patch('sagemaker.local.image._stream_output', side_effect=RuntimeError('this is expected'))
-@patch('sagemaker.local.image._SageMakerContainer._cleanup', Mock())
+@patch('sagemaker.local.image._SageMakerContainer._cleanup')
+@patch('sagemaker.local.image._SageMakerContainer.retrieve_artifacts')
 @patch('sagemaker.local.data.get_data_source_instance')
 @patch('subprocess.Popen', Mock())
-def test_train_error(get_data_source_instance, _stream_output, tmpdir, sagemaker_session):
+def test_train_error(get_data_source_instance, retrieve_artifacts, cleanup, _stream_output, tmpdir, sagemaker_session):
     data_source = Mock()
     data_source.get_root_dir.return_value = 'foo'
     get_data_source_instance.return_value = data_source
@@ -391,6 +396,9 @@ def test_train_error(get_data_source_instance, _stream_output, tmpdir, sagemaker
 
         assert 'this is expected' in str(e)
 
+    retrieve_artifacts.assert_called_once()
+    cleanup.assert_called_once()
+
 
 @patch('sagemaker.local.local_session.LocalSession', Mock())
 @patch('sagemaker.local.image._stream_output', Mock())