fix: Revert "change: stream stderr even when capture_error is True (#233)" (#268)

This reverts commit f4d3456.

Author: ajaykarpur

src/sagemaker_containers/_process.py

@@ -13,7 +13,6 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-import io
 import os
 import subprocess
 import sys
@@ -25,82 +24,30 @@
 
 
 def create(cmd, error_class, cwd=None, capture_error=False, **kwargs):
-    """Create subprocess.Popen object for the given command.
-
-    Args:
-        cmd (list): The command to be run.
-        error_class (cls): The class to use when raising an exception.
-        cwd (str): The location from which to run the command (default: None).
-            If None, this defaults to the ``code_dir`` of the environment.
-        capture_error (bool): whether or not to direct stderr to a stream
-            that can later be read (default: False).
-        **kwargs: Extra arguments that are passed to the subprocess.Popen constructor.
-
-    Returns:
-        subprocess.Popen: the process for the given command
-
-    Raises:
-        error_class: if there is an exception raised when creating the process
-    """
+    """Placeholder docstring"""
     try:
-        # Capture both so that we can control the order of when stdout and stderr are streamed
-        stdout = subprocess.PIPE if capture_error else None
         stderr = subprocess.PIPE if capture_error else None
-
         return subprocess.Popen(
-            cmd, env=os.environ, cwd=cwd or _env.code_dir, stdout=stdout, stderr=stderr, **kwargs
+            cmd, env=os.environ, cwd=cwd or _env.code_dir, stderr=stderr, **kwargs
         )
     except Exception as e:  # pylint: disable=broad-except
         six.reraise(error_class, error_class(e), sys.exc_info()[2])
 
 
 def check_error(cmd, error_class, capture_error=False, **kwargs):
     # type: (List[str], type, bool, Mapping[str, object]) -> subprocess.Popen
-    """Run a commmand, raising an exception if there is an error.
-
-    Args:
-        cmd (list): The command to be run.
-        error_class (cls): The class to use when raising an exception.
-        capture_error (bool): whether or not to include stderr in
-            the exception message (default: False). In either case,
-            stderr is streamed to the process's output.
-        **kwargs: Extra arguments that are passed to the subprocess.Popen constructor.
-
-    Returns:
-        subprocess.Popen: the process for the given command
-
-    Raises:
-        error_class: if there is an exception raised when creating the process
-    """
+    """Placeholder docstring"""
     process = create(cmd, error_class, capture_error=capture_error, **kwargs)
 
     if capture_error:
-        # Create a copy of stderr so that it can be read after being streamed
-        with io.BytesIO() as stderr_copy:
-            return_code = process.poll()
-            while return_code is None:
-                stdout = process.stdout.readline()
-                sys.stdout.write(stdout.decode("utf-8"))
-                stderr = process.stderr.readline()
-                sys.stdout.write(stderr.decode("utf-8"))
-
-                stderr_copy.write(stderr)
-                return_code = process.poll()
-
-            # Read the rest of stdout/stdin because readline() reads only one line at a time
-            stdout = process.stdout.read()
-            sys.stdout.write(stdout.decode("utf-8"))
-            stderr = process.stderr.read()
-            sys.stdout.write(stderr.decode("utf-8"))
-
-            stderr_copy.write(stderr)
-            full_stderr = stderr_copy.getvalue()
+        _, stderr = process.communicate()
+        return_code = process.poll()
     else:
-        full_stderr = None
+        stderr = None
         return_code = process.wait()
 
     if return_code:
-        raise error_class(return_code=return_code, cmd=" ".join(cmd), output=full_stderr)
+        raise error_class(return_code=return_code, cmd=" ".join(cmd), output=stderr)
     return process
 
 

test/unit/test_mpi.py

@@ -166,7 +166,6 @@ def test_mpi_master_run(training_env, popen, policy, ssh_client, path_exists):
             ],
             cwd=_env.code_dir,
             env=ANY,
-            stdout=None,
             stderr=None,
         )
 
@@ -262,7 +261,6 @@ def test_mpi_master_run_python(
             ],
             cwd=_env.code_dir,
             env=ANY,
-            stdout=None,
             stderr=None,
         )
 

test/unit/test_process.py

@@ -67,28 +67,20 @@ def test_run_bash(log, popen, entry_point_type_script):
         _process.ProcessRunner("launcher.sh", ["--lr", "1 3"], {}).run()
 
     cmd = ["/bin/sh", "-c", "./launcher.sh --lr '1 3'"]
-    popen.assert_called_with(cmd, cwd=_env.code_dir, env=os.environ, stdout=None, stderr=None)
+    popen.assert_called_with(cmd, cwd=_env.code_dir, env=os.environ, stderr=None)
     log.assert_called_with(cmd, {})
 
 
 @patch("subprocess.Popen")
 @patch("sagemaker_containers._logging.log_script_invocation")
-def test_run_python_capture_error(log, popen, entry_point_type_script):
-    mock_process = MagicMock()
-    mock_process.stdout.readline.return_value = b"stdout"
-    mock_process.stderr.readline.return_value = b"stderr"
-    mock_process.stdout.read.return_value = b"stdout"
-    mock_process.stderr.read.return_value = b"stderr"
-    mock_process.poll.return_value = 1
-    popen.return_value = mock_process
+def test_run_python(log, popen, entry_point_type_script):
+    popen().communicate.return_value = (None, 0)
 
     with pytest.raises(_errors.ExecuteUserScriptError):
         _process.ProcessRunner("launcher.py", ["--lr", "13"], {}).run(capture_error=True)
 
     cmd = [sys.executable, "launcher.py", "--lr", "13"]
-    popen.assert_called_with(
-        cmd, cwd=_env.code_dir, env=os.environ, stdout=subprocess.PIPE, stderr=subprocess.PIPE
-    )
+    popen.assert_called_with(cmd, cwd=_env.code_dir, env=os.environ, stderr=subprocess.PIPE)
     log.assert_called_with(cmd, {})
 
 
@@ -99,7 +91,7 @@ def test_run_module(log, popen, entry_point_type_module):
         _process.ProcessRunner("module.py", ["--lr", "13"], {}).run()
 
     cmd = [sys.executable, "-m", "module", "--lr", "13"]
-    popen.assert_called_with(cmd, cwd=_env.code_dir, env=os.environ, stdout=None, stderr=None)
+    popen.assert_called_with(cmd, cwd=_env.code_dir, env=os.environ, stderr=None)
     log.assert_called_with(cmd, {})