Spaces:
Build error
Build error
| #!/usr/bin/env python | |
| from __future__ import print_function | |
| ######### | |
| # Credit: https://github.com/RobotLocomotion/pytorch-dense-correspondence/blob/master/docker/docker_run.py | |
| ######### | |
| import argparse | |
| import os | |
| import socket | |
| import getpass | |
| import yaml | |
| if __name__=="__main__": | |
| user_name = getpass.getuser() | |
| default_image_name = user_name + '-cliport' | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("-i", "--image", type=str, | |
| help="(required) name of the image that this container is derived from", default=default_image_name) | |
| parser.add_argument("-nd", "--nvidia_docker", action='store_true', help="(optional) use nvidia-docker instead of docker") | |
| parser.add_argument("-c", "--container", type=str, default="cliport", help="(optional) name of the container") | |
| parser.add_argument("-d", "--data", type=str, default="data/", help="(optional) external data directory") | |
| parser.add_argument("-hl", "--headless", action='store_true', help="(optional) run in headless mode") | |
| parser.add_argument("-r", "--root", action='store_true', help="(optional) login as root instead of user") | |
| parser.add_argument("-g", "--gpus", type=str, default="all", help="(optional) gpus for nvidia docker") | |
| parser.add_argument("-dr", "--dry_run", action='store_true', help="(optional) perform a dry_run, print the command that would have been executed but don't execute it.") | |
| parser.add_argument("-p", "--passthrough", type=str, default="", help="(optional) extra string that will be tacked onto the docker run command, allows you to pass extra options. Make sure to put this in quotes and leave a space before the first character") | |
| args = parser.parse_args() | |
| print("running docker container derived from image %s" %args.image) | |
| source_dir = os.getcwd() | |
| image_name = args.image | |
| home_directory = '/home/' + user_name | |
| cmd = "" | |
| cmd += "xhost +local:root \n" if not args.headless else "" | |
| cmd += "docker run " | |
| if args.container: | |
| cmd += " --name %(container_name)s " % {'container_name': args.container} | |
| # gpus | |
| if args.nvidia_docker: | |
| cmd += "--gpus all " | |
| else: | |
| cmd += " --gpus %s" % (args.gpus) | |
| # display | |
| if args.headless: | |
| cmd += " -v /usr/bin/nvidia-xconfig:/usr/bin/nvidia-xconfig " | |
| else: # enable graphics | |
| cmd += " --env DISPLAY=unix$DISPLAY"\ | |
| " --env XAUTHORITY"\ | |
| " --env NVIDIA_DRIVER_CAPABILITIES=all"\ | |
| " --volume /tmp/.X11-unix:/tmp/.X11-unix"\ | |
| " --volume /dev/input:/dev/input" | |
| # bindings | |
| cmd += " -v %(source_dir)s:%(home_directory)s/cliport " \ | |
| % {'source_dir': source_dir, 'home_directory': home_directory} # mount source | |
| cmd += " -v ~/.ssh:%(home_directory)s/.ssh " % {'home_directory': home_directory} # mount ssh keys | |
| cmd += " -v ~/.torch:%(home_directory)s/.torch " % {'home_directory': home_directory} # mount torch folder | |
| cmd += " --user %s " % ("root" if args.root else user_name) # login | |
| # custom data path | |
| cmd += " -v %s:/data " %(os.path.join(source_dir, args.data)) | |
| # expose UDP ports | |
| cmd += " -p 8888:8888 " | |
| cmd += " --ipc=host " | |
| # share host machine network | |
| cmd += " --network=host " | |
| cmd += " " + args.passthrough + " " | |
| cmd += " --privileged" | |
| cmd += " --rm " # remove the image when you exit | |
| cmd += "-it " | |
| cmd += args.image | |
| cmd_endxhost = "xhost -local:root" | |
| print("command:\n", cmd) | |
| print("command = \n \n", cmd, "\n", cmd_endxhost) | |
| print("") | |
| # build the docker image | |
| if not args.dry_run: | |
| print("executing shell command") | |
| code = os.system(cmd) | |
| print("Executed with code ", code) | |
| if not args.headless: | |
| os.system(cmd_endxhost) | |
| # Squash return code to 0/1, as | |
| # Docker's very large return codes | |
| # were tricking Jenkins' failure | |
| # detection | |
| exit(code != 0) | |
| else: | |
| print("dry run, not executing command") | |
| exit(0) |