ezmsg 3.6.0__tar.gz → 3.6.2__tar.gz

{ezmsg-3.6.0 → ezmsg-3.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: ezmsg
-Version: 3.6.0
+Version: 3.6.2
 Summary: A simple DAG-based computation model
 Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>
 License: MIT License
@@ -12,17 +12,20 @@ License: MIT License
         The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
         
         THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+License-File: LICENSE
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.8
 Requires-Dist: typing-extensions>=4.9.0
+Provides-Extra: axisarray
+Requires-Dist: numpy>=1.24.4; extra == 'axisarray'
 Provides-Extra: docs
 Requires-Dist: ezmsg-sigproc>=1.2.3; extra == 'docs'
 Requires-Dist: sphinx-rtd-theme==2.0.0; extra == 'docs'
 Requires-Dist: sphinx<=7.2; extra == 'docs'
 Provides-Extra: test
 Requires-Dist: flake8>=5.0.4; extra == 'test'
-Requires-Dist: numpy>=1.24.2; extra == 'test'
+Requires-Dist: numpy>=1.24.4; extra == 'test'
 Requires-Dist: pytest-asyncio>=0.23.8; extra == 'test'
 Requires-Dist: pytest-cov>=5.0.0; extra == 'test'
 Requires-Dist: pytest>=7.0.0; extra == 'test'

{ezmsg-3.6.0 → ezmsg-3.6.2}/docs/source/other.rst

@@ -14,4 +14,10 @@ Still under construction!
 Notebooks
 ---------
 
-An ezmsg pipeline can be run in a notebook environment with an extra parameter. When using ``run()``, include the optional kwarg ``force_single_process=True``.
+An ezmsg pipeline can be run in a notebook environment with an extra parameter. When using ``run()``, include the optional kwarg ``force_single_process=True``.
+
+Media
+-----
+
+* `BCI Society 2025 Poster <https://github.com/ezmsg-org/ezmsg/releases/download/V3.6.1/ezmsg_poster_BCI2025.pdf>`_
+* `Society For Neuroscience 2024 Poster <https://github.com/ezmsg-org/ezmsg/releases/download/V3.6.1/ezmsg_poster_SfN2024.pdf>`_

{ezmsg-3.6.0 → ezmsg-3.6.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ezmsg"
-version = "3.6.0"
+version = "3.6.2"
 description = "A simple DAG-based computation model"
 authors = [
   { name = "Griffin Milsap", email = "griffin.milsap@gmail.com" },
@@ -14,18 +14,23 @@ classifiers = [
   "Operating System :: OS Independent",
 ]
 requires-python = ">=3.8"
-dependencies = ["typing-extensions >= 4.9.0"]
+dependencies = [
+    "typing-extensions >= 4.9.0",
+]
 
 [project.optional-dependencies]
 test = [
   "pytest>=7.0.0",
-  "numpy>=1.24.2",
   "pytest-asyncio>=0.23.8",
   "pytest-cov>=5.0.0",
   "flake8>=5.0.4",
   "xarray>=2023.1.0;python_version<'3.13'",
+  "numpy>=1.24.4",
 ]
 docs = ["sphinx<=7.2", "sphinx-rtd-theme==2.0.0", "ezmsg-sigproc>=1.2.3"]
+axisarray = [
+    "numpy>=1.24.4",
+]
 
 
 [project.scripts]

ezmsg-3.6.2/src/ezmsg/core/command.py

@@ -0,0 +1,178 @@
+import argparse
+import asyncio
+import base64
+import json
+import logging
+import os
+import subprocess
+import sys
+import typing
+import webbrowser
+import zlib
+
+from .graphserver import GraphService
+from .shmserver import SHMService
+from .netprotocol import (
+    Address,
+    GRAPHSERVER_ADDR_ENV,
+    GRAPHSERVER_PORT_DEFAULT,
+    SHMSERVER_ADDR_ENV,
+    SHMSERVER_PORT_DEFAULT,
+    PUBLISHER_START_PORT_ENV,
+    PUBLISHER_START_PORT_DEFAULT,
+    close_stream_writer,
+)
+
+logger = logging.getLogger("ezmsg")
+
+
+def cmdline() -> None:
+    parser = argparse.ArgumentParser(
+        "ezmsg.core",
+        description="start and stop core ezmsg server processes",
+        epilog=f"""
+            You can also change server configuration with environment variables.
+            GraphServer will be hosted on ${GRAPHSERVER_ADDR_ENV} (default port: {GRAPHSERVER_PORT_DEFAULT}).  
+            SHMServer will be hosted on ${SHMSERVER_ADDR_ENV} (default port: {SHMSERVER_PORT_DEFAULT}).
+            Publishers will be assigned available ports starting from {PUBLISHER_START_PORT_DEFAULT}. (Change with ${PUBLISHER_START_PORT_ENV})
+        """,
+    )
+
+    parser.add_argument(
+        "command",
+        help="command for ezmsg",
+        choices=["serve", "start", "shutdown", "graphviz", "mermaid"],
+    )
+
+    parser.add_argument("--address", help="Address for GraphServer", default=None)
+
+    parser.add_argument(
+        "--target",
+        help="Target for mermaid output. Options are 'ink', 'live', and 'play'.",
+        default="live",
+    )
+
+    class Args:
+        command: str
+        address: typing.Optional[str]
+        target: str
+
+    args = parser.parse_args(namespace=Args)
+
+    graph_address = Address("127.0.0.1", GRAPHSERVER_PORT_DEFAULT)
+    if args.address is not None:
+        graph_address = Address.from_string(args.address)
+    shm_address_str = os.environ.get(
+        SHMSERVER_ADDR_ENV, f"127.0.0.1:{SHMSERVER_PORT_DEFAULT}"
+    )
+    shm_address = Address.from_string(shm_address_str)
+
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    loop.run_until_complete(
+        run_command(args.command, graph_address, shm_address, args.target)
+    )
+
+
+async def run_command(
+    cmd: str, graph_address: Address, shm_address: Address, target: str = "live"
+) -> None:
+    shm_service = SHMService(shm_address)
+    graph_service = GraphService(graph_address)
+
+    if cmd == "serve":
+        logger.info(f"GraphServer Address: {graph_address}")
+        logger.info(f"SHMServer Address: {shm_address}")
+
+        shm_server = shm_service.create_server()
+        graph_server = graph_service.create_server()
+
+        try:
+            logger.info("Servers running...")
+            graph_server.join()
+
+        except KeyboardInterrupt:
+            logger.info("Interrupt detected; shutting down servers")
+
+        finally:
+            if graph_server is not None:
+                graph_server.stop()
+
+            if shm_server is not None:
+                shm_server.stop()
+
+    elif cmd == "start":
+        popen = subprocess.Popen(
+            [sys.executable, "-m", "ezmsg.core", "serve", f"--address={graph_address}"]
+        )
+
+        while True:
+            try:
+                _, writer = await graph_service.open_connection()
+                await close_stream_writer(writer)
+                _, writer = await shm_service.open_connection()
+                await close_stream_writer(writer)
+                break
+            except ConnectionRefusedError:
+                await asyncio.sleep(0.1)
+
+        logger.info(f"Forked ezmsg servers in PID: {popen.pid}")
+
+    elif cmd == "shutdown":
+        try:
+            await graph_service.shutdown()
+            logger.info(
+                f"Issued shutdown command to GraphServer @ {graph_service.address}"
+            )
+
+        except ConnectionRefusedError:
+            logger.warning(
+                f"Could not issue shutdown command to GraphServer @ {graph_service.address}; server not running?"
+            )
+
+    elif cmd in ["graphviz", "mermaid"]:
+        graph_out = await graph_service.get_formatted_graph(cmd)
+        print(graph_out)
+        if cmd == "mermaid" and target == "live":
+            print(
+                "%% If the graph does not render immediately, try toggling the 'Pan & Zoom' button."
+            )
+
+        if cmd == "mermaid":
+            webbrowser.open(mm(graph_out, target=target))
+
+
+def mm(graph: str, target="live") -> str:
+    if target != "ink":
+        jdict = {
+            "code": graph,
+            "mermaid": {"theme": "default"},
+            "updateDiagram": True,
+            "autoSync": True,
+            "rough": False,
+        }
+        graph = json.dumps(jdict)
+    graphbytes: bytes = graph.encode("utf8")
+
+    if target != "ink":
+        compress = zlib.compressobj(9, zlib.DEFLATED, 15, 8, zlib.Z_DEFAULT_STRATEGY)
+        graphbytes = compress.compress(graphbytes)
+        graphbytes += compress.flush()
+
+    base64_bytes = base64.b64encode(graphbytes)
+    base64_string = base64_bytes.decode("ascii")
+
+    if target == "ink":
+        prefix = "https://mermaid.ink/img/"
+    elif target in ["live", "play"]:
+        type_str = "pako"  # or "base64" if we skip compression above.
+        if target == "live":
+            prefix = f"https://mermaid.live/edit#{type_str}:"
+        else:  # "play"
+            prefix = f"https://www.mermaidchart.com/play#{type_str}:"
+    else:
+        raise ValueError(
+            f"Unknown mermaid target '{target}'. Available options are 'ink', 'live', or 'play'."
+        )
+    return prefix + base64_string

{ezmsg-3.6.0 → ezmsg-3.6.2}/src/ezmsg/core/graphserver.py

@@ -1,7 +1,10 @@
 import asyncio
+from collections import defaultdict
 import logging
 import pickle
+from textwrap import indent
 import typing
+from uuid import uuid4
 
 from contextlib import suppress
 
@@ -28,6 +31,8 @@ from .netprotocol import (
 
 logger = logging.getLogger("ezmsg")
 
+IND = "  "
+
 
 class GraphServer(ThreadedAsyncServer):
     """
@@ -309,3 +314,131 @@ class GraphService(ServiceManager[GraphServer]):
         await asyncio.wait_for(reader.read(1), timeout=timeout)  # Complete
         await close_stream_writer(writer)
         return dag
+
+    async def get_pruned_graph(self) -> typing.Optional[defaultdict[str, set[str]]]:
+        try:
+            dag: DAG = await self.dag()
+        except (ConnectionRefusedError, ConnectionResetError):
+            logger.info(
+                f"GraphServer not running @{self.address}, or host is refusing connections"
+            )
+            return None
+
+        graph_connections = dag.graph.copy()
+        # Let's eliminate proxy topics, i.e. connections with inputs and outputs.
+        source_nodes = []
+        for node, conns in graph_connections.items():
+            if len(conns) > 0:
+                source_nodes += [node]
+        proxy_topics = []
+        for conns in graph_connections.values():
+            for conn in conns:
+                if conn in source_nodes and conn not in proxy_topics:
+                    proxy_topics += [conn]
+        # Replace Proxy Topics with actual source and downstream
+        for proxy_topic in proxy_topics:
+            downstreams = graph_connections.pop(proxy_topic)
+            logger.info(f"{proxy_topic} downstream connections: {downstreams}")
+            for node, conns in graph_connections.items():
+                for conn in conns:
+                    if conn == proxy_topic:
+                        new_conns = conns.copy()
+                        new_conns.remove(proxy_topic)
+                        new_conns.union(downstreams)
+                        logger.info(
+                            f"Updating connections for {node} from {conns} to {new_conns}"
+                        )
+                        graph_connections[node] = new_conns
+
+        return dag.graph.copy()
+
+    async def get_formatted_graph(self, fmt: str, direction: str = "LR") -> str:
+        if fmt not in ["mermaid", "graphviz"]:
+            raise ValueError(
+                f"Invalid format '{fmt}'. Options are 'mermaid' or 'graphviz'"
+            )
+        graph_connections = await self.get_pruned_graph()
+
+        if graph_connections is None or not graph_connections:
+            return ""
+
+        # Let's come up with UUID node names
+        nodes = set(graph_connections.keys())
+        if fmt == "mermaid":
+            node_map = {name: f"{str(uuid4())}" for name in nodes}
+        else:
+            node_map = {name: f'"{str(uuid4())}"' for name in nodes}  # graphviz
+
+        # Construct the graph
+        def tree():
+            return defaultdict(tree)
+
+        graph: defaultdict = tree()
+
+        connections = ""
+        for node, conns in graph_connections.items():
+            subgraph = graph
+            path = node.split("/")
+            route = path[:-1]
+            stream = path[-1]
+            for seg in route:
+                subgraph = subgraph[seg]
+            subgraph[stream] = node
+
+            for sub in conns:
+                if fmt == "mermaid":
+                    connections += f"{node_map[node]} --> {node_map[sub]}" + "\n"
+                else:
+                    connections += f"{node_map[node]} -> {node_map[sub]};" + "\n"
+
+        if fmt == "graphviz":
+            header = [
+                "digraph EZ {",
+                indent(f'rankdir="{direction}"', IND),
+            ]
+            footer = ["}"]
+
+            def per_leaf(g, leaf):
+                out = None
+                if isinstance(g[leaf], defaultdict):
+                    out = [
+                        f"subgraph {leaf.lower()} {{",
+                        indent("cluster = true;", IND),
+                        indent(f'label = "{leaf}";', IND),
+                        f"{recurse_graph(g[leaf])}}};",
+                    ]
+                elif isinstance(g[leaf], str):
+                    out = [f"{node_map[g[leaf]]} [label={leaf}];"]
+                return out
+
+        else:  # fmt == mermaid
+            header = [f"flowchart {direction}"]
+            footer = []
+
+            def per_leaf(g, leaf):
+                out = None
+                if isinstance(g[leaf], defaultdict):
+                    out = [
+                        f"subgraph {leaf.lower()} [{leaf}]",
+                        # "direction LR",
+                        f"{recurse_graph(g[leaf])}",
+                        "end",
+                    ]
+                elif isinstance(g[leaf], str):
+                    out = [f"{node_map[g[leaf]]}[{leaf}]"]
+                return out
+
+        def recurse_graph(g: defaultdict):
+            out = ""
+            for leaf in g:
+                leaf_list = per_leaf(g, leaf)
+                if leaf_list is not None:
+                    leaf_list += [""]  # Append a newline
+                    out += indent("\n".join(leaf_list), IND)
+            return out[:-1]
+
+        graph_out = "\n".join(
+            header + [recurse_graph(graph), indent(connections, IND)] + footer
+        )
+
+        return graph_out

{ezmsg-3.6.0 → ezmsg-3.6.2}/src/ezmsg/core/netprotocol.py

@@ -177,6 +177,7 @@ def create_socket(
     ignore_ports: typing.List[int] = RESERVED_PORTS,
 ) -> socket.socket:
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
 
     if host is None:
         host = DEFAULT_HOST

{ezmsg-3.6.0 → ezmsg-3.6.2}/src/ezmsg/util/messagelogger.py

@@ -1,5 +1,6 @@
 import json
 import time
+import typing
 
 from io import TextIOWrapper
 from dataclasses import field
@@ -8,11 +9,10 @@ from pathlib import Path
 import ezmsg.core as ez
 
 from .messagecodec import MessageEncoder, LogStart
+from .rate import Rate
 
-from typing import Optional, Any, Dict, AsyncGenerator
 
-
-def log_object(obj: Any) -> str:
+def log_object(obj: typing.Any) -> str:
     return json.dumps({"ts": time.time(), "obj": obj}, cls=MessageEncoder)
 
 
@@ -23,13 +23,18 @@ class MessageLoggerSettings(ez.Settings):
     Args:
         output: :py:class:`pathlib.Path` for a file where the messages will be logged.
             If the file path already exists, the existing file will be truncated to 0 length.
+        write_period: Period in seconds for performing a write to disk.
+            If <=0, write incoming messages to disk immediately, applying backpressure as needed.
+            If >0, messages are buffered and periodically written to disk in a separate task
     """
 
-    output: Optional[Path] = None
+    output: typing.Optional[Path] = None
+    write_period: float = 0.0 # sec
 
 
 class MessageLoggerState(ez.State):
-    output_files: Dict[Path, TextIOWrapper] = field(default_factory=dict)
+    output_files: typing.Dict[Path, TextIOWrapper] = field(default_factory=dict)
+    write_buffer: typing.List[str] = field(default_factory=list)
 
 
 class MessageLogger(ez.Unit):
@@ -55,10 +60,10 @@ class MessageLogger(ez.Unit):
     to stop logging messages to that path.
     """
 
-    INPUT_MESSAGE = ez.InputStream(Any)
+    INPUT_MESSAGE = ez.InputStream(typing.Any)
     """Pass a piece of data to log it to every open file which the ``MessageLogger`` is using."""
 
-    OUTPUT_MESSAGE = ez.OutputStream(Any)
+    OUTPUT_MESSAGE = ez.OutputStream(typing.Any)
     """Messages which are sent to ``INPUT_MESSAGE`` will pass through and be published on ``OUTPUT_MESSAGE``."""
 
     OUTPUT_START = ez.OutputStream(Path)
@@ -69,7 +74,7 @@ class MessageLogger(ez.Unit):
     """If a file passed to ``INPUT_STOP`` is successfully closed, its path will be published to
     ``OUTPUT_STOP``, otherwise ``None``."""
 
-    def open_file(self, filepath: Path) -> Optional[Path]:
+    def open_file(self, filepath: Path) -> typing.Optional[Path]:
         """Returns file path if file successfully opened, otherwise None"""
         if filepath in self.STATE.output_files:
             # If the file is already open, we return None
@@ -85,7 +90,7 @@ class MessageLogger(ez.Unit):
 
         return filepath
 
-    def close_file(self, filepath: Path) -> Optional[Path]:
+    def close_file(self, filepath: Path) -> typing.Optional[Path]:
         """Returns file path if file successfully closed, otherwise None"""
         if filepath not in self.STATE.output_files:
             # We haven't opened this file
@@ -103,28 +108,53 @@ class MessageLogger(ez.Unit):
 
     @ez.subscriber(INPUT_START)
     @ez.publisher(OUTPUT_START)
-    async def start_file(self, message: Path) -> AsyncGenerator:
+    async def start_file(self, message: Path) -> typing.AsyncGenerator:
         out = self.open_file(message)
         if out is not None:
             yield (self.OUTPUT_START, out)
 
     @ez.subscriber(INPUT_STOP)
     @ez.publisher(OUTPUT_STOP)
-    async def stop_file(self, message: Path) -> AsyncGenerator:
+    async def stop_file(self, message: Path) -> typing.AsyncGenerator:
         out = self.close_file(message)
         if out is not None:
             yield (self.OUTPUT_STOP, out)
 
     @ez.subscriber(INPUT_MESSAGE)
     @ez.publisher(OUTPUT_MESSAGE)
-    async def on_message(self, message: Any) -> AsyncGenerator:
-        strmessage = log_object(message)
-        for output_f in self.STATE.output_files.values():
-            output_f.write(f"{strmessage}\n")
-            output_f.flush()
+    async def on_message(self, message: typing.Any) -> typing.AsyncGenerator:
+        strmessage = f"{log_object(message)}\n"
+
+        if self.SETTINGS.write_period <= 0:
+            for output_f in self.STATE.output_files.values():
+                output_f.write(strmessage)
+                output_f.flush()
+        else:
+            self.STATE.write_buffer.append(strmessage)
+
         yield (self.OUTPUT_MESSAGE, message)
 
+    @ez.task
+    async def buffered_write(self) -> None:
+        if self.SETTINGS.write_period <= 0:
+            return
+        
+        rate = Rate(1.0 / self.SETTINGS.write_period)
+        while True:
+            for output_f in self.STATE.output_files.values():
+                output_f.writelines(self.STATE.write_buffer)
+                output_f.flush()
+            self.STATE.write_buffer.clear()
+            await rate.sleep()
+
     async def shutdown(self) -> None:
         """Note that files that are closed at shutdown don't publish messages"""
+
+        if len(self.STATE.write_buffer):
+            for output_f in self.STATE.output_files.values():
+                output_f.writelines(self.STATE.write_buffer)
+                output_f.flush()
+            self.STATE.write_buffer.clear()
+
         for filepath in list(self.STATE.output_files):
             self.close_file(filepath)

{ezmsg-3.6.0 → ezmsg-3.6.2}/src/ezmsg/util/messages/axisarray.py

@@ -1,18 +1,25 @@
+from abc import abstractmethod, ABC
+from contextlib import contextmanager
+from dataclasses import field, dataclass
 import math
 import typing
 import warnings
 
-from abc import abstractmethod, ABC
-from contextlib import contextmanager
-from dataclasses import field, dataclass
+import ezmsg.core as ez
 
-import numpy as np
-import numpy.typing as npt
-import numpy.lib.stride_tricks as nps
+try:
+    import numpy as np
+    import numpy.typing as npt
+    import numpy.lib.stride_tricks as nps
+except ModuleNotFoundError:
+    ez.logger.error("Install ezmsg with the AxisArray extra:"
+                    'pip install "ezmsg[AxisArray]"')
+    raise
 
 from ezmsg.core.util import either_dict_or_kwargs
 from .util import replace
 
+
 if typing.TYPE_CHECKING:
     try:
         from xarray import DataArray

{ezmsg-3.6.0 → ezmsg-3.6.2}/uv.lock

@@ -240,6 +240,9 @@ dependencies = [
 ]
 
 [package.optional-dependencies]
+axisarray = [
+    { name = "numpy" },
+]
 docs = [
     { name = "ezmsg-sigproc" },
     { name = "sphinx" },
@@ -258,7 +261,8 @@ test = [
 requires-dist = [
     { name = "ezmsg-sigproc", marker = "extra == 'docs'", specifier = ">=1.2.3" },
     { name = "flake8", marker = "extra == 'test'", specifier = ">=5.0.4" },
-    { name = "numpy", marker = "extra == 'test'", specifier = ">=1.24.2" },
+    { name = "numpy", marker = "extra == 'axisarray'", specifier = ">=1.24.4" },
+    { name = "numpy", marker = "extra == 'test'", specifier = ">=1.24.4" },
     { name = "pytest", marker = "extra == 'test'", specifier = ">=7.0.0" },
     { name = "pytest-asyncio", marker = "extra == 'test'", specifier = ">=0.23.8" },
     { name = "pytest-cov", marker = "extra == 'test'", specifier = ">=5.0.0" },

ezmsg-3.6.0/.envrc

@@ -1,2 +0,0 @@
-export VIRTUAL_ENV=.venv
-layout python

ezmsg-3.6.0/src/ezmsg/core/command.py

@@ -1,319 +0,0 @@
-import os
-import sys
-import base64
-import asyncio
-import argparse
-import logging
-import subprocess
-import typing
-import webbrowser
-
-from collections import defaultdict
-from uuid import uuid4
-from textwrap import indent
-
-from .graphserver import GraphService
-from .shmserver import SHMService
-from .netprotocol import (
-    Address,
-    GRAPHSERVER_ADDR_ENV,
-    GRAPHSERVER_PORT_DEFAULT,
-    SHMSERVER_ADDR_ENV,
-    SHMSERVER_PORT_DEFAULT,
-    PUBLISHER_START_PORT_ENV,
-    PUBLISHER_START_PORT_DEFAULT,
-    close_stream_writer,
-)
-from .dag import DAG
-
-logger = logging.getLogger("ezmsg")
-
-IND = "  "
-
-
-def cmdline() -> None:
-    parser = argparse.ArgumentParser(
-        "ezmsg.core",
-        description="start and stop core ezmsg server processes",
-        epilog=f"""
-            You can also change server configuration with environment variables.
-            GraphServer will be hosted on ${GRAPHSERVER_ADDR_ENV} (default port: {GRAPHSERVER_PORT_DEFAULT}).  
-            SHMServer will be hosted on ${SHMSERVER_ADDR_ENV} (default port: {SHMSERVER_PORT_DEFAULT}).
-            Publishers will be assigned available ports starting from {PUBLISHER_START_PORT_DEFAULT}. (Change with ${PUBLISHER_START_PORT_ENV})
-        """,
-    )
-
-    parser.add_argument(
-        "command",
-        help="command for ezmsg",
-        choices=["serve", "start", "shutdown", "graphviz", "mermaid"],
-    )
-
-    parser.add_argument("--address", help="Address for GraphServer", default=None)
-
-    class Args:
-        command: str
-        address: typing.Optional[str]
-
-    args = parser.parse_args(namespace=Args)
-
-    graph_address = Address("127.0.0.1", GRAPHSERVER_PORT_DEFAULT)
-    if args.address is not None:
-        graph_address = Address.from_string(args.address)
-    shm_address_str = os.environ.get(
-        SHMSERVER_ADDR_ENV, f"127.0.0.1:{SHMSERVER_PORT_DEFAULT}"
-    )
-    shm_address = Address.from_string(shm_address_str)
-
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-
-    loop.run_until_complete(run_command(args.command, graph_address, shm_address))
-
-
-async def run_command(cmd: str, graph_address: Address, shm_address: Address) -> None:
-    shm_service = SHMService(shm_address)
-    graph_service = GraphService(graph_address)
-
-    if cmd == "serve":
-        logger.info(f"GraphServer Address: {graph_address}")
-        logger.info(f"SHMServer Address: {shm_address}")
-
-        shm_server = shm_service.create_server()
-        graph_server = graph_service.create_server()
-
-        try:
-            logger.info("Servers running...")
-            graph_server.join()
-
-        except KeyboardInterrupt:
-            logger.info("Interrupt detected; shutting down servers")
-
-        finally:
-            if graph_server is not None:
-                graph_server.stop()
-
-            if shm_server is not None:
-                shm_server.stop()
-
-    elif cmd == "start":
-        popen = subprocess.Popen(
-            [sys.executable, "-m", "ezmsg.core", "serve", f"--address={graph_address}"]
-        )
-
-        while True:
-            try:
-                _, writer = await graph_service.open_connection()
-                await close_stream_writer(writer)
-                _, writer = await shm_service.open_connection()
-                await close_stream_writer(writer)
-                break
-            except ConnectionRefusedError:
-                await asyncio.sleep(0.1)
-
-        logger.info(f"Forked ezmsg servers in PID: {popen.pid}")
-
-    elif cmd == "shutdown":
-        try:
-            await graph_service.shutdown()
-            logger.info(
-                f"Issued shutdown command to GraphServer @ {graph_service.address}"
-            )
-
-        except ConnectionRefusedError:
-            logger.warning(
-                f"Could not issue shutdown command to GraphServer @ {graph_service.address}; server not running?"
-            )
-
-    elif cmd == "graphviz":
-        try:
-            dag: DAG = await graph_service.dag()
-        except (ConnectionRefusedError, ConnectionResetError):
-            logger.info(
-                f"GraphServer not running @{graph_address}, or host is refusing connections"
-            )
-            return
-
-        graph_connections = dag.graph.copy()
-        # Let's eliminate proxy topics, i.e. connections with inputs and outputs.
-        source_nodes = []
-        for node, conns in graph_connections.items():
-            if len(conns) > 0:
-                source_nodes += [node]
-        proxy_topics = []
-        for conns in graph_connections.values():
-            for conn in conns:
-                if conn in source_nodes and conn not in proxy_topics:
-                    proxy_topics += [conn]
-        # Replace Proxy Topics with actual source and downstream
-        for proxy_topic in proxy_topics:
-            downstreams = graph_connections.pop(proxy_topic)
-            logger.info(f"{proxy_topic} downstream connetions: {downstreams}")
-            for node, conns in graph_connections.items():
-                for conn in conns:
-                    if conn == proxy_topic:
-                        new_conns = conns.copy()
-                        new_conns.remove(proxy_topic)
-                        new_conns.union(downstreams)
-                        logger.info(
-                            f"Updating connections for {node} from {conns} to {new_conns}"
-                        )
-                        graph_connections[node] = new_conns
-
-        graph_connections = dag.graph.copy()
-        # Let's come up with UUID node names
-        nodes = set(graph_connections.keys())
-        node_map = {name: f'"{str(uuid4())}"' for name in nodes}
-
-        # Construct the graph
-        def tree():
-            return defaultdict(tree)
-
-        graph: defaultdict = tree()
-
-        connections = ""
-        for node, conns in graph_connections.items():
-            subgraph = graph
-            path = node.split("/")
-            route = path[:-1]
-            stream = path[-1]
-            for seg in route:
-                subgraph = subgraph[seg]
-            subgraph[stream] = node
-
-            for sub in conns:
-                connections += f"{node_map[node]} -> {node_map[sub]};" + "\n"
-
-        # Now convert to dot syntax
-        def recurse_graph(g: defaultdict):
-            out = ""
-            for leaf in g:
-                if isinstance(g[leaf], defaultdict):
-                    out += indent(
-                        "\n".join(
-                            [
-                                f"subgraph {leaf.lower()} {{",
-                                indent("cluster = true;", IND),
-                                indent(f'label = "{leaf}";', IND),
-                                f"{recurse_graph(g[leaf])}}};",
-                                "",
-                            ]
-                        ),
-                        IND,
-                    )
-                elif isinstance(g[leaf], str):
-                    out += indent(f"{node_map[g[leaf]]} [label={leaf}];" + "\n", IND)
-            return out
-
-        subgraph_tree = recurse_graph(graph)
-
-        graphviz_out = "\n".join(
-            [
-                "digraph EZ {",
-                indent('rankdir="LR"', IND),
-                subgraph_tree,
-                indent(connections, IND),
-                "}",
-            ]
-        )
-
-        print(graphviz_out)
-    elif cmd == "mermaid":
-        try:
-            dag: DAG = await graph_service.dag()
-        except (ConnectionRefusedError, ConnectionResetError):
-            logger.info(
-                f"GraphServer not running @{graph_address}, or host is refusing connections"
-            )
-            return
-
-        graph_connections = dag.graph.copy()
-        # Let's eliminate proxy topics, i.e. connections with inputs and outputs.
-        source_nodes = []
-        for node, conns in graph_connections.items():
-            if len(conns) > 0:
-                source_nodes += [node]
-        proxy_topics = []
-        for conns in graph_connections.values():
-            for conn in conns:
-                if conn in source_nodes and conn not in proxy_topics:
-                    proxy_topics += [conn]
-        # Replace Proxy Topics with actual source and downstream
-        for proxy_topic in proxy_topics:
-            downstreams = graph_connections.pop(proxy_topic)
-            logger.info(f"{proxy_topic} downstream connetions: {downstreams}")
-            for node, conns in graph_connections.items():
-                for conn in conns:
-                    if conn == proxy_topic:
-                        new_conns = conns.copy()
-                        new_conns.remove(proxy_topic)
-                        new_conns.union(downstreams)
-                        logger.info(
-                            f"Updating connections for {node} from {conns} to {new_conns}"
-                        )
-                        graph_connections[node] = new_conns
-
-        graph_connections = dag.graph.copy()
-        # Let's come up with UUID node names
-        nodes = set(graph_connections.keys())
-        node_map = {name: f"{str(uuid4())}" for name in nodes}
-
-        # Construct the graph
-        def tree():
-            return defaultdict(tree)
-
-        graph: defaultdict = tree()
-
-        connections = ""
-        for node, conns in graph_connections.items():
-            subgraph = graph
-            path = node.split("/")
-            route = path[:-1]
-            stream = path[-1]
-            for seg in route:
-                subgraph = subgraph[seg]
-            subgraph[stream] = node
-
-            for sub in conns:
-                connections += f"{node_map[node]} --> {node_map[sub]}" + "\n"
-
-        # Now convert to dot syntax
-        def recurse_graph(g: defaultdict):
-            out = ""
-            for leaf in g:
-                if isinstance(g[leaf], defaultdict):
-                    out += indent(
-                        "\n".join(
-                            [
-                                f"subgraph {leaf.lower()} [{leaf}]",
-                                # "direction LR",
-                                f"{recurse_graph(g[leaf])}",
-                                "end",
-                                "",
-                            ]
-                        ),
-                        IND,
-                    )
-                elif isinstance(g[leaf], str):
-                    out += indent(f"{node_map[g[leaf]]}[{leaf}]" + "\n", IND)
-            return out
-
-        subgraph_tree = recurse_graph(graph)
-
-        mermaid = "\n".join(
-            [
-                "flowchart LR",
-                subgraph_tree,
-                indent(connections, IND),
-            ]
-        )
-
-        # print(mermaid)
-        webbrowser.open(mm(mermaid))
-
-
-def mm(graph):
-    graphbytes = graph.encode("utf8")
-    base64_bytes = base64.b64encode(graphbytes)
-    base64_string = base64_bytes.decode("ascii")
-    return f"https://mermaid.ink/img/{base64_string}"