zmq_url to zmq_port (#563)

* zmq_url to port

* lazy commit

nerfactory/configs/base.py

@@ -169,7 +169,7 @@ class ViewerConfig(PrintableConfig):
     start_train: bool = True
     """whether to immediately start training upon loading viewer
     if False, will just visualize dataset but you can toggle training in viewer"""
-    zmq_url: str = "tcp://127.0.0.1:6000"
+    zmq_port: int = 6000
     """the zmq port to connect to for communication"""
     launch_bridge_server: bool = True
     """whether or not to launch the zmq bridge server"""

nerfactory/viewer/server/server.py

@@ -14,11 +14,10 @@
 
 """Server bridge to faciliate interactions between python backend and javascript front end"""
 
-
-import argparse
 import sys
 from typing import Callable, List, Optional, Tuple
 
+import dcargs
 import msgpack
 import msgpack_numpy
 import tornado.gen
@@ -177,6 +176,7 @@ class ZMQWebSocketBridge:
         websocket_address: str = "0.0.0.0",
         websocket_port: Optional[int] = None,
     ):
+        self.zmq_port = zmq_port
         self.host = host
         self.websocket_address = websocket_address
         self.websocket_pool = set()
@@ -185,14 +185,15 @@ class ZMQWebSocketBridge:
         self.pcs = set()
         self.video_tracks = set()
 
-        if zmq_port is None:
+        if self.zmq_port is None:
 
             def f(port):
                 return self.setup_zmq(f"{DEFAULT_ZMQ_METHOD:s}://{self.host:s}:{port:d}")
 
             (self.zmq_socket, self.zmq_stream, self.zmq_url), _ = find_available_port(f, DEFAULT_ZMQ_PORT)
+            self.zmq_port = int(self.zmq_url.split(":")[-1])
         else:
-            zmq_url = f"{DEFAULT_ZMQ_METHOD:s}://{self.host:s}:{zmq_port:d}"
+            zmq_url = f"{DEFAULT_ZMQ_METHOD:s}://{self.host:s}:{self.zmq_port:d}"
             self.zmq_socket, self.zmq_stream, self.zmq_url = self.setup_zmq(zmq_url)
 
         listen_kwargs = {"address": self.websocket_address}
@@ -208,7 +209,7 @@ class ZMQWebSocketBridge:
 
     def __str__(self) -> str:
         class_name = self.__class__.__name__
-        return f'{class_name} using zmq_url="{self.zmq_url}" and websocket_url="{self.websocket_url}"'
+        return f'{class_name} using zmq_port="{self.zmq_port}" and websocket_port="{self.websocket_port}"'
 
     def make_app(self):
         """Create a tornado application for the websocket server."""
@@ -292,13 +293,14 @@ class ZMQWebSocketBridge:
         self.ioloop.start()
 
 
-def run_viewer_bridge_server():
-    """Run the viewer bridge server"""
-    parser = argparse.ArgumentParser(description="Listen for ZeroMQ commands")
-    parser.add_argument("--zmq-port", "-z", type=int, nargs="?", default=6000)
-    parser.add_argument("--websocket-port", "-wp", type=int, nargs="?", default=7007)
-    args = parser.parse_args()
-    bridge = ZMQWebSocketBridge(zmq_port=args.zmq_port, websocket_port=args.websocket_port)
+def run_viewer_bridge_server(zmq_port: int = 6000, websocket_port: int = 7007):
+    """Run the viewer bridge server.
+
+    Args:
+        zmq_port: port to use for zmq
+        websocket_port: port to use for websocket
+    """
+    bridge = ZMQWebSocketBridge(zmq_port=zmq_port, websocket_port=websocket_port)
     print(bridge)
     try:
         bridge.run()
@@ -306,5 +308,11 @@ def run_viewer_bridge_server():
         pass
 
 
+def main():
+    """The main entrypoint."""
+    dcargs.extras.set_accent_color("bright_yellow")
+    dcargs.cli(run_viewer_bridge_server)
+
+
 if __name__ == "__main__":
-    run_viewer_bridge_server()
+    main()

