def spawn_worker(ip, num_containers, run_async): """Creates worker container(s) on the desired IP. Args: ip (str): IP of the machine to run the worker container. num_containers (int): Number of containers to be run. run_async (bool): Whether the spawning should happen synchronously or not. """ print(f"Spawning worker on: {ip}...") remote_image = f"{FLAGS.master}:{config.DOCKER_REGISTRY_PORT}/{config.DOCKER_IMAGE}" configure_worker_daemon(ip) cmds = ["docker stop $(docker ps -a -q)", f"docker pull {remote_image}"] cmds += [docker_run_cmd(ip, remote_image)] * num_containers nc = NetcatClient(ip, config.NETCAT_PORT) os_type = get_os_type(ip) if os_type == OSType.LINUX: nc.run_script("setup_gpu.sh") if run_async: nc.run_async(cmds) else: nc.run(cmds) print(f"Completed setup of {ip}!")
def configure_worker_daemon(ip): """Configures the Docker daemon to accept HTTP connections for using the local registry. Args: ip (str): IP of the worker. """ os_type = get_os_type(ip) os_paths = { OSType.MAC: "~/.docker/", OSType.WINDOWS: "$env:userprofile\.docker", OSType.LINUX: "/etc/docker/", } os_restarts = { OSType.MAC: [ """osascript -e 'quit app "Docker"'""", "open -a Docker", "until docker ps; sleep 2; done", ], OSType.WINDOWS: [ "net stop docker", "net stop com.docker.service", 'taskkill /IM "dockerd.exe" /F', 'taskkill /IM "Docker for Windows.exe" /F', "net start docker", "net start com.docker.service", '& "c:\\Program Files\\Docker\\Docker\\Docker for Windows.exe"', "while (!(docker ps)) { sleep 2 };", ], OSType.LINUX: ["systemctl restart docker"], } registry = f"{FLAGS.master}:{config.DOCKER_REGISTRY_PORT}" daemon_json = os.path.join(os_paths[os_type], config.DOCKER_DAEMON_JSON) nc = NetcatClient(ip, config.NETCAT_PORT) results = nc.run([f"cat {daemon_json}"]) try: relevant_part = r"\{[^\}]*\}" # extracts section inside braces m = re.search(relevant_part, results) daemon_config = json.loads(m.group(0)) except Exception: daemon_config = {} if "insecure-registries" in daemon_config: if registry in daemon_config["insecure-registries"]: return else: daemon_config["insecure-registries"] = [] daemon_config["insecure-registries"].append(registry) new_daemon_config = json.dumps(daemon_config) configure_cmds = [f"echo '{new_daemon_config}' > {daemon_json}"] configure_cmds += os_restarts[os_type] nc.run(configure_cmds)
def setup_local_gpu(): # Check if we are using Linux and we have an NVIDIA card, and we are not rendering in AWS if not FLAGS.project_root.startswith("s3://"): host_os = get_os_type(config.LOCALHOST) if host_os == OSType.LINUX and pyvidia.get_nvidia_device() is not None: gpu_script = os.path.join(dir_scripts, "render", "setup_gpu.sh") print(glog.green("Setting up GPU environment...")) run_command(f"/bin/bash {gpu_script}", run_silently=not FLAGS.verbose) else: print( glog.yellow( "We can only access an Nvidia GPU from a Linux host. Skipping Docker GPU setup" ))
def run_ui(client, docker_img): """Starts the UI. Args: client (DockerClient): Docker client configured to the host environment. docker_img (str): Name of the Docker image. """ if not FLAGS.verbose: print(glog.green("Initializing container"), end="") loading_context = RepeatedTimer(1, lambda: print(glog.green("."), end="")) host_os = get_os_type(config.LOCALHOST) # Setup steps for X11 forwarding vary slightly per the host operating system volumes = { "/var/run/docker.sock": { "bind": "/var/run/docker.sock", "mode": "ro" } } if host_os == OSType.MAC or host_os == OSType.LINUX: volumes.update( {"/tmp/.X11-unix": { "bind": "/tmp/.X11-unix", "mode": "ro" }}) if host_os == OSType.MAC or host_os == OSType.LINUX: run_command(f"xhost + {config.LOCALHOST}", run_silently=not FLAGS.verbose) if host_os == OSType.LINUX: run_command(f"xhost + {config.DOCKER_LOCAL_HOSTNAME}", run_silently=not FLAGS.verbose) host_to_docker_path = {FLAGS.project_root: config.DOCKER_INPUT_ROOT} project = Project( FLAGS.project_root, FLAGS.cache, FLAGS.csv_path, FLAGS.s3_sample_frame, FLAGS.s3_ignore_fullsize_color, FLAGS.verbose, ) project.verify() cmds = [ "cd scripts/ui", f"""python3 -u dep.py \ --host_os={get_os_type(config.LOCALHOST)} \ --local_bin={FLAGS.local_bin} \ --master={FLAGS.master} \ --password={FLAGS.password} \ --project_root={FLAGS.project_root} \ --s3_ignore_fullsize_color={FLAGS.s3_ignore_fullsize_color} \ --s3_sample_frame={FLAGS.s3_sample_frame} \ --username={FLAGS.username} \ --verbose={FLAGS.verbose}""", ] docker_networks = client.networks.list() network_names = [docker_network.name for docker_network in docker_networks] if config.DOCKER_NETWORK not in network_names: client.networks.create(config.DOCKER_NETWORK, driver="bridge") project_address = Address(FLAGS.project_root) project_protocol = project_address.protocol if project_protocol == "smb": mounts = docker_mounts(FLAGS.project_root, host_to_docker_path, FLAGS.username, FLAGS.password) cmds = [f"mkdir {config.DOCKER_INPUT_ROOT}"] + mounts + cmds local_project_root = None elif project_protocol == "s3": glog.check_ne(FLAGS.csv_path, "", "csv_path cannot be empty if rendering on AWS") aws_util = AWSUtil(FLAGS.csv_path, s3_url=FLAGS.project_root) glog.check( aws_util.s3_bucket_is_valid(FLAGS.project_root), f"Invalid S3 project path: {FLAGS.project_root}", ) volumes.update({ FLAGS.csv_path: { "bind": config.DOCKER_AWS_CREDENTIALS, "mode": "rw" } }) project_name = project_address.path cache_path = os.path.join(os.path.expanduser(FLAGS.cache), project_name) os.makedirs(cache_path, exist_ok=True) volumes.update( {cache_path: { "bind": config.DOCKER_INPUT_ROOT, "mode": "rw" }}) local_project_root = cache_path else: glog.check( os.path.isdir(FLAGS.project_root), f"Invalid project path: {FLAGS.project_root}", ) volumes.update({ host_path: { "bind": docker_path, "mode": "rw" } for host_path, docker_path in host_to_docker_path.items() }) local_project_root = FLAGS.project_root ipc_dir = os.path.join(local_project_root, "ipc") os.makedirs(ipc_dir, exist_ok=True) volumes.update({ipc_dir: {"bind": config.DOCKER_IPC_ROOT, "mode": "rw"}}) cmd = f'/bin/bash -c "{" && ".join(cmds)}"' global container_name display = ":0" if host_os == OSType.LINUX else "host.docker.internal:0" runtime = "nvidia" if which("nvidia-docker") else "" if host_os != OSType.LINUX: display = "host.docker.internal:0" if not FLAGS.verbose: loading_context.stop() print("") try: container = client.containers.run( docker_img, command=cmd, detach=True, environment={"DISPLAY": display}, runtime=runtime, network=config.DOCKER_NETWORK, ports={ config.RABBITMQ_PORT: config.RABBITMQ_PORT, config.RABBITMQ_MANAGE_PORT: config.RABBITMQ_MANAGE_PORT, }, privileged=True, volumes=volumes, stderr=True, ) except docker.errors.APIError as e: if "port is already allocated" in str(e): raise Exception( "Failed to launch UI! Ensure: \n" "(1) No other instance of UI is running (check: docker ps) and\n" "(2) RabbitMQ is not running on your machine (check: ps aux | grep 'rabbitmq')" ) from None raise e container_name = container.name create_viewer_watchdog(client, ipc_dir, local_project_root)
def on_modified(self, event): """When a viewer file is modified from the UI, the appropriate viewer runs on the host. Args: event (watchdog.FileSystemEvent): Watchdog event for when viewer file has been modified. """ if isinstance(event, DirModifiedEvent): return ipc_name = os.path.basename(event.src_path) host_os = get_os_type(config.LOCALHOST) if ipc_name == config.DOCKER_RIFT_VIEWER_IPC and host_os != OSType.WINDOWS: print(glog.yellow("RiftViewer is only supported on Windows!")) return app_name = config.get_app_name(ipc_name) if not app_name: print(glog.red(f"Invalid IPC name: {ipc_name}")) return try: output_dir = posixpath.join(self.local_project_root, config.OUTPUT_ROOT_NAME) if ipc_name == config.DOCKER_RIFT_VIEWER_IPC: fused_dir = posixpath.join(output_dir, image_type_paths["fused"]) fused_json = self.get_fused_json(fused_dir) if not fused_json: print( glog.red(f"Cannot find fused rig json in {fused_dir}")) return cmd_flags = f"""--rig={posixpath.join(fused_dir, fused_json)} \ --catalog={posixpath.join(fused_dir, "fused.json")} \ --strip_files={posixpath.join(fused_dir, "fused_0.bin")}""" elif ipc_name in [ config.DOCKER_SMR_IPC, config.DOCKER_SMR_ONSCREEN_IPC ]: flags_render = self.get_render_flags("export") if ipc_name == config.DOCKER_SMR_IPC: flags_render["output"] = posixpath.join( output_dir, image_type_paths["exports"]) flags_smr = [flag["name"] for flag in bin_to_flags[app_name]] cmd_flags = "" ignore_onscreen = ["format", "output"] for flag in flags_render: if flag in flags_smr: if (flag in ignore_onscreen and ipc_name == config.DOCKER_SMR_ONSCREEN_IPC): continue cmd_flags += f" --{flag}={flags_render[flag]}" cmd_flags = cmd_flags.replace(config.DOCKER_INPUT_ROOT, self.local_project_root) cmd = f"{posixpath.join(FLAGS.local_bin, app_name)} {cmd_flags}" if os.name != "nt": # GLOG initiatives don't work in Powershell/cmd cmd = f"GLOG_alsologtostderr=1 GLOG_stderrthreshold=0 {cmd}" run_command(cmd) except Exception as e: print(glog.red(e))