Merge pull request #228 from ezmsg-org/fix/collection-topics-relays

griffinmilsap · web-flow · commit b283f795de6e · 2026-03-16T20:52:42.000-04:00
Fix: Input/OutputTopic/Relay for Collections
diff --git a/examples/ezmsg_configs.py b/examples/ezmsg_configs.py
@@ -82,8 +82,8 @@ async def listen(self, msg: int) -> None:
 
 
 class PassthroughCollection(ez.Collection):
-    INPUT = ez.InputStream(int)
-    OUTPUT = ez.OutputStream(int)
+    INPUT = ez.InputTopic(int)
+    OUTPUT = ez.OutputTopic(int)
 
     def network(self) -> ez.NetworkDefinition:
         return ((self.INPUT, self.OUTPUT),)
@@ -136,7 +136,7 @@ def configure(self) -> None:
 
 
 class PubNoSubCollection(ez.Collection):
-    OUTPUT = ez.OutputStream(int)
+    OUTPUT = ez.OutputTopic(int)
     GENERATE = Generator()
     LOG = DebugLog()
 
@@ -148,7 +148,7 @@ def network(self) -> ez.NetworkDefinition:
 
 
 class SubNoPubCollection(ez.Collection):
-    INPUT = ez.InputStream(int)
+    INPUT = ez.InputTopic(int)
     LISTEN = Listener()
 
     def network(self) -> ez.NetworkDefinition:
@@ -175,7 +175,7 @@ class PubNoSubPassthroughCollection(ez.Collection):
     COLLECTION = PubNoSubCollection()
     PASSTHROUGH = PassthroughCollection()
 
