Created
August 21, 2019 00:01
-
-
Save mlin/23ad39e9529d7cb6e274699b499caf98 to your computer and use it in GitHub Desktop.
swarmsub: run command in a docker image using swarm scheduler
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/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()) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment