Repository: DataBiosphere/toil

Issue:
Change apiDockerCall to handle stream=True.
Currently, apiDockerCall returns one output of stdout and stderr joined together:

```
>>> cmd = '/bin/sh -c "echo hello stdout ; echo hello stderr >&2"'
>>> res = client.containers.run(cmd, stream=False)
>>> res.output
b'hello stderr\nhello stdout\n'
```

The API inside can be changed with stream=True, which returns an iterator:

```
>>> cmd = '/bin/sh -c "echo hello stdout ; echo hello stderr >&2"'
>>> res = client.containers.run(cmd, stream=True)
>>> next(res.output)
b'hello stdout\n'
>>> next(res.output)
b'hello stderr\n'
>>> next(res.output)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
```

Our apiDockerCall would need to change to handle the case where stream=True, so that we can return separate stdout and stderr streams.

┆Issue is synchronized with this [Jira Task](https://ucsc-cgl.atlassian.net/browse/TOIL-628)
┆Issue Number: TOIL-628



Discussion:
I'm not clear on what this means. You illustrate the underlying Docker API combining output and error streams. Do you want us to add a `stream=True` to apiDockerCall that streams output and also un-combines them somehow?
@DailyDreaming I'm not sure if using `stream=True` would fully separate stdout and stderr. For example, the `hello-world` image outputs multiple lines to stdout: 

```python
output = client.containers.run(image='hello-world', stream=True)
print(next(output))
print(next(output))
print(next(output))
```
This outputs: 
```
b'\n'
b'Hello from Docker!\n'
b'This message shows that your installation appears to be working correctly.\n'
```
with more to read.
@w-gao I see two ways of doing this.

The first is that the docker iterator appears to have a hidden "content" attribute with the `stderr` and `stdout` prefixed with `\x02\x00\x00\x00\x00\x00\x00\r` and `\x01\x00\x00\x00\x00\x00\x00\r` respectively:

```
import docker


client = docker.from_env(version='auto', timeout=3600)

cmd = '/bin/bash -c "' \
      'echo hello stdout ; ' \
      'echo hello stderr >&2 ; ' \
      'echo hello stdout ; ' \
      'echo hello stderr >&2 ; ' \
      'echo hello stdout ; ' \
      'echo hello stdout"'

stdout = client.containers.run(image='ubuntu', command=cmd, stream=True, stderr=True)

print(stdout._response.content)
```

```
b'\x01\x00\x00\x00\x00\x00\x00\rhello stdout\n\x02\x00\x00\x00\x00\x00\x00\rhello stderr\n\x02\x00\x00\x00\x00\x00\x00\rhello stderr\n\x01\x00\x00\x00\x00\x00\x00\rhello stdout\n\x01\x00\x00\x00\x00\x00\x00\rhello stdout\n\x01\x00\x00\x00\x00\x00\x00\rhello stdout\n'
```

OR... we add another function that wraps `container.exec_run()` which will separate the streams internally when `demux=True` (which I honestly was probably looking at by mistake when I created the issue):

```
>>> res = container.exec_run(cmd, stream=True, demux=True)
>>> next(res.output)
(b'hello stdout\n', None)
>>> next(res.output)
(None, b'hello stderr\n')
>>> next(res.output)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
```

@DailyDreaming I was looking at `container.exec_run()` though that requires to run in detach mode to get the container object. I'll tinker with it more or look into the first way. Thanks for the suggestions! 

Resolving patch:
diff --git a/src/toil/lib/docker.py b/src/toil/lib/docker.py
--- a/src/toil/lib/docker.py
+++ b/src/toil/lib/docker.py
@@ -1,18 +1,14 @@
 from __future__ import absolute_import
 
-try:
-    # In Python 3 we have this quote
-    from shlex import quote
-except ImportError:
-    # But in 2.7 we have this deprecated one
-    from pipes import quote
 
 import docker
 import base64
 import requests
 import logging
 import os
-import sys
+import struct
+from shlex import quote
+from docker.utils.socket import consume_socket_output, demux_adaptor
 from docker.errors import create_api_error_from_http_exception
 from docker.errors import ContainerError
 from docker.errors import ImageNotFound
@@ -56,6 +52,8 @@ def apiDockerCall(job,
                   environment=None,
                   stdout=None,
                   stderr=False,
+                  stream=False,
+                  demux=False,
                   streamfile=None,
                   timeout=365 * 24 * 60 * 60,
                   **kwargs):
@@ -79,7 +77,7 @@ def toil_job(job):
                       image='quay.io/ucgc_cgl/samtools:latest',
                       working_dir=working_dir,
                       parameters=parameters)
-                      
+
     Note that when run with detatch=False, or with detatch=True and stdout=True
     or stderr=True, this is a blocking call. When run with detatch=True and
     without output capture, the container is started and returned without
@@ -114,7 +112,13 @@ def toil_job(job):
                         Block and capture stderr to a file when detach=True
                         (default: False). Output capture defaults to output.log,
                         and can be specified with the "streamfile" kwarg.
