docker CLI wrapper which adds the current cgroup as the --cgroup-parent

docker

#!/usr/bin/env python3

import os
import os.path
import sys


def get_memory_cgroup(data):
    for line in data.splitlines():
        if not line:
            continue
        hid, controllers, path = line.split(":", 2)
        if "memory" in controllers.split(","):
            return path


def read_current_process_cgroups():
    pid = os.getpid()
    return open(f"/proc/{pid}/cgroups").read()


def docker_env(env, script_path=None):
    if script_path is None:
        script_path = __file__

    # Need to remove our own path from env so we find the correct docker
    script_dir = os.path.dirname(os.path.abspath(script_path))
    path = [p for p in os.get_exec_path(env) if os.path.abspath(p) != script_dir]

    env = env.copy()
    env["PATH"] = os.pathsep.join(path)
    return env


def docker_add_cgroup_parent(args, cgroup_parent):
    # Find the subcommand
    for i, arg in enumerate(args):
        # Skip global flags
        if arg.startswith("-"):
            continue
        # We only add the argument to run and build
        if arg not in ("run", "build"):
            break
        return args[:i] + [arg, "--cgroup-parent", cgroup_parent] + args[i + 1 :]
    return args


def main():
    # Update PATH so we don't include this script when looking for docker
    # (preventing executing it again)
    env = docker_env(os.environ.copy())
    try:
        # Extract the cgroup for this process to pass to docker
        cgroup = get_memory_cgroup(read_current_process_cgroups())
        # include --cgroup-parent in args
        args = docker_add_cgroup_parent(sys.argv[1:], cgroup)
    except FileNotFoundError:
        if os.getenv("CI"):
            print("ERROR: Failed to find cgroup on CI", file=sys.stderr)
            sys.exit(1)
        args = sys.argv[1:]
    print("INFO: docker %s" % " ".join(args), file=sys.stderr)
    os.execvpe("docker", ["docker"] + args, env)


if __name__ == "__main__":
    main()

docker_test.py

import os.path

# Hack to load "docker" script since pytest can't discover it since it isn't a
# legit python script.
def load_source(name, path):
    import importlib.machinery
    import importlib.util

    loader = importlib.machinery.SourceFileLoader(name, path)
    spec = importlib.util.spec_from_loader(loader.name, loader)
    mod = importlib.util.module_from_spec(spec)
    loader.exec_module(mod)
    return mod


docker = load_source("docker_wrapper", os.path.dirname(__file__) + "/docker")


def test_get_memory_cgroup():
    data = """12:net_cls,net_prio:/kubepods/burstable/pod828291bc-1962-11ea-bc13-42010a800046/3b6bb8b8c5848a5894dfea613736249bcdd85a51421e36d3ce46bd7a140376ab
5:memory:/kubepods/burstable/pod828291bc-1962-11ea-bc13-42010a800046/3b6bb8b8c5848a5894dfea613736249bcdd85a51421e36d3ce46bd7a140376ab
"""
    assert (
        docker.get_memory_cgroup(data)
        == "/kubepods/burstable/pod828291bc-1962-11ea-bc13-42010a800046/3b6bb8b8c5848a5894dfea613736249bcdd85a51421e36d3ce46bd7a140376ab"
    )


def test_docker_env():
    env = {"PATH": "/builds/dev/ci/bin:/usr/bin"}
    assert docker.docker_env(env, script_path="/builds/dev/ci/bin/docker") == {
        "PATH": "/usr/bin"
    }


def test_docker_add_cgroup_parent():
    import shlex

    cases = {
        "run foo": "run --cgroup-parent cgroup_test foo",
        "-D run --rm foo": "-D run --cgroup-parent cgroup_test --rm foo",
        "rm foo": "rm foo",
        "build": "build --cgroup-parent cgroup_test",
    }
    for argv, want in cases.items():
        argv = shlex.split(argv)
        want = shlex.split(want)
        got = docker.docker_add_cgroup_parent(argv, "cgroup_test")
        assert got == want