fix: provide option to override capture error

Author: Satish Pasumarthi

src/sagemaker_pytorch_container/training.py

@@ -81,12 +81,17 @@ def train(training_environment):
             runner_type = runner.PyTorchXLARunnerType
             logger.info('Invoking PT-XLA Runner')
     logger.info('Invoking user training script.')
+
+    # get capture_error from framework parameters
+    capture_error = training_environment.additional_framework_parameters.get("sagemaker_toolkit_capture_error", True)
+    logger.info(f'capture_error is {capture_error}. Default is True')
+
     try:
         entry_point.run(uri=training_environment.module_dir,
                         user_entry_point=training_environment.user_entry_point,
                         args=training_environment.to_cmd_args(),
                         env_vars=training_environment.to_env_vars(),
-                        capture_error=True,
+                        capture_error=capture_error,
                         runner_type=runner_type)
     except errors.ExecuteUserScriptError as err:
         message = str(err)

test/unit/test_train.py

@@ -74,6 +74,20 @@ def test_train(run_entry_point, training_env):
                                        runner_type=runner.ProcessRunnerType)
 
 
+@patch('sagemaker_training.entry_point.run')
+@patch('socket.gethostbyname', MagicMock())
+def test_train_no_capture_error(run_entry_point, training_env):
+    training_env.additional_framework_parameters["sagemaker_toolkit_capture_error"] = False
+    train(training_env)
+
+    run_entry_point.assert_called_with(uri=training_env.module_dir,
+                                       user_entry_point=training_env.user_entry_point,
+                                       args=training_env.to_cmd_args(),
+                                       env_vars=training_env.to_env_vars(),
+                                       capture_error=False,
+                                       runner_type=runner.ProcessRunnerType)
+
+
 @patch("sagemaker_training.entry_point.run")
 @patch('socket.gethostbyname', MagicMock())
 def test_train_smdataparallel(run_module, training_env):