-    :param str streamfile: Collect container output to this file if detach=True and 
+    :param bool stream: If True and detach=False, return a log generator instead
+                        of a string. Ignored if detach=True. (default: False).
+    :param bool demux: Similar to `demux` in container.exec_run(). If True and
+                       detach=False, returns a tuple of (stdout, stderr). If
+                       stream=True, returns a log generator with tuples of
+                       (stdout, stderr). Ignored if detach=True. (default: False).
+    :param str streamfile: Collect container output to this file if detach=True and
                         stderr and/or stdout are True. Defaults to "output.log".
     :param dict log_config: Specify the logs to return from the container.  See:
                       https://docker-py.readthedocs.io/en/stable/containers.html
@@ -138,8 +142,8 @@ def toil_job(job):
                    run command.  The list is 75 keywords total, for examples
                    and full documentation see:
                    https://docker-py.readthedocs.io/en/stable/containers.html
-                   
-    :returns: Returns the standard output/standard error text, as requested, when 
+
+    :returns: Returns the standard output/standard error text, as requested, when
               detatch=False. Returns the underlying
               docker.models.containers.Container object from the Docker API when
               detatch=True.
@@ -235,16 +239,31 @@ def toil_job(job):
                                         auto_remove=auto_remove,
                                         stdout=stdout,
                                         stderr=stderr,
+                                        # to get the generator if demux=True
+                                        stream=stream or demux,
                                         remove=remove,
                                         log_config=log_config,
                                         user=user,
                                         environment=environment,
                                         **kwargs)
-            return out
+
+            if demux is False:
+                return out
+
+            # If demux is True (i.e.: we want STDOUT and STDERR separated), we need to decode
+            # the raw response from the docker API and preserve the stream type this time.
+            response = out._response
+            gen = (demux_adaptor(*frame) for frame in _multiplexed_response_stream_helper(response))
+
+            if stream:
+                return gen
+            else:
+                return consume_socket_output(frames=gen, demux=True)
+
         else:
             if (stdout or stderr) and log_config is None:
                 logger.warning('stdout or stderr specified, but log_config is not set.  '
-                            'Defaulting to "journald".')
+                               'Defaulting to "journald".')
                 log_config = dict(type='journald')
 
             if stdout is None:
@@ -263,6 +282,7 @@ def toil_job(job):
                                               auto_remove=auto_remove,
                                               stdout=stdout,
                                               stderr=stderr,
+                                              stream=stream,
                                               remove=remove,
                                               log_config=log_config,
                                               user=user,
@@ -276,12 +296,12 @@ def toil_job(job):
                     # the container stops and there is no more output.
                     with open(streamfile, 'w') as f:
                         f.write(line)
-                        
+
             # If we didn't capture output, the caller will need to .wait() on
             # the container to know when it is done. Even if we did capture
             # output, the caller needs the container to get at the exit code.
             return container
-            
+
     except ContainerError:
         logger.error("Docker had non-zero exit.  Check your command: " + repr(command))
         raise
@@ -316,10 +336,10 @@ def dockerKill(container_name, gentleKill=False, timeout=365 * 24 * 60 * 60):
             this_container = client.containers.get(container_name)
     except NotFound:
         logger.debug("Attempted to stop container, but container != exist: ",
-                      container_name)
+                     container_name)
     except requests.exceptions.HTTPError as e:
         logger.debug("Attempted to stop container, but server gave an error: ",
-                      container_name)
+                     container_name)
         raise create_api_error_from_http_exception(e)
 
 
@@ -357,12 +377,35 @@ def containerIsRunning(container_name, timeout=365 * 24 * 60 * 60):
         return None
     except requests.exceptions.HTTPError as e:
         logger.debug("Server error attempting to call container: ",
-                      container_name)
+                     container_name)
         raise create_api_error_from_http_exception(e)
 
 
 def getContainerName(job):
     """Create a random string including the job name, and return it."""
     return '--'.join([str(job),
-                      base64.b64encode(os.urandom(9), b'-_').decode('utf-8')])\
-                      .replace("'", '').replace('"', '').replace('_', '')
+                      base64.b64encode(os.urandom(9), b'-_').decode('utf-8')]) \
+        .replace("'", '').replace('"', '').replace('_', '')
+
+
+def _multiplexed_response_stream_helper(response):
+    """
+    A generator of multiplexed data blocks coming from a response stream modified from:
+    https://github.com/docker/docker-py/blob/4.3.1-release/docker/api/client.py#L370
+
+    :param response: requests.Response
+    :return: a generator with tuples of (stream_type, data)
+    """
+    while True:
+        header = response.raw.read(8)
+        if not header:
+            break
+        # header is 8 bytes with format: {STREAM_TYPE, 0, 0, 0, SIZE1, SIZE2, SIZE3, SIZE4}
+        # protocol: https://docs.docker.com/engine/api/v1.24/#attach-to-a-container
+        stream_type, length = struct.unpack('>BxxxL', header)
+        if not length:
+            continue
+        data = response.raw.read(length)
+        if not data:
+            break
+        yield stream_type, data


