swarmsub: run command in a docker image using swarm scheduler

swarmsub.py

#!/usr/bin/env python3

import sys
import time
import docker
import multiprocessing
from argparse import ArgumentParser, REMAINDER

def swarmsub(image, command=None, cpu=1, mounts=None):
    client = docker.from_env()
    svc = None
    try:
        # launch swarm "service"
        svc = client.services.create(
            image,
            command,
            # restart_policy 'none' so that swarm runs the container just once
            restart_policy=docker.types.RestartPolicy("none"),
            resources=docker.types.Resources(
                # cpu_limit throttles container to desired # of cpus.
                # the unit expected by swarm is "NanoCPUs"
                cpu_limit=cpu*1_000_000_000,
                # cpu_reservation makes swarm delay starting the container until the desired # of
                # cpus are available (considering other running services)
                cpu_reservation=cpu*1_000_000_000,
            ),
            mounts=mounts,
        )

        # poll for completion of the "service" consisting of one "task"
        exit_code = None
        while exit_code is None:
            time.sleep(1)
            svc.reload()
            tasks = svc.tasks()
            assert len(tasks) == 1
            if tasks[0]["Status"]["State"] in ["complete", "failed"]:
                exit_code = tasks[0]["Status"]["ContainerStatus"]["ExitCode"]
                assert isinstance(exit_code, int)

        return exit_code
    finally:
        if svc:
            # clean up the service which otherwise sticks around e.g. in 'docker service ls'
            svc.remove()
        client.close()

def main():
    parser = ArgumentParser()
    parser.add_argument("-v", "--volume", dest="volume", metavar="LIST",
                        help="bind mounts, e.g. /tmp:/tmp,/host_path:/container_path:ro")
    parser.add_argument("-c", "--cpu", dest="cpu", metavar="N", type=int, default=1,
                        help="number of CPUs to reserve & limit to, default %(default)s")
    parser.add_argument("image", metavar="IMAGE", help="docker image")
    parser.add_argument("command", metavar="...", nargs=REMAINDER, default=[], help="command to run in container")
    args = parser.parse_args(sys.argv[1:])

    kwargs = {}
    if args.command:
        kwargs["command"] = args.command
    if args.volume:
        kwargs["mounts"] = args.volume.split(",")
    if args.cpu:
        # swarm won't start the container if we let cpu_reservation > multiprocessing.cpu_count()
        kwargs["cpu"] = min(args.cpu, multiprocessing.cpu_count())

    return swarmsub(args.image, **kwargs)

if __name__ == "__main__":
    sys.exit(main())