Repository: DataBiosphere/toil

Issue:
Move Docker call functions in Toil-lib into Toil
See https://github.com/BD2KGenomics/toil-lib/blob/master/src/toil_lib/programs.py

Proposal is that docker_call function (and necessary functions) moves to src/toil/lib/docker.py




Discussion:
Note, this needs a synched PR to remove from toil-lib.
> Note, this needs a synched PR to remove from toil-lib.

and toil-scripts and toil-rnaseq
@jvivian do you want to take this on? I figure I would take on the testing.
Thanks @jvivian 

Resolving patch:
diff --git a/src/toil/lib/docker.py b/src/toil/lib/docker.py
new file mode 100644
--- /dev/null
+++ b/src/toil/lib/docker.py
@@ -0,0 +1,248 @@
+"""
+    Module for calling Docker. Assumes `docker` is on the PATH.
+
+    Contains two user-facing functions: dockerCall and dockerCheckOutput
+
+    Uses Toil's defer functionality to ensure containers are shutdown even in case of job or pipeline failure
+
+    Example of using dockerCall in a Toil pipeline to index a FASTA file with SAMtools:
+        def toil_job(job):
+            work_dir = job.fileStore.getLocalTempDir()
+            path = job.fileStore.readGlobalFile(ref_id, os.path.join(work_dir, 'ref.fasta')
+            parameters = ['faidx', path]
+            dockerCall(job, tool='quay.io/ucgc_cgl/samtools:latest', work_dir=work_dir, parameters=parameters)
+"""
+import base64
+import logging
+import subprocess
+
+import os
+from bd2k.util.exceptions import require
+
+_logger = logging.getLogger(__name__)
+
+
+def dockerCall(job,
+               tool,
+               parameters=None,
+               workDir=None,
+               dockerParameters=None,
+               outfile=None,
+               defer=None):
+    """
+    Throws CalledProcessorError if the Docker invocation returns a non-zero exit code
+    This function blocks until the subprocess call to Docker returns
+
+    :param toil.Job.job job: The Job instance for the calling function.
+    :param str tool: Name of the Docker image to be used (e.g. quay.io/ucsc_cgl/samtools:latest).
+    :param list[str] parameters: Command line arguments to be passed to the tool.
+    :param str workDir: Directory to mount into the container via `-v`. Destination convention is /data
+    :param list[str] dockerParameters: Parameters to pass to Docker. Default parameters are `--rm`,
+            `--log-driver none`, and the mountpoint `-v work_dir:/data` where /data is the destination convention.
+             These defaults are removed if docker_parmaters is passed, so be sure to pass them if they are desired.
+    :param file outfile: Pipe output of Docker call to file handle
+    :param int defer: What action should be taken on the container upon job completion?
+           FORGO (0) will leave the container untouched.
+           STOP (1) will attempt to stop the container with `docker stop` (useful for debugging).
+           RM (2) will stop the container and then forcefully remove it from the system
+           using `docker rm -f`. This is the default behavior if defer is set to None.
+    """
+    _docker(job, tool=tool, parameters=parameters, workDir=workDir, dockerParameters=dockerParameters,
+            outfile=outfile, checkOutput=False, defer=defer)
+
+
+def dockerCheckOutput(job,
+                      tool,
+                      parameters=None,
+                      workDir=None,
+                      dockerParameters=None,
+                      defer=None):
+    """
+    Returns the stdout from the Docker invocation (via subprocess.check_output)
+    Throws CalledProcessorError if the Docker invocation returns a non-zero exit code
+    This function blocks until the subprocess call to Docker returns
+
+    :param toil.Job.job job: The Job instance for the calling function.
+    :param str tool: Name of the Docker image to be used (e.g. quay.io/ucsc_cgl/samtools:latest).
+    :param list[str] parameters: Command line arguments to be passed to the tool.
+    :param str workDir: Directory to mount into the container via `-v`. Destination convention is /data
+    :param list[str] dockerParameters: Parameters to pass to Docker. Default parameters are `--rm`,
+            `--log-driver none`, and the mountpoint `-v work_dir:/data` where /data is the destination convention.
+             These defaults are removed if docker_parmaters is passed, so be sure to pass them if they are desired.
+    :param int defer: What action should be taken on the container upon job completion?
+           FORGO (0) will leave the container untouched.
+           STOP (1) will attempt to stop the container with `docker stop` (useful for debugging).
+           RM (2) will stop the container and then forcefully remove it from the system
+           using `docker rm -f`. This is the default behavior if defer is set to None.
+    :returns: Stdout from the docker call
+    :rtype: str
+    """
+    return _docker(job, tool=tool, parameters=parameters, workDir=workDir,
+                   dockerParameters=dockerParameters, checkOutput=True, defer=defer)
+
+
+def _docker(job,
+            tool,
+            parameters=None,
+            workDir=None,
+            dockerParameters=None,
+            outfile=None,
+            checkOutput=False,
+            defer=None):
+    """
+    :param toil.Job.job job: The Job instance for the calling function.
+    :param str tool: Name of the Docker image to be used (e.g. quay.io/ucsc_cgl/samtools).
+    :param list[str] parameters: Command line arguments to be passed to the tool.
+    :param str workDir: Directory to mount into the container via `-v`. Destination convention is /data
+    :param list[str] dockerParameters: Parameters to pass to Docker. Default parameters are `--rm`,
+            `--log-driver none`, and the mountpoint `-v work_dir:/data` where /data is the destination convention.
+             These defaults are removed if docker_parmaters is passed, so be sure to pass them if they are desired.
+    :param file outfile: Pipe output of Docker call to file handle
+    :param bool checkOutput: When True, this function returns docker's output.
+    :param int defer: What action should be taken on the container upon job completion?
+           FORGO (0) will leave the container untouched.
+           STOP (1) will attempt to stop the container with `docker stop` (useful for debugging).
+           RM (2) will stop the container and then forcefully remove it from the system
+           using `docker rm -f`. This is the default behavior if defer is set to None.
+    """
+    if parameters is None:
+        parameters = []
+    if workDir is None:
+        workDir = os.getcwd()
+
+    # Setup the outgoing subprocess call for docker
+    baseDockerCall = ['docker', 'run']
+    if dockerParameters:
+        baseDockerCall += dockerParameters
+    else:
+        baseDockerCall += ['--rm', '--log-driver', 'none', '-v',
+                           os.path.abspath(workDir) + ':/data']
+
+    # Ensure the user has passed a valid value for defer
+    require(defer in (None, FORGO, STOP, RM),
+            'Please provide a valid value for defer.')
+
+    # Get container name which is needed for _dockerKill
+    try:
+        if any('--name' in x for x in baseDockerCall):
+            if any('--name=' in x for x in baseDockerCall):
+                containerName = [x.split('=')[1] for x in baseDockerCall if '--name' in x][0]
+            else:
+                containerName = baseDockerCall[baseDockerCall.index('--name') + 1]
+        else:
+            containerName = _getContainerName(job)
+    except ValueError:
+        containerName = _getContainerName(job)
+        baseDockerCall.extend(['--name', containerName])
+    except IndexError:
+        raise RuntimeError("Couldn't parse Docker's `--name=` option, check parameters: " + str(dockerParameters))
+
+    # Defer the container on-exit action
+    if '--rm' in baseDockerCall and defer is None:
+        defer = RM
+    if '--rm' in baseDockerCall and defer is not RM:
+        _logger.warn('--rm being passed to docker call but defer not set to dockerCall.RM, defer set to: ' + str(defer))
+    job.defer(_dockerKill, containerName, action=defer)
+    # Defer the permission fixing function which will run after this job concludes.
+    # We call this explicitly later on in this function, but we defer it as well to handle unexpected job failure.
+    job.defer(_fixPermissions, tool, workDir)
+
+    # Make subprocess call
+    call = baseDockerCall + [tool] + parameters
+    _logger.info("Calling docker with " + repr(call))
+
+    if outfile:
+        subprocess.check_call(call, stdout=outfile)
+    else:
+        if checkOutput:
+            return subprocess.check_output(call)
+        else:
+            subprocess.check_call(call)
+
+
+FORGO = 0
+STOP = 1
+RM = 2
+
+
+def _dockerKill(containerName, action):
+    """
+    Kills the specified container.
+    :param str containerName: The name of the container created by docker_call
+    :param int action: What action should be taken on the container?  See `defer=` in
+           :func:`docker_call`
+    """
+    running = _containerIsRunning(containerName)
+    if running is None:
+        # This means that the container doesn't exist.  We will see this if the container was run
+        # with --rm and has already exited before this call.
+        _logger.info('The container with name "%s" appears to have already been removed.  Nothing to '
+                  'do.', containerName)
+    else:
+        if action in (None, FORGO):
+            _logger.info('The container with name %s continues to exist as we were asked to forgo a '
+                      'post-job action on it.', containerName)
+        else:
+            _logger.info('The container with name %s exists. Running user-specified defer functions.',
+                         containerName)
+            if running and action >= STOP:
+                _logger.info('Stopping container "%s".', containerName)
+                subprocess.check_call(['docker', 'stop', containerName])
+            else:
+                _logger.info('The container "%s" was not found to be running.', containerName)
+            if action >= RM:
+                # If the container was run with --rm, then stop will most likely remove the
+                # container.  We first check if it is running then remove it.
+                running = _containerIsRunning(containerName)
+                if running is not None:
+                    _logger.info('Removing container "%s".', containerName)
+                    subprocess.check_call(['docker', 'rm', '-f', containerName])
+                else:
+                    _logger.info('The container "%s" was not found on the system.  Nothing to remove.',
+                                 containerName)
+
+
+def _fixPermissions(tool, workDir):
+    """
+    Fix permission of a mounted Docker directory by reusing the tool to change ownership.
+    Docker natively runs as a root inside the container, and files written to the
+    mounted directory are implicitly owned by root.
+
+    :param list baseDockerCall: Docker run parameters
+    :param str tool: Name of tool
+    :param str workDir: Path of work directory to recursively chown
+    """
+    baseDockerCall = ['docker', 'run', '--log-driver=none',
+                      '-v', os.path.abspath(workDir) + ':/data', '--rm', '--entrypoint=chown']
+    stat = os.stat(workDir)
+    command = baseDockerCall + [tool] + ['-R', '{}:{}'.format(stat.st_uid, stat.st_gid), '/data']
+    subprocess.check_call(command)
+
+
+def _getContainerName(job):
+    return '--'.join([str(job),
+                      job.fileStore.jobID,
+                      base64.b64encode(os.urandom(9), '-_')]).replace("'", '').replace('_', '')
+
+
+def _containerIsRunning(container_name):
+    """
+    Checks whether the container is running or not.
+    :param container_name: Name of the container being checked.
+    :returns: True if running, False if not running, None if the container doesn't exist.
+    :rtype: bool
+    """
+    try:
+        output = subprocess.check_output(['docker', 'inspect', '--format', '{{.State.Running}}',
+                                          container_name]).strip()
+    except subprocess.CalledProcessError:
+        # This will be raised if the container didn't exist.
+        _logger.debug("'docker inspect' failed. Assuming container %s doesn't exist.", container_name,
+                   exc_info=True)
+        return None
+    if output == 'true':
+        return True
+    elif output == 'false':
+        return False
+    else:
+        raise RuntimeError("Got unexpected value for State.Running (%s)" % output)