Test patch:
diff --git a/src/toil/test/lib/dockerTest.py b/src/toil/test/lib/dockerTest.py
--- a/src/toil/test/lib/dockerTest.py
+++ b/src/toil/test/lib/dockerTest.py
@@ -1,4 +1,4 @@
-# Copyright (C) 2015-2018 Regents of the University of California
+# Copyright (C) 2015-2020 Regents of the University of California
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import absolute_import
 import logging
 import signal
 import time
@@ -286,6 +285,41 @@ def testNonCachingDockerChain(self):
     def testNonCachingDockerChainErrorDetection(self):
         self.testDockerPipeChainErrorDetection(disableCaching=False)
 
+    def testDockerLogs(self, stream=False, demux=False):
+        """Test for the different log outputs when deatch=False."""
+
+        working_dir = os.path.join(self.tempDir, 'working')
+        script_file = os.path.join(working_dir, 'script.sh')
+        os.makedirs(working_dir, exist_ok=True)
+
+        options = Job.Runner.getDefaultOptions(os.path.join(self.tempDir, 'jobstore'))
+        options.logLevel = self.dockerTestLogLevel
+        options.workDir = working_dir
+        options.clean = 'always'
+        A = Job.wrapJobFn(_testDockerLogsFn,
+                          working_dir=working_dir,
+                          script_file=script_file,
+                          stream=stream,
+                          demux=demux)
+
+        try:
+            rv = Job.Runner.startToil(A, options)
+            assert rv == True
+        finally:
+            try:
+                os.remove(script_file)
+            except:
+                pass
+
+    def testDockerLogs_Stream(self):
+        self.testDockerLogs(stream=True, demux=False)
+
+    def testDockerLogs_Demux(self):
+        self.testDockerLogs(stream=False, demux=True)
+
+    def testDockerLogs_Demux_Stream(self):
+        self.testDockerLogs(stream=True, demux=True)
+
 
 def _testDockerCleanFn(job,
                        working_dir,
@@ -338,7 +372,7 @@ def _testDockerPipeChainFn(job):
 
 def _testDockerPipeChainErrorFn(job):
     """Return True if the command exit 1 | wc -l raises a ContainerError."""
-    parameters = [ ['exit', '1'], ['wc', '-l'] ]
+    parameters = [['exit', '1'], ['wc', '-l']]
     try:
         apiDockerCall(job,
                       image='quay.io/ucsc_cgl/spooky_test',
@@ -346,3 +380,72 @@ def _testDockerPipeChainErrorFn(job):
     except ContainerError:
         return True
     return False
+
+
+def _testDockerLogsFn(job,
+                      working_dir,
+                      script_file,
+                      stream=False,
+                      demux=False):
+    """Return True if the test succeeds. Otherwise Exception is raised."""
+
+    # we write a script file because the redirection operator, '>&2', is wrapped
+    # in quotes when passed as parameters.
+    import textwrap
+    bash_script = textwrap.dedent('''
+    #!/bin/bash
+    echo hello stdout ;
+    echo hello stderr >&2 ;
+    echo hello stdout ;
+    echo hello stderr >&2 ;
+    echo hello stdout ;
+    echo hello stdout ;
+    ''')
+
+    with open(script_file, 'w') as file:
+        file.write(bash_script)
+
+    out = apiDockerCall(job,
+                        image='ubuntu:latest',
+                        working_dir=working_dir,
+                        parameters=[script_file],
+                        volumes={working_dir: {'bind': working_dir, 'mode': 'rw'}},
+                        entrypoint="/bin/bash",
+                        stdout=True,
+                        stderr=True,
+                        stream=stream,
+                        demux=demux)
+
+    # we check the output length because order is not guaranteed.
+    if stream:
+        if demux:
+            # a generator with tuples of (stdout, stderr)
+            assert hasattr(out, '__iter__')
+            for _ in range(6):
+                stdout, stderr = next(out)
+                if stdout:
+                    # len('hello stdout\n') == 13
+                    assert len(stdout) == 13
+                elif stderr:
+                    assert len(stderr) == 13
+                else:
+                    assert False
+        else:
+            # a generator with bytes
+            assert hasattr(out, '__iter__')
+            for _ in range(6):
+                assert len(next(out)) == 13
+    else:
+        if demux:
+            # a tuple of (stdout, stderr)
+            stdout, stderr = out
+            # len('hello stdout\n' * 4) == 52
+            assert len(stdout) == 52
+            # len('hello stderr\n' * 2) == 26
+            assert len(stderr) == 26
+        else:
+            # a bytes object
+            # len('hello stdout\n' * 4 + 'hello stderr\n' * 2) == 78
+            assert len(out) == 78
+
+    return True