Fix local mode error message handling (#272)

Author: laurenyu

CHANGELOG.rst

@@ -6,6 +6,7 @@ CHANGELOG
 ========
 
 * bug-fix: Session: include role path in ``get_execution_role()`` result
+* bug-fix: Local Mode: fix RuntimeError handling
 
 1.5.2
 =====

src/sagemaker/local/image.py

@@ -112,7 +112,7 @@ def train(self, input_data_config, hyperparameters):
         except RuntimeError as e:
             # _stream_output() doesn't have the command line. We will handle the exception
             # which contains the exit code and append the command line to it.
-            msg = "Failed to run: %s, %s" % (compose_command, e.message)
+            msg = "Failed to run: %s, %s" % (compose_command, str(e))
             raise RuntimeError(msg)
 
         s3_artifacts = self.retrieve_artifacts(compose_data)
@@ -516,7 +516,7 @@ def run(self):
         except RuntimeError as e:
             # _stream_output() doesn't have the command line. We will handle the exception
             # which contains the exit code and append the command line to it.
-            msg = "Failed to run: %s, %s" % (self.command, e.message)
+            msg = "Failed to run: %s, %s" % (self.command, str(e))
             raise RuntimeError(msg)
 
     def down(self):

tests/unit/test_image.py

@@ -246,6 +246,25 @@ def test_train(_download_folder, _cleanup, popen, _stream_output, LocalSession,
                 assert config['services'][h]['command'] == 'train'
 
 
+@patch('sagemaker.local.local_session.LocalSession')
+@patch('sagemaker.local.image._stream_output', side_effect=RuntimeError('this is expected'))
+@patch('subprocess.Popen')
+@patch('sagemaker.local.image._SageMakerContainer._cleanup')
+@patch('sagemaker.local.image._SageMakerContainer._download_folder')
+def test_train_error(_download_folder, _cleanup, popen, _stream_output, LocalSession, tmpdir, sagemaker_session):
+    directories = [str(tmpdir.mkdir('container-root')), str(tmpdir.mkdir('data'))]
+
+    with patch('sagemaker.local.image._SageMakerContainer._create_tmp_folder', side_effect=directories):
+        instance_count = 2
+        image = 'my-image'
+        sagemaker_container = _SageMakerContainer('local', instance_count, image, sagemaker_session=sagemaker_session)
+
+        with pytest.raises(RuntimeError) as e:
+            sagemaker_container.train(INPUT_DATA_CONFIG, HYPERPARAMETERS)
+
+        assert 'this is expected' in str(e)
+
+
 @patch('sagemaker.local.local_session.LocalSession')
 @patch('sagemaker.local.image._stream_output')
 @patch('subprocess.Popen')