Add option to auto select a websocket port (#1628)

* option to auto select a websocket port

* keep default port

* add docs

* default port

* Update base_config.py

---------

Co-authored-by: Matthew Tancik <tancik@berkeley.edu>

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

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.

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}")