diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index 9d95cc269..fda6cc10f 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -135,8 +135,10 @@ class ViewerConfig(PrintableConfig): """The zmq port to connect to for communication. If None, find an available port.""" launch_bridge_server: bool = True """whether or not to launch the bridge server""" - websocket_port: Optional[int] = 7007 - """the default websocket port to connect to""" + websocket_port: Optional[int] = None + """The websocket port to connect to. If None, find an available port.""" + websocket_port_default: int = 7007 + """The default websocket port to connect to if websocket_port is not specified""" ip_address: str = "127.0.0.1" """the ip address where the bridge server is running""" num_rays_per_chunk: int = 32768 diff --git a/nerfstudio/viewer/server/subprocess.py b/nerfstudio/viewer/server/subprocess.py index 3e07e2d6b..b50038561 100644 --- a/nerfstudio/viewer/server/subprocess.py +++ b/nerfstudio/viewer/server/subprocess.py @@ -33,6 +33,42 @@ from nerfstudio.viewer.server import server CONSOLE = Console() +def is_port_open(port: int): + """Returns True if the port is open. + + Args: + port: Port to check. + + Returns: + True if the port is open, False otherwise. + """ + try: + sock = socket.socket() + _ = sock.bind(("", port)) + sock.close() + return True + except OSError: + return False + + +def get_free_port(default_port: int = None): + """Returns a free port on the local machine. Try to use default_port is possible. + + Args: + default_port: Port to try to use. + + Returns: + A free port on the local machine. + """ + if default_port: + if is_port_open(default_port): + return default_port + sock = socket.socket() + sock.bind(("", 0)) + port = sock.getsockname()[1] + return port + + def run_viewer_bridge_server_as_subprocess( websocket_port: int, zmq_port: Optional[int] = None, @@ -54,9 +90,7 @@ def run_viewer_bridge_server_as_subprocess( # find an available port for zmq if zmq_port is None: - sock = socket.socket() - sock.bind(("", 0)) - zmq_port = sock.getsockname()[1] + zmq_port = get_free_port() string = f"Using ZMQ port: {zmq_port}" CONSOLE.print(f"[bold yellow]{string}") @@ -90,7 +124,6 @@ def run_viewer_bridge_server_as_subprocess( "You likely have to modify --viewer.zmq-port and/or --viewer.websocket-port in the " "config to avoid conflicting ports.\n" ) - string += "Try modifying --viewer.websocket-port 7007\n" CONSOLE.print(f"[bold red]{string}") cleanup(process) # This exists the entire program. sys.exit() will only kill the thread that this runs in. diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index e13084f8e..b781bea08 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -45,7 +45,10 @@ from nerfstudio.utils import colormaps, profiler, writer from nerfstudio.utils.decorators import check_main_thread, decorate_all from nerfstudio.utils.io import load_from_json, write_to_json from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter -from nerfstudio.viewer.server.subprocess import run_viewer_bridge_server_as_subprocess +from nerfstudio.viewer.server.subprocess import ( + get_free_port, + run_viewer_bridge_server_as_subprocess, +) from nerfstudio.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h from nerfstudio.viewer.server.visualizer import Viewer @@ -256,10 +259,13 @@ class ViewerState: self.datapath = datapath.parent if datapath.is_file() else datapath if self.config.launch_bridge_server: # start the viewer bridge server - assert self.config.websocket_port is not None + if self.config.websocket_port is None: + websocket_port = get_free_port(default_port=self.config.websocket_port_default) + else: + websocket_port = self.config.websocket_port self.log_filename.parent.mkdir(exist_ok=True) zmq_port = run_viewer_bridge_server_as_subprocess( - self.config.websocket_port, + websocket_port, zmq_port=self.config.zmq_port, ip_address=self.config.ip_address, log_filename=str(self.log_filename), @@ -267,7 +273,7 @@ class ViewerState: # TODO(ethan): log the output of the viewer bridge server in a file where the training logs go CONSOLE.line() version = get_viewer_version() - websocket_url = f"ws://localhost:{self.config.websocket_port}" + websocket_url = f"ws://localhost:{websocket_port}" self.viewer_url = f"https://viewer.nerf.studio/versions/{version}/?websocket_url={websocket_url}" CONSOLE.rule(characters="=") CONSOLE.print(f"[Public] Open the viewer at {self.viewer_url}")