Run command in a container

def container_run(platform: str,
                  nvidia_runtime: bool,
                  docker_registry: str,
                  shared_memory_size: str,
                  local_ccache_dir: str,
                  command: List[str],
                  cleanup: Cleanup,
                  environment: Dict[str, str],
                  dry_run: bool = False) -> int:
    """Run command in a container"""
    container_wait_s = 600
    #
    # Environment setup
    #
    environment.update({
        'CCACHE_MAXSIZE': '500G',
        'CCACHE_TEMPDIR': '/tmp/ccache',  # temp dir should be local and not shared
        'CCACHE_DIR': '/work/ccache',  # this path is inside the container as /work/ccache is
                                       # mounted
        'CCACHE_LOGFILE': '/tmp/ccache.log',  # a container-scoped log, useful for ccache
                                              # verification.
    })
    # These variables are passed to the container to the process tree killer can find runaway
    # process inside the container
    # https://wiki.jenkins.io/display/JENKINS/ProcessTreeKiller
    # https://github.com/jenkinsci/jenkins/blob/578d6bacb33a5e99f149de504c80275796f0b231/core/src/main/java/hudson/model/Run.java#L2393
    #
    jenkins_env_vars = ['BUILD_NUMBER', 'BUILD_ID', 'BUILD_TAG']
    environment.update({k: os.environ[k] for k in jenkins_env_vars if k in os.environ})
    environment.update({k: os.environ[k] for k in ['CCACHE_MAXSIZE'] if k in os.environ})

    tag = get_docker_tag(platform=platform, registry=docker_registry)
    mx_root = get_mxnet_root()
    local_build_folder = buildir()
    # We need to create it first, otherwise it will be created by the docker daemon with root only permissions
    os.makedirs(local_build_folder, exist_ok=True)
    os.makedirs(local_ccache_dir, exist_ok=True)
    logging.info("Using ccache directory: %s", local_ccache_dir)
    docker_client = docker.from_env()
    # Equivalent command
    docker_cmd_list = [
        get_docker_binary(nvidia_runtime),
        'run',
        "--cap-add",
        "SYS_PTRACE", # Required by ASAN
        '--rm',
        '--shm-size={}'.format(shared_memory_size),
        # mount mxnet root
        '-v', "{}:/work/mxnet".format(mx_root),
        # mount mxnet/build for storing build
        '-v', "{}:/work/build".format(local_build_folder),
        '-v', "{}:/work/ccache".format(local_ccache_dir),
        '-u', '{}:{}'.format(os.getuid(), os.getgid()),
        '-e', 'CCACHE_MAXSIZE={}'.format(environment['CCACHE_MAXSIZE']),
        # temp dir should be local and not shared
        '-e', 'CCACHE_TEMPDIR={}'.format(environment['CCACHE_TEMPDIR']),
        # this path is inside the container as /work/ccache is mounted
        '-e', "CCACHE_DIR={}".format(environment['CCACHE_DIR']),
        # a container-scoped log, useful for ccache verification.
        '-e', "CCACHE_LOGFILE={}".format(environment['CCACHE_LOGFILE']),
        '-ti',
        tag]
    docker_cmd_list.extend(command)
    docker_cmd = ' \\\n\t'.join(docker_cmd_list)
    logging.info("Running %s in container %s", command, tag)
    logging.info("Executing the equivalent of:\n%s\n", docker_cmd)
    # return code of the command inside docker
    ret = 0
    if not dry_run:
        #############################
        #
        signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGINT, signal.SIGTERM})
        # noinspection PyShadowingNames
        runtime = None
        if nvidia_runtime:
            # noinspection PyShadowingNames
            # runc is default (docker info | grep -i runtime)
            runtime = 'nvidia'
        container = docker_client.containers.run(
            tag,
            runtime=runtime,
            detach=True,
            command=command,
            shm_size=shared_memory_size,
            user='{}:{}'.format(os.getuid(), os.getgid()),
            cap_add='SYS_PTRACE',
            volumes={
                mx_root:
                    {'bind': '/work/mxnet', 'mode': 'rw'},
                local_build_folder:
                    {'bind': '/work/build', 'mode': 'rw'},
                local_ccache_dir:
                    {'bind': '/work/ccache', 'mode': 'rw'},
            },
            environment=environment)
        try:
            logging.info("Started container: %s", trim_container_id(container.id))
            # Race condition:
            # If the previous call is interrupted then it's possible that the container is not cleaned up
            # We avoid by masking the signals temporarily
            cleanup.add_container(container)
            signal.pthread_sigmask(signal.SIG_UNBLOCK, {signal.SIGINT, signal.SIGTERM})
            #
            #############################

            stream = container.logs(stream=True, stdout=True, stderr=True)
            sys.stdout.flush()
            for chunk in stream:
                sys.stdout.buffer.write(chunk)
                sys.stdout.buffer.flush()
            sys.stdout.flush()
            stream.close()
            try:
                logging.info("Waiting for status of container %s for %d s.",
                            trim_container_id(container.id),
                            container_wait_s)
                wait_result = container.wait(timeout=container_wait_s)
                logging.info("Container exit status: %s", wait_result)
                ret = wait_result.get('StatusCode', 200)
                if ret != 0:
                    logging.error("Container exited with an error 😞")
                else:
                    logging.info("Container exited with success 👍")
            except Exception as e:
                logging.exception(e)
                ret = 150

            # Stop
            try:
                logging.info("Stopping container: %s", trim_container_id(container.id))
                container.stop()
            except Exception as e:
                logging.exception(e)
                ret = 151

            # Remove
            try:
                logging.info("Removing container: %s", trim_container_id(container.id))
                container.remove()
            except Exception as e:
                logging.exception(e)
                ret = 152
            cleanup.remove_container(container)
            containers = docker_client.containers.list()
            if containers:
                logging.info("Other running containers: %s", [trim_container_id(x.id) for x in containers])
        except docker.errors.NotFound as e:
            logging.info("Container was stopped before cleanup started: %s", e)
    return ret