diff --git a/dimos/core/coordination/coordinator_rpc.py b/dimos/core/coordination/coordinator_rpc.py index ac4fe9960b..3c64fdc85c 100644 --- a/dimos/core/coordination/coordinator_rpc.py +++ b/dimos/core/coordination/coordinator_rpc.py @@ -19,6 +19,8 @@ from dimos.core.global_config import global_config from dimos.core.transport_factory import rpc_backend +from dimos.protocol.rpc.zenohrpc import ZenohRPC +from dimos.protocol.service.zenohservice import ZENOH_LOCAL_ROUTER_ENDPOINT from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -51,7 +53,12 @@ def serve(cls, coordinator: RPCInspectable) -> CoordinatorRPC: @classmethod def connect(cls, *, timeout: float) -> CoordinatorRPC: """Attach to a running Coordinator, raising `TimeoutError` if none answers.""" - rpc = rpc_backend()() + backend = rpc_backend() + rpc = ( + ZenohRPC(mode="client", connect=[ZENOH_LOCAL_ROUTER_ENDPOINT]) + if backend is ZenohRPC + else backend() + ) rpc.start() client = cls(rpc) deadline = time.monotonic() + timeout diff --git a/dimos/core/coordination/module_coordinator.py b/dimos/core/coordination/module_coordinator.py index 24b7a70ce2..6187d73a99 100644 --- a/dimos/core/coordination/module_coordinator.py +++ b/dimos/core/coordination/module_coordinator.py @@ -19,6 +19,7 @@ import dataclasses import importlib import inspect +import os import shutil import sys import threading @@ -41,6 +42,11 @@ pZenohTransport, ) from dimos.core.transport_factory import make_transport +from dimos.protocol.service.zenohservice import ( + ZENOH_LOCAL_ROUTER_ENDPOINT, + ZENOH_ROUTER_ENDPOINT_ENV, + ZenohRouter, +) from dimos.spec.utils import is_spec, spec_annotation_compliance, spec_structural_compliance from dimos.utils.generic import short_id from dimos.utils.logging_config import setup_logger @@ -89,13 +95,24 @@ def __init__( self._modules_lock = threading.RLock() self._rpc_lock = threading.RLock() self._coordinator_rpc: CoordinatorRPC | None = None + self._zenoh_router: ZenohRouter | None = None + self._previous_zenoh_router_endpoint: str | None = None def start(self) -> None: from dimos.core.o3dpickle import register_picklers register_picklers() - for m in self._managers.values(): - m.start() + if self._global_config.transport == "zenoh": + self._zenoh_router = ZenohRouter() + self._zenoh_router.start() + self._previous_zenoh_router_endpoint = os.environ.get(ZENOH_ROUTER_ENDPOINT_ENV) + os.environ[ZENOH_ROUTER_ENDPOINT_ENV] = ZENOH_LOCAL_ROUTER_ENDPOINT + try: + for m in self._managers.values(): + m.start() + except BaseException: + self._stop_zenoh_router() + raise self._started = True def stop(self) -> None: @@ -119,9 +136,21 @@ def _stop_manager(m: WorkerManager) -> None: logger.error("Error stopping manager", manager=type(m).__name__, exc_info=True) safe_thread_map(tuple(self._managers.values()), _stop_manager) + self._stop_zenoh_router() + + def _stop_zenoh_router(self) -> None: + if self._zenoh_router is not None: + self._zenoh_router.stop() + self._zenoh_router = None + if self._global_config.transport != "zenoh": + return + if self._previous_zenoh_router_endpoint is None: + os.environ.pop(ZENOH_ROUTER_ENDPOINT_ENV, None) + else: + os.environ[ZENOH_ROUTER_ENDPOINT_ENV] = self._previous_zenoh_router_endpoint def start_rpc_service(self) -> None: - """Expose the coordinator's API as @rpc methods over LCM.""" + """Expose the coordinator's API over the configured RPC transport.""" with self._rpc_lock: if self._coordinator_rpc is not None: return diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index 461bd4d0ed..d7b712a66f 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -60,6 +60,10 @@ class MyCppModule(NativeModule): from dimos.core.core import rpc from dimos.core.global_config import global_config from dimos.core.module import Module, ModuleConfig +from dimos.protocol.service.zenohservice import ( + ZENOH_LOCAL_ROUTER_ENDPOINT, + ZENOH_ROUTER_ENDPOINT_ENV, +) from dimos.utils.logging_config import setup_logger if sys.platform.startswith("linux"): @@ -254,6 +258,8 @@ def start(self) -> None: # set transport so native modules know which one to spawn env["DIMOS_TRANSPORT"] = global_config.transport + if global_config.transport == "zenoh": + env[ZENOH_ROUTER_ENDPOINT_ENV] = ZENOH_LOCAL_ROUTER_ENDPOINT # set Rust logging to match Python level env["RUST_LOG"] = _PYTHON_TO_RUST_LEVELS.get( diff --git a/dimos/protocol/service/test_zenohservice.py b/dimos/protocol/service/test_zenohservice.py index 621c370485..c1406023b4 100644 --- a/dimos/protocol/service/test_zenohservice.py +++ b/dimos/protocol/service/test_zenohservice.py @@ -16,7 +16,13 @@ import pytest -from dimos.protocol.service.zenohservice import ZenohConfig, ZenohService, ZenohSessionPool +from dimos.protocol.service.zenohservice import ( + ZENOH_LOCAL_ROUTER_ENDPOINT, + ZENOH_ROUTER_ENDPOINT_ENV, + ZenohConfig, + ZenohService, + ZenohSessionPool, +) @pytest.fixture() @@ -33,6 +39,17 @@ def test_different_modes_produce_different_keys() -> None: assert peer.session_key != client.session_key +def test_default_config_uses_local_router_when_configured( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv(ZENOH_ROUTER_ENDPOINT_ENV, ZENOH_LOCAL_ROUTER_ENDPOINT) + + config = ZenohConfig() + + assert config.mode == "client" + assert config.connect == [ZENOH_LOCAL_ROUTER_ENDPOINT] + + def test_start_creates_session(session_pool) -> None: svc = ZenohService(session_pool=session_pool) svc.start() diff --git a/dimos/protocol/service/zenohservice.py b/dimos/protocol/service/zenohservice.py index 22bc364377..227df0c29b 100644 --- a/dimos/protocol/service/zenohservice.py +++ b/dimos/protocol/service/zenohservice.py @@ -15,6 +15,7 @@ from __future__ import annotations import json +import os import platform import socket import threading @@ -31,6 +32,10 @@ logger = setup_logger() +ZENOH_ROUTER_ENDPOINT_ENV = "DIMOS_ZENOH_ROUTER_ENDPOINT" +ZENOH_LOCAL_ROUTER_ENDPOINT = "tcp/127.0.0.1:7447" +ZENOH_LOCAL_ROUTER_LISTEN = "tcp/[::]:7447" + # Robot-side bridges (e.g. go2web) listen here so a remote dimos can dial in # when multicast discovery fails. Zenoh's own default port. ROBOT_ZENOH_PORT = 7447 @@ -43,14 +48,12 @@ LOOPBACK_INTERFACE = "lo0" if platform.system() == "Darwin" else "lo" -def _default_connect_endpoints() -> list[str]: - """Dial known robots directly instead of trusting multicast scouting. +def _robot_connect_endpoints() -> list[str]: + """Return explicit robot endpoints when multicast scouting is insufficient. Many APs filter multicast between WiFi clients, so a robot that is - perfectly reachable over TCP never answers a scout. When the session is - zenoh-transported and a robot IP is configured, it becomes an explicit - endpoint; scouting stays on for everything else. An IP carrying its own - ``:port`` is used as given. + reachable over TCP may never answer a scout. An IP carrying its own port is + used as given. """ from dimos.core.global_config import global_config @@ -67,6 +70,17 @@ def _default_connect_endpoints() -> list[str]: return out +def _default_mode() -> str: + return "client" if os.getenv(ZENOH_ROUTER_ENDPOINT_ENV) else "peer" + + +def _default_connect_endpoints() -> list[str]: + router_endpoint = os.getenv(ZENOH_ROUTER_ENDPOINT_ENV) + if router_endpoint: + return [router_endpoint] + return _robot_connect_endpoints() + + def _default_scouting() -> bool: from dimos.core.global_config import global_config @@ -98,10 +112,37 @@ def endpoint_addresses(endpoint: str) -> set[str]: return out +def _await_connect_endpoints( + session: zenoh.Session, + endpoints: list[str], + timeout: float, +) -> None: + pending = {endpoint: endpoint_addresses(endpoint) for endpoint in endpoints} + if not pending or timeout <= 0: + return + + deadline = time.monotonic() + timeout + while pending: + linked = {str(link.dst).rpartition("/")[2] for link in session.info.links()} + for endpoint in [item for item, addresses in pending.items() if addresses & linked]: + logger.debug("Zenoh linked", endpoint=endpoint) + del pending[endpoint] + if not pending: + return + if time.monotonic() >= deadline: + logger.warning( + "Zenoh endpoints not linked; continuing", + timeout=timeout, + endpoints=sorted(pending), + ) + return + time.sleep(_CONNECT_POLL_INTERVAL) + + class ZenohConfig(BaseConfig): - mode: str = "peer" + mode: str = Field(default_factory=_default_mode) connect: list[str] = Field(default_factory=_default_connect_endpoints) - listen: list[str] = [] + listen: list[str] = Field(default_factory=list) # Discover peers across the network. Off keeps discovery on loopback. scouting: bool = Field(default_factory=_default_scouting) # Seconds to block in start() waiting for `connect` endpoints to link. @@ -156,6 +197,44 @@ def close_all(self) -> None: default_session_pool = ZenohSessionPool() +class ZenohRouter: + def __init__( + self, + listen: str = ZENOH_LOCAL_ROUTER_LISTEN, + connect: list[str] | None = None, + ) -> None: + self._listen = listen + self._connect = _robot_connect_endpoints() if connect is None else list(connect) + self._session: zenoh.Session | None = None + + def start(self) -> None: + config = zenoh.Config() + config.insert_json5("mode", '"router"') + config.insert_json5("listen/endpoints", json.dumps([self._listen])) + if self._connect: + config.insert_json5("connect/endpoints", json.dumps(self._connect)) + if not _default_scouting(): + config.insert_json5("scouting/multicast/interface", json.dumps(LOOPBACK_INTERFACE)) + config.insert_json5("scouting/gossip/enabled", "false") + try: + self._session = zenoh.open(config) + _await_connect_endpoints( + self._session, + self._connect, + _default_connect_timeout(), + ) + logger.info("Local Zenoh router started", endpoint=self._listen) + except zenoh.ZError as exc: + if "Address already in use" not in str(exc): + raise + logger.info("Using existing local Zenoh router", endpoint=self._listen) + + def stop(self) -> None: + if self._session is not None: + self._session.close() + self._session = None + + class ZenohService(Service): config: ZenohConfig @@ -167,6 +246,9 @@ def __init__(self, *, session_pool: ZenohSessionPool | None = None, **kwargs: An self._session: zenoh.Session | None = None def start(self) -> None: + endpoint = os.getenv(ZENOH_ROUTER_ENDPOINT_ENV) + if endpoint and not self.config.model_fields_set: + self.config = ZenohConfig(mode="client", connect=[endpoint]) self._session = self._session_pool.acquire(self.config) self._await_connect(self._session) super().start() @@ -182,24 +264,11 @@ def _await_connect(self, session: zenoh.Session) -> None: Unreachable endpoints are a warning, not an error: one robot being down should not stop the rest of the graph from coming up. """ - pending = {ep: endpoint_addresses(ep) for ep in self.config.connect} - if not pending or self.config.connect_timeout <= 0: - return - deadline = time.monotonic() + self.config.connect_timeout - while pending: - linked = {str(link.dst).rpartition("/")[2] for link in session.info.links()} - for endpoint in [e for e, addrs in pending.items() if addrs & linked]: - logger.debug(f"Zenoh linked {endpoint}") - del pending[endpoint] - if not pending: - return - if time.monotonic() >= deadline: - logger.warning( - f"Zenoh endpoints not linked after {self.config.connect_timeout}s: " - f"{sorted(pending)} - continuing, published messages may be dropped" - ) - return - time.sleep(_CONNECT_POLL_INTERVAL) + _await_connect_endpoints( + session, + self.config.connect, + self.config.connect_timeout, + ) @property def session(self) -> zenoh.Session: diff --git a/docs/usage/transports/index.md b/docs/usage/transports/index.md index f27734fed9..27e34d93f3 100644 --- a/docs/usage/transports/index.md +++ b/docs/usage/transports/index.md @@ -351,6 +351,8 @@ Use Zenoh when: At the stream level, the transport wrappers are `ZenohTransport` and `pZenohTransport`. Install, defaults, and CLI versus environment overrides are in the [Zenoh quickstart](#zenoh-quickstart) above. +For a local `dimos run`, the coordinator starts or reuses one Zenoh router listening on port `7447`. Python workers, native modules, and local coordinator clients connect to it through `127.0.0.1`. This avoids an all-to-all peer mesh and keeps a local run independent of Wi-Fi, VPN, and Docker interface changes. The router remains reachable on other interfaces for explicitly configured remote participants. + Performance note: zenoh's session-to-session path (modules in different processes, the common case) benchmarks faster than LCM for small messages and for >=2MiB ones. Delivery *within* one shared session (co-located modules in one worker) is its slow path for 256KiB-1MiB messages (a few GiB/s); pin shared memory transports for heavy co-located streams. The benchmark has both cases (`Zenoh` = shared session, `ZenohPeers` = separate sessions). The Rerun bridge also follows the global transport. When `transport=zenoh`, the bridge listens on Zenoh and on LCM for TF data. diff --git a/native/rust/dimos-module/src/zenoh.rs b/native/rust/dimos-module/src/zenoh.rs index 9ff5209793..198ca89ce5 100644 --- a/native/rust/dimos-module/src/zenoh.rs +++ b/native/rust/dimos-module/src/zenoh.rs @@ -63,9 +63,16 @@ pub struct ZenohTransport { impl ZenohTransport { pub async fn new() -> io::Result { - let session = ::zenoh::open(::zenoh::Config::default()) - .await - .map_err(to_io)?; + let mut config = ::zenoh::Config::default(); + if let Ok(endpoint) = std::env::var("DIMOS_ZENOH_ROUTER_ENDPOINT") { + config.insert_json5("mode", r#""client""#).map_err(to_io)?; + let endpoints = serde_json::to_string(&[endpoint]) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + config + .insert_json5("connect/endpoints", &endpoints) + .map_err(to_io)?; + } + let session = ::zenoh::open(config).await.map_err(to_io)?; Ok(Self { session, qos: OnceLock::new(),