Merge pull request #240 from ezmsg-org/feature/dashboard-integration

Feature: Dashboard integration

Author: griffinmilsap

README.md

@@ -95,4 +95,4 @@ These publications provide insights into the practical applications and impact o
 
 ## Financial Support
 
-`ezmsg` is supported by Johns Hopkins University (JHU), the JHU Applied Physics Laboratory (APL), and by the Wyss Center for Bio and Neuro Engineering.
+`ezmsg` is supported by Johns Hopkins University (JHU), the JHU Applied Physics Laboratory (APL), Blackrock Neurotech and by the Wyss Center for Bio and Neuro Engineering.

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ezmsg"
-version = "3.8.0"
+version = "3.9.0"
 description = "A simple DAG-based computation model"
 authors = [
   { name = "Griffin Milsap", email = "griffin.milsap@gmail.com" },
@@ -52,9 +52,6 @@ docs = [
 axisarray = [
     "numpy>=2.2.6",
 ]
-perf = [
-  "xarray",
-]
 
 [project.scripts]
 ezmsg = "ezmsg.core.command:cmdline"
@@ -63,6 +60,12 @@ ezmsg = "ezmsg.core.command:cmdline"
 axisarray = [
     "numpy>=2.2.6",
 ]
+perf = [
+  "xarray>=2025.6.1",
+]
+dashboard = [
+  "ezmsg-dashboard; python_version >= '3.11'",
+]
 
 [tool.pytest.ini_options]
 addopts = ["--import-mode=importlib"]

src/ezmsg/core/command.py

@@ -10,11 +10,17 @@
 from .commands.start import handle_start
 from .netprotocol import (
     Address,
+    DEFAULT_HOST,
     GRAPHSERVER_ADDR_ENV,
     GRAPHSERVER_PORT_DEFAULT,
     PUBLISHER_START_PORT_ENV,
     PUBLISHER_START_PORT_DEFAULT,
 )
+from .commands.dashboard import (
+    DASHBOARD_ADDR_ENV,
+    DASHBOARD_INSTALL_HINT,
+    DASHBOARD_PORT_DEFAULT,
+)
 
 
 def build_parser() -> argparse.ArgumentParser:
@@ -30,6 +36,7 @@ def build_parser() -> argparse.ArgumentParser:
         epilog=f"""
             You can also change server configuration with environment variables.
             GraphServer will be hosted on ${GRAPHSERVER_ADDR_ENV} (default port: {GRAPHSERVER_PORT_DEFAULT}).  
+            Dashboard will be hosted on ${DASHBOARD_ADDR_ENV} (default: {DEFAULT_HOST}:{DASHBOARD_PORT_DEFAULT}, or graph port + 1).
             Publishers will be assigned available ports starting from {PUBLISHER_START_PORT_DEFAULT}. (Change with ${PUBLISHER_START_PORT_ENV})
         """,
     )
@@ -55,7 +62,12 @@ def cmdline(argv: list[str] | None = None) -> None:
 
     result = args._handler(args)
     if inspect.isawaitable(result):
-        asyncio.run(result)
+        try:
+            asyncio.run(result)
+        except KeyboardInterrupt:
+            # asyncio.run() re-raises KeyboardInterrupt after cancelling the main
+            # task on Ctrl+C, even when command cleanup has already completed.
+            pass
 
 
 async def run_command(
@@ -64,8 +76,10 @@ async def run_command(
     target: str = "live",
     compact: int | None = None,
     nobrowser: bool = False,
+    dashboard: int | bool | None = None,
 ) -> None:
     handlers = {
+        "dashboard": None,
         "serve": handle_serve,
         "start": handle_start,
         "shutdown": handle_shutdown,
@@ -74,11 +88,25 @@ async def run_command(
     }
     if cmd not in handlers:
         raise ValueError(f"Unknown ezmsg command '{cmd}'")
+    if cmd == "dashboard":
+        try:
+            from ezmsg.dashboard.server import handle_dashboard
+        except ImportError as exc:
+            raise RuntimeError(DASHBOARD_INSTALL_HINT) from exc
+        handlers["dashboard"] = handle_dashboard
     args = argparse.Namespace(
         command=cmd,
         address=str(graph_address),
+        graph_address=str(graph_address),
         target=target,
         compact=compact,
         nobrowser=nobrowser,
+        dashboard=dashboard,
+        host="127.0.0.1",
+        port=8000,
+        open_browser=False,
+        log_level="info",
     )
-    await handlers[cmd](args)
+    result = handlers[cmd](args)
+    if inspect.isawaitable(result):
+        await result

src/ezmsg/core/commands/__init__.py

@@ -1,5 +1,6 @@
 import argparse
 
+from .dashboard_cmd import setup_dashboard_cmdline
 from .graphviz import setup_graphviz_cmdline
 from .mermaid import setup_mermaid_cmdline
 from .serve import setup_serve_cmdline
@@ -8,6 +9,7 @@
 
 
 def setup_core_cmdline(subparsers: argparse._SubParsersAction) -> None:
+    setup_dashboard_cmdline(subparsers)
     setup_serve_cmdline(subparsers)
     setup_start_cmdline(subparsers)
     setup_shutdown_cmdline(subparsers)

src/ezmsg/core/commands/common.py

@@ -1,6 +1,7 @@
 import argparse
 
-from ..netprotocol import Address, GRAPHSERVER_PORT_DEFAULT
+from ..graphserver import GraphService
+from ..netprotocol import Address
 
 
 def add_address_argument(parser: argparse.ArgumentParser) -> None:
@@ -21,5 +22,5 @@ def add_compact_argument(parser: argparse.ArgumentParser) -> None:
 
 def graph_address_from_args(args: argparse.Namespace) -> Address:
     if args.address is None:
-        return Address("127.0.0.1", GRAPHSERVER_PORT_DEFAULT)
+        return GraphService.default_address()
     return Address.from_string(args.address)

src/ezmsg/core/commands/dashboard.py

@@ -0,0 +1,77 @@
+import argparse
+import os
+from typing import Any
+
+from ..netprotocol import (
+    Address,
+    DEFAULT_HOST,
+    GRAPHSERVER_ADDR_ENV,
+    GRAPHSERVER_PORT_DEFAULT,
+)
+
+DASHBOARD_ADDR_ENV = "EZMSG_DASHBOARD_ADDR"
+DASHBOARD_PORT_DEFAULT = GRAPHSERVER_PORT_DEFAULT + 1
+DASHBOARD_INSTALL_HINT = (
+    "Dashboard support requires the optional `ezmsg-dashboard` package. "
+    "Install it with `pip install ezmsg-dashboard`."
+)
+
+
+class DashboardDependencyError(RuntimeError):
+    pass
+
+
+def add_dashboard_argument(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--dashboard",
+        nargs="?",
+        const=True,
+        default=None,
+        type=int,
+        metavar="PORT",
+        help=(
+            "Serve the optional ezmsg dashboard alongside the graph server. "
+            "If PORT is omitted, ezmsg uses the configured dashboard address or graph port + 1."
+        ),
+    )
+
+
+def default_graph_address() -> Address:
+    address_str = os.environ.get(
+        GRAPHSERVER_ADDR_ENV, f"{DEFAULT_HOST}:{GRAPHSERVER_PORT_DEFAULT}"
+    )
+    return Address.from_string(address_str)
+
+
+def dashboard_address(
+    graph_address: Address | None = None, dashboard_port: int | None = None
+) -> Address:
+    if DASHBOARD_ADDR_ENV in os.environ:
+        address = Address.from_string(os.environ[DASHBOARD_ADDR_ENV])
+    else:
+        resolved_graph_address = graph_address or default_graph_address()
+        address = Address(resolved_graph_address.host, resolved_graph_address.port + 1)
+
+    if dashboard_port is not None:
+        return Address(address.host, dashboard_port)
+    return address
+
+
+def require_dashboard_dependency() -> Any:
+    try:
+        from ezmsg.dashboard.server import start_dashboard_server
+    except ImportError as exc:
+        raise DashboardDependencyError(DASHBOARD_INSTALL_HINT) from exc
+    return start_dashboard_server
+
+
+def start_dashboard(graph_address: Address, dashboard_port: int | None = None) -> Any:
+    start_dashboard_server = require_dashboard_dependency()
+
+    address = dashboard_address(graph_address, dashboard_port=dashboard_port)
+    return start_dashboard_server(
+        graph_address=graph_address,
+        host=address.host,
+        port=address.port,
+        log_level="warning",
+    )

src/ezmsg/core/commands/dashboard_cmd.py

@@ -0,0 +1,43 @@
+import argparse
+import logging
+
+from .dashboard import DASHBOARD_INSTALL_HINT
+
+logger = logging.getLogger("ezmsg")
+
+
+def _warn_dashboard_dependency_missing(_: argparse.Namespace) -> None:
+    logger.warning(DASHBOARD_INSTALL_HINT)
+
+
+def _setup_dashboard_fallback(subparsers: argparse._SubParsersAction) -> None:
+    parser = subparsers.add_parser(
+        "dashboard",
+        help="launch the optional ezmsg dashboard server",
+        description="Launch the optional ezmsg dashboard server.",
+    )
+    parser.add_argument("--graph-address", default=None, help="Address of the ezmsg graph server.")
+    parser.add_argument("--host", default="127.0.0.1", help="HTTP bind host for the dashboard.")
+    parser.add_argument("--port", type=int, default=8000, help="HTTP bind port for the dashboard.")
+    parser.add_argument(
+        "--open-browser",
+        action="store_true",
+        help="Open the dashboard in a browser after startup.",
+    )
+    parser.add_argument(
+        "--log-level",
+        default="info",
+        choices=["critical", "error", "warning", "info", "debug", "trace"],
+        help="Uvicorn log verbosity.",
+    )
+    parser.set_defaults(_handler=_warn_dashboard_dependency_missing)
+
+
+def setup_dashboard_cmdline(subparsers: argparse._SubParsersAction) -> None:
+    try:
+        from ezmsg.dashboard.server import setup_dashboard_cmdline as setup_optional_dashboard
+    except ImportError:
+        _setup_dashboard_fallback(subparsers)
+        return
+
+    setup_optional_dashboard(subparsers)

src/ezmsg/core/commands/serve.py

@@ -4,6 +4,11 @@
 
 from ..graphserver import GraphService
 from .common import add_address_argument, graph_address_from_args
+from .dashboard import (
+    DashboardDependencyError,
+    add_dashboard_argument,
+    start_dashboard,
+)
 
 logger = logging.getLogger("ezmsg")
 
@@ -14,17 +19,29 @@ async def handle_serve(args: argparse.Namespace) -> None:
 
     logger.info(f"GraphServer Address: {graph_address}")
     graph_server = graph_service.create_server()
+    dashboard_server = None
 
     try:
+        if args.dashboard is not None:
+            dashboard_port = args.dashboard if type(args.dashboard) is int else None
+            dashboard_server = start_dashboard(
+                graph_service.address, dashboard_port=dashboard_port
+            )
+            logger.info(f"Dashboard Address: {dashboard_server.url}")
         logger.info("Servers running...")
         await asyncio.to_thread(graph_server.join)
-    except KeyboardInterrupt:
+    except (KeyboardInterrupt, asyncio.CancelledError):
         logger.info("Interrupt detected; shutting down servers")
+    except DashboardDependencyError as exc:
+        logger.warning(str(exc))
     finally:
+        if dashboard_server is not None:
+            dashboard_server.stop()
         graph_server.stop()
 
 
 def setup_serve_cmdline(subparsers: argparse._SubParsersAction) -> None:
     parser = subparsers.add_parser("serve")
     add_address_argument(parser)
+    add_dashboard_argument(parser)
     parser.set_defaults(_handler=handle_serve)

src/ezmsg/core/commands/start.py

@@ -7,17 +7,30 @@
 from ..graphserver import GraphService
 from ..netprotocol import close_stream_writer
 from .common import add_address_argument, graph_address_from_args
+from .dashboard import (
+    DashboardDependencyError,
+    add_dashboard_argument,
+    require_dashboard_dependency,
+)
 
 logger = logging.getLogger("ezmsg")
 
 
 async def handle_start(args: argparse.Namespace) -> None:
     graph_address = graph_address_from_args(args)
     graph_service = GraphService(graph_address)
+    cmd = [sys.executable, "-m", "ezmsg.core", "serve", f"--address={graph_address}"]
+    if args.dashboard is not None:
+        try:
+            require_dashboard_dependency()
+        except DashboardDependencyError as exc:
+            logger.warning(str(exc))
+            return
+        cmd.append("--dashboard")
+        if type(args.dashboard) is int:
+            cmd.append(str(args.dashboard))
 
-    popen = subprocess.Popen(
-        [sys.executable, "-m", "ezmsg.core", "serve", f"--address={graph_address}"]
-    )
+    popen = subprocess.Popen(cmd)
 
     while True:
         try:
@@ -33,4 +46,5 @@ async def handle_start(args: argparse.Namespace) -> None:
 def setup_start_cmdline(subparsers: argparse._SubParsersAction) -> None:
     parser = subparsers.add_parser("start")
     add_address_argument(parser)
+    add_dashboard_argument(parser)
     parser.set_defaults(_handler=handle_start)

src/ezmsg/util/perf/run.py

@@ -145,9 +145,8 @@ def benchmark(
         )
 
     try:
-        communications = (
-            DEFAULT_COMMS if comms is None else [Communication(c) for c in comms]
-        )
+        communication_names = DEFAULT_COMMS if comms is None else list(comms)
+        communications = [Communication(c) for c in communication_names]
     except ValueError:
         ez.logger.error(
             f"Invalid test communications requested. Valid communications: {', '.join([c.value for c in Communication])}"