Test patch:
diff --git a/src/toil/test/lib/dockerTest.py b/src/toil/test/lib/dockerTest.py
new file mode 100644
--- /dev/null
+++ b/src/toil/test/lib/dockerTest.py
@@ -0,0 +1,129 @@
+import logging
+import signal
+import time
+import uuid
+from threading import Thread
+
+import os
+from bd2k.util.files import mkdir_p
+from toil.job import Job
+from toil.leader import FailedJobsException
+from toil.lib.docker import dockerCall, _containerIsRunning, _dockerKill, STOP, FORGO, RM
+from toil.test import ToilTest
+
+_log = logging.getLogger(__name__)
+
+
+class DockerTest(ToilTest):
+    """
+    Tests dockerCall and ensures no containers are left around.
+
+
+    When running tests you may optionally set the TOIL_TEST_TEMP environment variable to the path
+    of a directory where you want temporary test files be placed. The directory will be created
+    if it doesn't exist. The path may be relative in which case it will be assumed to be relative
+    to the project root. If TOIL_TEST_TEMP is not defined, temporary files and directories will
+    be created in the system's default location for such files and any temporary files or
+    directories left over from tests will be removed automatically removed during tear down.
+    Otherwise, left-over files will not be removed.
+    """
+    def setUp(self):
+        self.tempDir = self._createTempDir(purpose='tempDir')
+
+    def testDockerClean(self):
+        """
+        Run the test container that creates a file in the work dir, and sleeps for 5 minutes.  Ensure
+        that the calling job gets SIGKILLed after a minute, leaving behind the spooky/ghost/zombie
+        container. Ensure that the container is killed on batch system shutdown (through the defer
+        mechanism).
+        This inherently also tests _docker
+        :returns: None
+        """
+        # We need to test the behaviour of `defer` with `rm` and `detached`. We do not look at the case
+        # where `rm` and `detached` are both True.  This is the truth table for the different
+        # combinations at the end of the test. R = Running, X = Does not exist, E = Exists but not
+        # running.
+        #              None     FORGO     STOP    RM
+        #    rm        X         R         X      X
+        # detached     R         R         E      X
+        #  Neither     R         R         E      X
+        assert os.getuid() != 0, "Cannot test this if the user is root."
+        data_dir = os.path.join(self.tempDir, 'data')
+        work_dir = os.path.join(self.tempDir, 'working')
+        test_file = os.path.join(data_dir, 'test.txt')
+        mkdir_p(data_dir)
+        mkdir_p(work_dir)
+        options = Job.Runner.getDefaultOptions(os.path.join(self.tempDir, 'jobstore'))
+        options.logLevel = 'INFO'
+        options.workDir = work_dir
+        options.clean = 'always'
+        for rm in (True, False):
+            for detached in (True, False):
+                if detached and rm:
+                    continue
+                for defer in (FORGO, STOP, RM, None):
+                    # Not using base64 logic here since it might create a name starting with a `-`.
+                    container_name = uuid.uuid4().hex
+                    print rm, detached, defer
+                    A = Job.wrapJobFn(_testDockerCleanFn, data_dir, detached, rm, defer,
+                                      container_name)
+                    try:
+                        Job.Runner.startToil(A, options)
+                    except FailedJobsException:
+                        # The file created by spooky_container would remain in the directory, and since
+                        # it was created inside the container, it would have had uid and gid == 0 (root)
+                        # upon creation. If the defer mechanism worked, it should now be non-zero and we
+                        # check for that.
+                        file_stats = os.stat(test_file)
+                        assert file_stats.st_gid != 0
+                        assert file_stats.st_uid != 0
+                        if (rm and defer != FORGO) or defer == RM:
+                            # These containers should not exist
+                            assert _containerIsRunning(container_name) is None, \
+                                'Container was not removed.'
+                        elif defer == STOP:
+                            # These containers should exist but be non-running
+                            assert _containerIsRunning(container_name) == False, \
+                                'Container was not stopped.'
+                        else:
+                            # These containers will be running
+                            assert _containerIsRunning(container_name) == True, \
+                                'Container was not running.'
+                    finally:
+                        # Prepare for the next test.
+                        _dockerKill(container_name, RM)
+                        os.remove(test_file)
+
+
+def _testDockerCleanFn(job, workDir, detached=None, rm=None, defer=None, containerName=None):
+    """
+    Test function for test docker_clean.  Runs a container with given flags and then dies leaving
+    behind a zombie container
+    :param toil.job.Job job: job
+    :param workDir: See `work_dir=` in :func:`dockerCall`
+    :param bool rm: See `rm=` in :func:`dockerCall`
+    :param bool detached: See `detached=` in :func:`dockerCall`
+    :param int defer: See `defer=` in :func:`dockerCall`
+    :param str containerName: See `container_name=` in :func:`dockerCall`
+    :return:
+    """
+    dockerParameters = ['--log-driver=none', '-v', os.path.abspath(workDir) + ':/data',
+                        '--name', containerName]
+    if detached:
+        dockerParameters.append('-d')
+    if rm:
+        dockerParameters.append('--rm')
+
+    def killSelf():
+        test_file = os.path.join(workDir, 'test.txt')
+        # This will kill the worker once we are sure the docker container started
+        while not os.path.exists(test_file):
+            _log.debug('Waiting on the file created by spooky_container.')
+            time.sleep(1)
+        # By the time we reach here, we are sure the container is running.
+        os.kill(os.getpid(), signal.SIGKILL)  # signal.SIGINT)
+    t = Thread(target=killSelf)
+    # Make it a daemon thread so that thread failure doesn't hang tests.
+    t.daemon = True
+    t.start()
+    dockerCall(job, tool='quay.io/ucsc_cgl/spooky_test', workDir=workDir, defer=defer, dockerParameters=dockerParameters)