nerfactory/viewer/server/subprocess.py

@@ -69,10 +69,10 @@ def run_viewer_bridge_server_as_subprocess(zmq_port: int, websocket_port: int, l
             time.sleep(0.5)
         string = f"\nThe viewer bridge server subprocess failed. Please check the log file {log_filename}.\n"
         string += (
-            "You likely have to modify --viewer.zmq-url and/or --viewer.websocket-port in the "
+            "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.zmq-url tcp://127.0.0.1:6000 --viewer.websocket-port 7007\n"
+        string += "Try modifying --viewer.zmq-port 6000 --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.

nerfactory/viewer/server/viewer_utils.py

@@ -209,10 +209,11 @@ class ViewerState:
         self.log_filename = log_filename
         if self.config.launch_bridge_server:
             # start the viewer bridge server
-            zmq_port = int(self.config.zmq_url.split(":")[-1])
             websocket_port = self.config.websocket_port
             self.log_filename.parent.mkdir(exist_ok=True)
-            run_viewer_bridge_server_as_subprocess(zmq_port, websocket_port, log_filename=str(self.log_filename))
+            run_viewer_bridge_server_as_subprocess(
+                self.config.zmq_port, websocket_port, log_filename=str(self.log_filename)
+            )
             # TODO(ethan): move this into the writer such that it's at the bottom
             # of the logging stack and easy to see and click
             # TODO(ethan): log the output of the viewer bridge server in a file where the training logs go
@@ -228,7 +229,7 @@ class ViewerState:
             console.print(dev_open_viewer_instructions_string)
             console.rule(characters="=")
             console.line()
-        self.vis = Viewer(zmq_url=self.config.zmq_url)
+        self.vis = Viewer(zmq_port=self.config.zmq_port)
 
         # viewer specific variables
         self.prev_camera_matrix = None

nerfactory/viewer/server/visualizer.py

@@ -33,10 +33,11 @@ class ViewerWindow:
 
     context = zmq.Context()
 
-    def __init__(self, zmq_url):
-        self.zmq_url = zmq_url
+    def __init__(self, zmq_port):
+        self.zmq_port = zmq_port
         self.client = self.context.socket(zmq.REQ)
-        self.client.connect(self.zmq_url)
+        zmq_url = f"tcp://127.0.0.1:{self.zmq_port}"
+        self.client.connect(zmq_url)
         self.assert_connected()
 
     def send(self, command):
@@ -90,15 +91,15 @@ class Viewer:
     """Viewer class for connecting to the bridge server.
 
     Args:
-        zmq_url: Where to connect with ZMQ.
+        zmq_port: Where to connect with ZMQ.
         window: An already existing ViewerWindow.
     """
 
-    def __init__(self, zmq_url: Optional[str] = None, window: Optional[ViewerWindow] = None):
-        if zmq_url is None and window is None:
-            raise ValueError("Must specify either zmq_url or window.")
+    def __init__(self, zmq_port: Optional[int] = None, window: Optional[ViewerWindow] = None):
+        if zmq_port is None and window is None:
+            raise ValueError("Must specify either zmq_port or window.")
         if window is None:
-            self.window = ViewerWindow(zmq_url=zmq_url)
+            self.window = ViewerWindow(zmq_port=zmq_port)
         else:
             self.window = window
         self.path = Path(())

pyproject.toml

@@ -74,7 +74,7 @@ docs = [
 ]
 
 [project.scripts]
-nerfactory-viewer-bridge-server = "nerfactory.viewer.server.server:run_viewer_bridge_server"
+nf-scripts-viewer-bridge-server = "nerfactory.viewer.server.server:main"
 
 [options]
 # equivalent to using --extra-index-url with pip, which is needed for specifying the CUDA version torch and torchvision

scripts/run_viewer_bridge_server.py

@@ -1,7 +0,0 @@
-#!/usr/bin/env python
-"""View Bridge Server"""
-
-from nerfactory.viewer.server.server import run_viewer_bridge_server
-
-if __name__ == "__main__":
-    run_viewer_bridge_server()