change: stream stderr even when capture_error is True (#233)

Author: laurenyu

src/sagemaker_containers/_process.py

@@ -13,6 +13,7 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
+import io
 import os
 import subprocess
 import sys
@@ -24,30 +25,82 @@
 
 
 def create(cmd, error_class, cwd=None, capture_error=False, **kwargs):
-    """Placeholder docstring"""
+    """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
+    """
     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, stderr=stderr, **kwargs
+            cmd, env=os.environ, cwd=cwd or _env.code_dir, stdout=stdout, 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
-    """Placeholder docstring"""
+    """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
+    """
     process = create(cmd, error_class, capture_error=capture_error, **kwargs)
 
     if capture_error:
-        _, stderr = process.communicate()
-        return_code = process.poll()
+        # 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()
     else:
-        stderr = None
+        full_stderr = None
         return_code = process.wait()
 
     if return_code:
-        raise error_class(return_code=return_code, cmd=" ".join(cmd), output=stderr)
+        raise error_class(return_code=return_code, cmd=" ".join(cmd), output=full_stderr)
     return process
 
 

test/unit/test_mpi.py

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

test/unit/test_process.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2018-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
@@ -67,20 +67,28 @@ def test_run_bash(log, popen, entry_point_type_script):
         _process.ProcessRunner("launcher.sh", ["--lr", "13"], {}).run()
 
     cmd = ["/bin/sh", "-c", "./launcher.sh --lr 13"]
-    popen.assert_called_with(cmd, cwd=_env.code_dir, env=os.environ, stderr=None)
+    popen.assert_called_with(cmd, cwd=_env.code_dir, env=os.environ, stdout=None, stderr=None)
     log.assert_called_with(cmd, {})
 
 
 @patch("subprocess.Popen")
 @patch("sagemaker_containers._logging.log_script_invocation")
-def test_run_python(log, popen, entry_point_type_script):
-    popen().communicate.return_value = (None, 0)
+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
 
     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, stderr=subprocess.PIPE)
+    popen.assert_called_with(
+        cmd, cwd=_env.code_dir, env=os.environ, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+    )
     log.assert_called_with(cmd, {})
 
 
@@ -91,7 +99,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, stderr=None)
+    popen.assert_called_with(cmd, cwd=_env.code_dir, env=os.environ, stdout=None, stderr=None)
     log.assert_called_with(cmd, {})