-    OUTPUT = ez.OutputStream(int)
+    OUTPUT = ez.OutputTopic(int)
 
     def network(self) -> ez.NetworkDefinition:
         return (
@@ -188,7 +188,7 @@ class SubNoPubPassthroughCollection(ez.Collection):
     COLLECTION = SubNoPubCollection()
     PASSTHROUGH = PassthroughCollection()
 
-    INPUT = ez.InputStream(int)
+    INPUT = ez.InputTopic(int)
 
     def network(self) -> ez.NetworkDefinition:
         return (
diff --git a/examples/ezmsg_toy.py b/examples/ezmsg_toy.py
@@ -123,8 +123,8 @@ class ModifierCollection(ez.Collection):
     """This collection will subscribe to messages
     and append the most recent LFO output"""
 
-    INPUT = ez.InputStream(str)
-    OUTPUT = ez.OutputStream(str)
+    INPUT = ez.InputTopic(str)
+    OUTPUT = ez.OutputTopic(str)
 
     SIN = LFO()
     # SIN2 = LFO()
diff --git a/src/ezmsg/core/__init__.py b/src/ezmsg/core/__init__.py
@@ -14,6 +14,11 @@
     "Settings",
     "Collection",
     "NetworkDefinition",
+    "Topic",
+    "InputTopic",
+    "OutputTopic",
+    "InputRelay",
+    "OutputRelay",
     "InputStream",
     "OutputStream",
     "Unit",
@@ -44,7 +49,15 @@
 from .settings import Settings
 from .collection import Collection, NetworkDefinition
 from .unit import Unit, task, publisher, subscriber, main, timeit, process, thread
-from .stream import InputStream, OutputStream
+from .stream import (
+    Topic,
+    InputTopic,
+    OutputTopic,
+    InputRelay,
+    OutputRelay,
+    InputStream,
+    OutputStream,
+)
 from .backend import run, GraphRunner, GraphRunnerStartError
 from .backendprocess import Complete, NormalTermination
 from .graphserver import GraphServer
diff --git a/src/ezmsg/core/backend.py b/src/ezmsg/core/backend.py
@@ -5,6 +5,7 @@
 import logging
 import os
 import signal
+from dataclasses import dataclass
 from threading import BrokenBarrierError
 from multiprocessing import Event, Barrier
 from multiprocessing.synchronize import Event as EventType
@@ -16,8 +17,9 @@
 
 from .collection import Collection, NetworkDefinition
 from .component import Component
-from .stream import Stream
+from .stream import Stream, InputRelay, OutputRelay
 from .unit import Unit, PROCESS_ATTR
+from .relay import _CollectionRelayUnit, _RelaySettings
 
 from .graphserver import GraphService
 from .graphcontext import GraphContext
@@ -33,6 +35,16 @@
 logger = logging.getLogger("ezmsg")
 
 
+@dataclass
+class _RelayBinding:
+    kind: str  # "input" or "output"
+    endpoint_topic: str
+    relay_in_topic: str
+    relay_out_topic: str
+    endpoint: InputRelay | OutputRelay
+    relay_unit: _CollectionRelayUnit
+
+
 class ExecutionContext:
     _process_units: list[list[Unit]]
     _processes: list[BackendProcess] | None
@@ -95,22 +107,32 @@ def setup(
         start_participant: bool = False,
     ) -> "ExecutionContext | None":
         graph_connections: list[tuple[str, str]] = []
+        relay_bindings: dict[str, _RelayBinding] = {}
 
         for name, component in components.items():
             component._set_name(name)
             component._set_location([root_name] if root_name is not None else [])
 
+        def normalize_topic(endpoint: Stream | str | enum.Enum, where: str) -> str:
+            if isinstance(endpoint, Stream):
+                return endpoint.address
+            if isinstance(endpoint, enum.Enum):
+                return endpoint.name
+            if isinstance(endpoint, str):
+                return endpoint
+            raise TypeError(
+                f"Invalid endpoint type in {where}: {type(endpoint)}. "
+                "Expected Stream, str, or Enum."
+            )
+
         if connections is not None:
             for from_topic, to_topic in connections:
-                if isinstance(from_topic, Stream):
-                    from_topic = from_topic.address
-                if isinstance(to_topic, Stream):
-                    to_topic = to_topic.address
-                if isinstance(to_topic, enum.Enum):
-                    to_topic = to_topic.name
-                if isinstance(from_topic, enum.Enum):
-                    from_topic = from_topic.name
-                graph_connections.append((from_topic, to_topic))
+                graph_connections.append(
+                    (
+                        normalize_topic(from_topic, "connections"),
+                        normalize_topic(to_topic, "connections"),
+                    )
+                )
 
         def crawl_components(
             component: Component, callback: Callable[[Component], None]
@@ -121,23 +143,115 @@ def crawl_components(
                 search += list(comp.components.values())
                 callback(comp)
 
+        def input_relay_settings(relay: InputRelay) -> _RelaySettings:
+            return _RelaySettings(
+                leaky=relay.leaky,
+                max_queue=relay.max_queue,
+                copy_on_forward=relay.copy_on_forward,
+            )
+
+        def output_relay_settings(relay: OutputRelay) -> _RelaySettings:
+            return _RelaySettings(
+                host=relay.host,
+                port=relay.port,
+                num_buffers=relay.num_buffers,
+                buf_size=relay.buf_size,
+                force_tcp=relay.force_tcp,
+                copy_on_forward=relay.copy_on_forward,
+            )
+
+        def add_collection_relay_units(comp: Component) -> None:
+            if not isinstance(comp, Collection):
+                return
+
+            for endpoint_name, endpoint in comp.streams.items():
+                if isinstance(endpoint, InputRelay):
+                    relay_name = f"__relay_in_{endpoint_name}"
+                    if relay_name in comp.components:
+                        raise ValueError(
+                            f"{comp.address} already defines component '{relay_name}'."
+                        )
+
+                    relay_unit = _CollectionRelayUnit(input_relay_settings(endpoint))
+                    relay_unit._set_name(relay_name)
+                    relay_unit._set_location(comp.location + [comp.name])
+                    comp.components[relay_name] = relay_unit
+                    setattr(comp, relay_name, relay_unit)
+
+                    relay_bindings[endpoint.address] = _RelayBinding(
+                        kind="input",
+                        endpoint_topic=endpoint.address,
+                        relay_in_topic=relay_unit.INPUT.address,
+                        relay_out_topic=relay_unit.OUTPUT.address,
+                        endpoint=endpoint,
+                        relay_unit=relay_unit,
+                    )
+
+                elif isinstance(endpoint, OutputRelay):
+                    relay_name = f"__relay_out_{endpoint_name}"
+                    if relay_name in comp.components:
+                        raise ValueError(
+                            f"{comp.address} already defines component '{relay_name}'."
+                        )
+
+                    relay_unit = _CollectionRelayUnit(output_relay_settings(endpoint))
+                    relay_unit._set_name(relay_name)
+                    relay_unit._set_location(comp.location + [comp.name])
+                    comp.components[relay_name] = relay_unit
+                    setattr(comp, relay_name, relay_unit)
+
+                    relay_bindings[endpoint.address] = _RelayBinding(
+                        kind="output",
+                        endpoint_topic=endpoint.address,
+                        relay_in_topic=relay_unit.INPUT.address,
+                        relay_out_topic=relay_unit.OUTPUT.address,
+                        endpoint=endpoint,
+                        relay_unit=relay_unit,
+                    )
+
+        for component in components.values():
+            if isinstance(component, Collection):
+                crawl_components(component, add_collection_relay_units)
+
         def gather_edges(comp: Component):
             if isinstance(comp, Collection):
                 for from_stream, to_stream in comp.network():
-                    if isinstance(from_stream, Stream):
-                        from_stream = from_stream.address
-                    if isinstance(to_stream, Stream):
-                        to_stream = to_stream.address
-                    if isinstance(to_stream, enum.Enum):
-                        to_stream = to_stream.name
-                    if isinstance(from_stream, enum.Enum):
-                        from_stream = from_stream.name
-                    graph_connections.append((from_stream, to_stream))
+                    graph_connections.append(
+                        (
+                            normalize_topic(from_stream, f"{comp.address}.network"),
+                            normalize_topic(to_stream, f"{comp.address}.network"),
+                        )
+                    )
 
         for component in components.values():
             if isinstance(component, Collection):
                 crawl_components(component, gather_edges)
 
+        if relay_bindings:
+            rewritten_connections: list[tuple[str, str]] = []
+            for from_topic, to_topic in graph_connections:
+                to_binding = relay_bindings.get(to_topic, None)
+                if to_binding is not None and to_binding.kind == "output":
+                    to_topic = to_binding.relay_in_topic
+
+                from_binding = relay_bindings.get(from_topic, None)
+                if from_binding is not None and from_binding.kind == "input":
+                    from_topic = from_binding.relay_out_topic
+
+                rewritten_connections.append((from_topic, to_topic))
+
+            for binding in relay_bindings.values():
+                if binding.kind == "input":
+                    rewritten_connections.append(
+                        (binding.endpoint_topic, binding.relay_in_topic)
+                    )
+                else:
+                    rewritten_connections.append(
+                        (binding.relay_out_topic, binding.endpoint_topic)
+                    )
+
+            graph_connections = rewritten_connections
+
         processes = collect_processes(components.values(), process_components)
 
         for component in components.values():
@@ -149,6 +263,14 @@ def configure_collections(comp: Component):
 
                 crawl_components(component, configure_collections)
 
+        for binding in relay_bindings.values():
+            if isinstance(binding.endpoint, InputRelay):
+                binding.relay_unit.apply_settings(input_relay_settings(binding.endpoint))
+            elif isinstance(binding.endpoint, OutputRelay):
+                binding.relay_unit.apply_settings(
+                    output_relay_settings(binding.endpoint)
+                )
+
         if force_single_process:
             processes = [[u for pu in processes for u in pu]]
 
diff --git a/src/ezmsg/core/collection.py b/src/ezmsg/core/collection.py
@@ -2,8 +2,9 @@
 from collections.abc import Collection as AbstractCollection
 import typing
 from copy import deepcopy
+import warnings
 
-from .stream import Stream
+from .stream import Stream, InputStream, OutputStream
 from .component import ComponentMeta, Component
 from .settings import Settings
 
@@ -34,6 +35,16 @@ def __init__(
             if isinstance(field_value, Component):
                 field_value._set_name(field_name)
                 cls.__components__[field_name] = field_value
+            elif isinstance(field_value, (InputStream, OutputStream)):
+                warnings.warn(
+                    f"{name}.{field_name} uses {type(field_value).__name__} as a "
+                    "Collection boundary endpoint. This behavior is deprecated and "
+                    "will change in a future release. Use InputTopic / OutputTopic "
+                    "for zero-overhead topic shortcuts, or InputRelay / OutputRelay "
+                    "for explicit boundary republishers.",
+                    FutureWarning,
+                    stacklevel=2,
+                )
 
 
 class Collection(Component, metaclass=CollectionMeta):
diff --git a/src/ezmsg/core/relay.py b/src/ezmsg/core/relay.py
@@ -0,0 +1,42 @@
+from copy import deepcopy
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from .settings import Settings
+from .netprotocol import DEFAULT_SHM_SIZE
+from .stream import InputStream, OutputStream
+from .unit import Unit, publisher, subscriber
+
+
+class _RelaySettings(Settings):
+    leaky: bool = False
+    max_queue: int | None = None
+    host: str | None = None
+    port: int | None = None
+    num_buffers: int = 32
+    buf_size: int = DEFAULT_SHM_SIZE
+    force_tcp: bool = False
+    copy_on_forward: bool = True
+
+
+class _CollectionRelayUnit(Unit):
+    SETTINGS = _RelaySettings
+
+    INPUT = InputStream(Any)
+    OUTPUT = OutputStream(Any)
+
+    async def initialize(self) -> None:
+        self.INPUT.leaky = self.SETTINGS.leaky
+        self.INPUT.max_queue = self.SETTINGS.max_queue
+        self.OUTPUT.host = self.SETTINGS.host
+        self.OUTPUT.port = self.SETTINGS.port
+        self.OUTPUT.num_buffers = self.SETTINGS.num_buffers
+        self.OUTPUT.buf_size = self.SETTINGS.buf_size
+        self.OUTPUT.force_tcp = self.SETTINGS.force_tcp
+
+    @subscriber(INPUT)
+    @publisher(OUTPUT)
+    async def relay(self, msg: Any) -> AsyncGenerator:
+        if self.SETTINGS.copy_on_forward:
+            msg = deepcopy(msg)
+        yield self.OUTPUT, msg
diff --git a/src/ezmsg/core/stream.py b/src/ezmsg/core/stream.py
diff --git a/src/ezmsg/core/unit.py b/src/ezmsg/core/unit.py
diff --git a/tests/test_topics.py b/tests/test_topics.py
