From 04dd0e628887d88455e46f68985b4ce401155e42 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 29 Aug 2026 12:21:15 -0700 Subject: [PATCH 01/16] feat: one handler for N same-typed topics in rust native modules A port's launch-line value can now be an array of topic names instead of a single name, so a module can subscribe a group of same-typed topics to one callback that is told which topic each message arrived on. Rust: `InputGroup` port with `recv() -> Option<(usize, T)>` and a `#[input_group(decode = ..., handler = ...)]` derive attribute. Python: `NativeModuleConfig.topic_groups: dict[str, TopicGroup]`. Existing single-topic ports are unaffected. --- dimos/core/native_module.py | 40 ++- dimos/core/test_native_module.py | 49 +++- dimos/core/test_transport_factory.py | 9 + dimos/core/transport_factory.py | 17 ++ native/rust/README.md | 29 +++ native/rust/dimos-module-macros/src/lib.rs | 68 +++-- native/rust/dimos-module/src/host.rs | 8 +- native/rust/dimos-module/src/lib.rs | 4 +- native/rust/dimos-module/src/module.rs | 280 ++++++++++++++++++--- 9 files changed, 442 insertions(+), 62 deletions(-) diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index 84fea6b685..ff75597d82 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -54,13 +54,13 @@ class MyCppModule(NativeModule): import time from typing import IO, Any -from pydantic import Field, model_validator +from pydantic import BaseModel, Field, model_validator from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.global_config import global_config from dimos.core.module import Module, ModuleConfig -from dimos.core.transport_factory import session_config +from dimos.core.transport_factory import channel_for, session_config from dimos.protocol.service.spec import SessionConfig from dimos.utils.logging_config import setup_logger @@ -114,6 +114,18 @@ class LogFormat(enum.Enum): } +class TopicGroup(BaseModel): + """Several same-typed channels a native port fans into one handler. + + The module declares no In stream for these, so nothing here participates in + blueprint autoconnection: the names are resolved to wire channels and handed + to the native process, which subscribes them itself. + """ + + names: list[str] + msg_type: type[Any] | None = None + + class NativeModuleConfig(ModuleConfig): """Configuration for a native subprocess module.""" @@ -131,6 +143,9 @@ class NativeModuleConfig(ModuleConfig): stdin_config: bool = False + # Port name -> the group of channels that port's single handler receives. + topic_groups: dict[str, TopicGroup] = Field(default_factory=dict) + cli_exclude: frozenset[str] = frozenset() cli_name_override: dict[str, str] = Field(default_factory=dict) @@ -201,6 +216,9 @@ class NativeModule(Module): ``DIMOS_TRANSPORT`` env var. With ``stdin_config``, the topics, config, publisher QoS and session settings also arrive as one JSON line on stdin. + A port named in ``topic_groups`` gets a list of channels instead of one, so + several same-typed topics reach a single native handler. + The native process should parse whichever it uses and pub/sub on the given topics directly. On ``stop()``, the process receives SIGTERM. """ @@ -258,16 +276,17 @@ def _session(self) -> SessionConfig: # A blueprint builds its config before global config is settled. return pinned.rebased() - def _argv(self, topics: dict[str, str]) -> list[str]: + def _argv(self, topics: dict[str, str | list[str]]) -> list[str]: """The command line the native process is spawned with.""" cmd = [self.config.executable] for name, topic_str in topics.items(): - cmd.extend([f"--{name}", topic_str]) + joined = ",".join(topic_str) if isinstance(topic_str, list) else topic_str + cmd.extend([f"--{name}", joined]) cmd.extend(self.config.to_cli_args()) cmd.extend(self.config.extra_args) return cmd - def _stdin_blob(self, topics: dict[str, str]) -> bytes: + def _stdin_blob(self, topics: dict[str, str | list[str]]) -> bytes: """The JSON line the native process reads its launch from.""" config_dict = self.config.to_config_dict() blob: dict[str, Any] = { @@ -513,8 +532,8 @@ def _maybe_build(self) -> None: duration_sec=round(build_elapsed, 3), ) - def _collect_topics(self) -> dict[str, str]: - topics: dict[str, str] = {} + def _collect_topics(self) -> dict[str, str | list[str]]: + topics: dict[str, str | list[str]] = {} for name in list(self.inputs) + list(self.outputs) + list(self.ios): stream = getattr(self, name, None) if stream is None: @@ -525,6 +544,13 @@ def _collect_topics(self) -> dict[str, str]: channel = getattr(transport, "channel", None) if channel is not None: topics[name] = channel + for port, group in self.config.topic_groups.items(): + if port in topics: + raise ValueError( + f"[{self._module_label}] topic group {port!r} collides with the " + "port of the same name declared as a stream" + ) + topics[port] = [channel_for(n, group.msg_type) for n in group.names] return topics def _collect_output_qos(self) -> dict[str, dict[str, str]]: diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index a58742c3c5..fbf30871bb 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -36,7 +36,7 @@ from dimos.core.core import rpc from dimos.core.global_config import GlobalConfig, TransportBackend from dimos.core.module import Module -from dimos.core.native_module import LogFormat, NativeModule, NativeModuleConfig +from dimos.core.native_module import LogFormat, NativeModule, NativeModuleConfig, TopicGroup from dimos.core.stream import IO, In, Out from dimos.core.transport import LCMTransport, ZenohTransport from dimos.core.transport_factory import make_transport, transport_topic @@ -211,6 +211,53 @@ def test_tf_topic_comes_from_the_declared_port_only() -> None: transport.stop() +def test_a_topic_group_reaches_the_native_process_as_a_list(monkeypatch) -> None: + """A group's channels are resolved without the module declaring a stream each.""" + monkeypatch.setattr(native_module_mod.global_config, "transport", "lcm") + module = StubNativeModule( + executable=_ECHO, + topic_groups={"cams": TopicGroup(names=["/cam0/imu", "/cam1/imu"], msg_type=Imu)}, + ) + try: + topics = module._collect_topics() + assert topics["cams"] == ["/cam0/imu#sensor_msgs.Imu", "/cam1/imu#sensor_msgs.Imu"] + assert module._argv(topics)[1:3] == [ + "--cams", + "/cam0/imu#sensor_msgs.Imu,/cam1/imu#sensor_msgs.Imu", + ] + finally: + module.stop() + + +def test_a_topic_group_cannot_shadow_a_declared_port() -> None: + module = StubNativeModule( + executable=_ECHO, + topic_groups={"cmd_vel": TopicGroup(names=["/other"], msg_type=Twist)}, + ) + transport = LCMTransport("/cmd_vel", Twist) + try: + module.set_transport("cmd_vel", transport) + with pytest.raises(ValueError, match="collides"): + module._collect_topics() + finally: + module.stop() + with contextlib.suppress(Exception): + transport.stop() + + +def test_a_topic_group_is_not_a_native_config_field() -> None: + """The group is wiring, so it belongs in `topics`, not the config struct.""" + module = StubNativeModule( + executable=_ECHO, + topic_groups={"cams": TopicGroup(names=["/cam0/imu"], msg_type=Imu)}, + ) + try: + assert "topic_groups" not in module.config.to_config_dict() + assert "--topic_groups" not in module._argv({}) + finally: + module.stop() + + def test_io_port_publisher_qos_reaches_the_native_process() -> None: module = StubIoModule(executable=_ECHO) transport = ZenohTransport(ZenohTopic("/tf", TFMessage, qos=QOS_NEVER_DROP)) diff --git a/dimos/core/test_transport_factory.py b/dimos/core/test_transport_factory.py index 69e8e7d519..f6491bad09 100644 --- a/dimos/core/test_transport_factory.py +++ b/dimos/core/test_transport_factory.py @@ -25,6 +25,7 @@ ) from dimos.core.transport_factory import ( apply_transport_arg, + channel_for, default_zenoh_qos, make_transport, rpc_backend, @@ -80,6 +81,14 @@ def test_make_transport_zenoh_pickled() -> None: assert t.topic == "dimos/human_input" +@pytest.mark.parametrize("g", [LCM, ZENOH]) +@pytest.mark.parametrize(("name", "msg_type"), [("/camera/color", Image), ("/human_input", None)]) +def test_channel_for_matches_the_transport_it_skips_building( + g: GlobalConfig, name: str, msg_type: type | None +) -> None: + assert channel_for(name, msg_type, g=g) == make_transport(name, msg_type, g=g).channel + + def test_default_zenoh_qos_high_rate_sensor_types_drop() -> None: assert default_zenoh_qos("/camera/color", Image) == QOS_LATEST_WINS diff --git a/dimos/core/transport_factory.py b/dimos/core/transport_factory.py index 597c0c344f..abfb901d2a 100644 --- a/dimos/core/transport_factory.py +++ b/dimos/core/transport_factory.py @@ -27,6 +27,7 @@ pLCMTransport, pZenohTransport, ) +from dimos.protocol.pubsub.impl.lcmpubsub import Topic as LCMTopic from dimos.protocol.pubsub.impl.zenohpubsub import ( QOS_LATEST_WINS, QOS_NEVER_DROP, @@ -112,6 +113,22 @@ def make_transport( return LCMTransport(topic, msg_type) +def channel_for(name: str, msg_type: type | None = None, *, g: GlobalConfig = global_config) -> str: + """The wire channel `make_transport` would land on, without building one. + + Lets a caller name a channel it never subscribes to itself, such as the extra + topics a native module fans into one handler. + """ + use_pickled = msg_type is None or getattr(msg_type, "lcm_encode", None) is None + topic = transport_topic(name, g) + if g.transport == "zenoh": + return str(ZenohTopic(topic, None if use_pickled else msg_type).key_expr) + if use_pickled: + return topic + assert msg_type is not None + return str(LCMTopic(topic, msg_type)) + + def _transport_arg_error(argv: list[str], message: str) -> NoReturn: """Print an argparse-style CLI error for `--transport` and exit(2).""" prog = os.path.basename(argv[0]) if argv else "dimos" diff --git a/native/rust/README.md b/native/rust/README.md index 4c31fc1d81..9d6ad5c995 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -62,6 +62,7 @@ Every transport is compiled into the binary. `run_with_transport` opens the one - `#[derive(Module)]`: on the struct. Required. - `#[module(setup = fn, teardown = fn)]`: on the struct. Both optional. Names methods on `Self`. `setup` runs once before the input dispatch loop starts (use it to spawn background tasks or initialize resources); `teardown` runs once after the loop exits (use it for cleanup). - `#[input(decode = fn, handler = fn)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. +- `#[input_group(decode = fn, handler = fn)]`: on a field of type `InputGroup`, one port fed by several topics of the same message type (see [Topic groups](#topic-groups)). `decode` is required; `handler` defaults to `handle_` and takes `(index, msg)`. - `#[output(encode = fn)]`: on a field of type `Output`. `encode` is required. - `#[io(decode = fn, encode = fn, handler = fn)]`: on a field of type `Io`, a port that publishes to and subscribes on one topic. `decode` and `encode` are required; `handler` defaults to `handle_`. The transports deliver a message back to its own sender, so the handler also sees what the module publishes. Use `#[output]` instead when the module only publishes. - `#[config]`: on one field. The type must be defined with `#[native_config]` (see [Config](#config)). At most one per struct. If absent, `Config` defaults to `dimos_module::NoConfig`. @@ -108,6 +109,34 @@ At runtime `run()` enforces the mapping on the Python payload: deserialization r Field name = port name. Ports map to topics via the stdin JSON; unmapped ports fall back to `/{port}`. +## Topic groups + +A rig with N identical sensors would otherwise need N ports and N near-identical handlers. An `InputGroup` is one port wired to a list of topics that all carry `T`, delivered to one handler in arrival order. Each message is tagged with the index of the topic it arrived on, so the handler can tell the sources apart. + +```rust +#[derive(Module)] +struct MultiCam { + #[input_group(decode = Image::decode)] + cameras: InputGroup, +} + +impl MultiCam { + async fn handle_cameras(&mut self, index: usize, image: Image) { + let topic = self.cameras.topic(index); + } +} +``` + +The Python wrapper supplies the topics with `topic_groups`, keyed by port name: + +```python +MultiCam.blueprint( + topic_groups={"cameras": TopicGroup(names=["/cam0/color", "/cam1/color"], msg_type=Image)}, +) +``` + +Those channels are not declared as `In` streams, so they take no part in blueprint autoconnection: the names are resolved to wire channels and the native process subscribes them itself. On the launch line the port's value is an array rather than a string. A group configured with no topics still claims its port but never yields. + ## Transforms A `#[tf]` field gives a module a view of the transform graph, the Rust counterpart to Python's `tf.get()` and `tf.publish()`. It subscribes to the `tf` topic (mapped like any other port, default `/tf`), buffers each `parent -> child` edge it sees, and answers queries by composing transforms along the shortest path through the graph. diff --git a/native/rust/dimos-module-macros/src/lib.rs b/native/rust/dimos-module-macros/src/lib.rs index 94afe00b31..da468f0afe 100644 --- a/native/rust/dimos-module-macros/src/lib.rs +++ b/native/rust/dimos-module-macros/src/lib.rs @@ -17,7 +17,7 @@ use proc_macro2::TokenStream as TokenStream2; use quote::{format_ident, quote}; use syn::{parse_macro_input, Data, DeriveInput, Field, Fields, Ident, LitStr, Path, Type}; -#[proc_macro_derive(Module, attributes(input, output, io, config, tf, module))] +#[proc_macro_derive(Module, attributes(input, input_group, output, io, config, tf, module))] pub fn derive_module(input: TokenStream) -> TokenStream { let input = parse_macro_input!(input as DeriveInput); match expand(input) { @@ -175,7 +175,7 @@ fn is_option(ty: &Type) -> bool { } const ONE_ATTR_ONLY: &str = "field has multiple module attributes; only one of #[input], \ - #[output], #[io], #[config], #[tf] is allowed"; + #[input_group], #[output], #[io], #[config], #[tf] is allowed"; /// What a `#[tf]` port carries, for the Cargo.toml registry cross-check. const TF_PAYLOAD_TYPE: &str = "TFMessage"; @@ -185,6 +185,10 @@ enum FieldKind { decode: Path, handler: Ident, }, + InputGroup { + decode: Path, + handler: Ident, + }, Output { encode: Path, }, @@ -323,6 +327,9 @@ fn expand(input: DeriveInput) -> syn::Result { FieldKind::Input { decode, .. } => { quote!(#name: builder.input(#name_str, #decode)) } + FieldKind::InputGroup { decode, .. } => { + quote!(#name: builder.input_group(#name_str, #decode)) + } FieldKind::Output { encode } => { quote!(#name: builder.output(#name_str, #encode)) } @@ -336,26 +343,29 @@ fn expand(input: DeriveInput) -> syn::Result { }); // Every port that receives messages gets an arm in the select! loop. - let handled_fields: Vec<(&Ident, &Ident)> = classified + let handle_arms: Vec = classified .iter() - .filter_map(|f| match &f.kind { - FieldKind::Input { handler, .. } | FieldKind::Io { handler, .. } => { - Some((f.name, handler)) + .filter_map(|f| { + let name = f.name; + match &f.kind { + FieldKind::Input { handler, .. } | FieldKind::Io { handler, .. } => Some(quote!( + ::core::option::Option::Some(msg) = self.#name.recv() => { + self.#handler(msg).await + } + )), + FieldKind::InputGroup { handler, .. } => Some(quote!( + ::core::option::Option::Some((index, msg)) = self.#name.recv() => { + self.#handler(index, msg).await + } + )), + _ => None, } - _ => None, }) .collect(); - let handle_body = if handled_fields.is_empty() { + let handle_body = if handle_arms.is_empty() { quote!(::std::future::pending::<()>().await) } else { - let handle_arms = handled_fields.iter().map(|(name, handler)| { - quote!( - ::core::option::Option::Some(msg) = self.#name.recv() => { - self.#handler(msg).await - } - ) - }); quote! { loop { ::tokio::select! { @@ -434,7 +444,7 @@ fn port_decls(classified: &[ClassifiedField], want_input: bool) -> Vec classified .iter() .filter(|f| match f.kind { - FieldKind::Input { .. } => want_input, + FieldKind::Input { .. } | FieldKind::InputGroup { .. } => want_input, FieldKind::Output { .. } => !want_input, // A `#[tf]` field is a port like any other as far as the registry // goes: bake has to put the topic in the host's map or the module @@ -577,6 +587,32 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { .ok_or_else(|| syn::Error::new_spanned(attr, "#[input] requires `decode = ...`"))?; let handler = handler.unwrap_or_else(|| format_ident!("handle_{}", name)); found = Some(FieldKind::Input { decode, handler }); + } else if path.is_ident("input_group") { + if found.is_some() { + return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); + } + let mut decode: Option = None; + let mut handler: Option = None; + attr.parse_nested_meta(|meta| { + if meta.path.is_ident("decode") { + decode = Some(meta.value()?.parse()?); + } else if meta.path.is_ident("handler") { + handler = Some(meta.value()?.parse()?); + } else if meta.path.is_ident("msg") { + msg = Some(meta.value()?.parse::()?.value()); + } else { + return Err(meta.error( + "unrecognized #[input_group] argument; expected `decode = ...`, \ + `handler = ...` or `msg = ...`", + )); + } + Ok(()) + })?; + let decode = decode.ok_or_else(|| { + syn::Error::new_spanned(attr, "#[input_group] requires `decode = ...`") + })?; + let handler = handler.unwrap_or_else(|| format_ident!("handle_{}", name)); + found = Some(FieldKind::InputGroup { decode, handler }); } else if path.is_ident("output") { if found.is_some() { return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); diff --git a/native/rust/dimos-module/src/host.rs b/native/rust/dimos-module/src/host.rs index cc7982b7e5..07d219950c 100644 --- a/native/rust/dimos-module/src/host.rs +++ b/native/rust/dimos-module/src/host.rs @@ -19,7 +19,6 @@ //! [`host_main`]. Everything a host needs beyond that lives here, so the //! generated crate stays a table of module entries and three `include_str!`s. -use std::collections::HashMap; use std::future::Future; use std::io; use std::pin::Pin; @@ -33,7 +32,7 @@ use tracing::{error, info, warn}; use crate::lcm::LcmTransport; use crate::module::{ init_tracing, log_wiring, parse_config_value, read_launch_config, run_module_core, - validate_config, Module, + validate_config, Module, Topics, }; use crate::transport::{SharedTransport, Transport}; use crate::zenoh::{ZenohTransport, SESSION_KEY}; @@ -52,7 +51,7 @@ type RunFn = Box, watch::Receiver) -> Modu /// once the transport is open. Produced before anything is spawned so a bad /// config kills the host instead of half-starting it. struct Prepared { - topics: HashMap, + topics: Topics, run: RunFn, } @@ -450,7 +449,7 @@ fn run_host_fallible(spec: &HostSpec) -> io::Result<()> { let prepared = prepare_all(spec, &stdin)?; let known: Vec = prepared .iter() - .flat_map(|p| p.topics.values().cloned()) + .flat_map(|p| p.topics.channels().cloned()) .collect(); let suppress = resolve_suppress(spec, &stdin, &known)?; if let Ok(transport) = std::env::var("DIMOS_TRANSPORT") { @@ -587,6 +586,7 @@ fn supervise( #[cfg(test)] mod tests { use super::*; + use std::collections::HashMap; #[test] fn thread_name_is_prefixed_and_capped() { diff --git a/native/rust/dimos-module/src/lib.rs b/native/rust/dimos-module/src/lib.rs index a1a6f6efd6..5b108b95ae 100644 --- a/native/rust/dimos-module/src/lib.rs +++ b/native/rust/dimos-module/src/lib.rs @@ -29,7 +29,9 @@ pub mod zenoh; pub use dimos_module_macros::{native_config, Module}; pub use host::{host_main, HostSpec, ModuleEntry}; pub use lcm::LcmTransport; -pub use module::{run, Builder, Input, Io, Module, ModuleConfig, NativeConfig, NoConfig, Output}; +pub use module::{ + run, Builder, Input, InputGroup, Io, Module, ModuleConfig, NativeConfig, NoConfig, Output, +}; pub use tf::{Lookup, Tf, Transform}; pub use transport::{SharedTransport, Transport}; pub use workers::worker_pool; diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index e9a5f014fa..c526af6a20 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -74,9 +74,13 @@ pub(crate) trait Route: Send + Sync { fn try_dispatch(&self, data: &[u8]); } +/// Decodes a frame into whatever the port's channel carries. Boxed because a +/// group route tags the message with the index of the topic it arrived on. +type Decode = Box io::Result + Send + Sync>; + struct TypedRoute { topic: String, - decode: fn(&[u8]) -> io::Result, + decode: Decode, sender: mpsc::Sender, drop_count: AtomicU64, last_log_ns: AtomicU64, @@ -120,6 +124,34 @@ impl Input { } } +/// Several topics of one message type, fanned into a single handler. +/// +/// The launch line gives the port an array of topics instead of one, and every +/// message is tagged with the index of the topic it arrived on so a handler can +/// tell a rig's cameras apart. A group configured with no topics never yields. +pub struct InputGroup { + pub topics: Vec, + receiver: mpsc::Receiver<(usize, T)>, +} + +impl InputGroup { + pub async fn recv(&mut self) -> Option<(usize, T)> { + self.receiver.recv().await + } + + pub fn topic(&self, index: usize) -> &str { + &self.topics[index] + } + + pub fn len(&self) -> usize { + self.topics.len() + } + + pub fn is_empty(&self) -> bool { + self.topics.is_empty() + } +} + #[derive(Clone)] pub struct Output { pub topic: String, @@ -164,19 +196,64 @@ pub(crate) async fn publish_encoded( .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "background task gone")) } +/// The port-to-topic wiring from the launch line. A port names one topic, or an +/// array of same-typed topics that one handler sees. +#[derive(Clone, Debug, Default)] +pub(crate) struct Topics { + single: HashMap, + grouped: HashMap>, +} + +impl Topics { + fn ports(&self) -> BTreeSet<&String> { + self.single.keys().chain(self.grouped.keys()).collect() + } + + /// Every wire channel the module touches, groups flattened. + pub(crate) fn channels(&self) -> impl Iterator { + self.single.values().chain(self.grouped.values().flatten()) + } +} + +fn parse_topics(json: &serde_json::Value) -> io::Result { + let mut topics = Topics::default(); + let Some(table) = json.get("topics").and_then(|v| v.as_object()) else { + return Ok(topics); + }; + for (port, value) in table { + let invalid = |detail: &str| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("topic for port '{port}' {detail}"), + ) + }; + match value { + serde_json::Value::String(topic) => { + topics.single.insert(port.clone(), topic.clone()); + } + serde_json::Value::Array(items) => { + let group = items + .iter() + .map(|item| { + item.as_str() + .map(str::to_string) + .ok_or_else(|| invalid("is a group, so every entry must be a string")) + }) + .collect::>>()?; + topics.grouped.insert(port.clone(), group); + } + _ => return Err(invalid("must be a string or an array of strings")), + } + } + Ok(topics) +} + /// Extract `(topics, config)` from an already-parsed config object. `run` /// parses the line once and also reads `qos` from it, so this takes the value. pub(crate) fn parse_config_value( json: &serde_json::Value, -) -> io::Result<(HashMap, C)> { - let mut topics = HashMap::new(); - if let Some(t) = json.get("topics").and_then(|v| v.as_object()) { - for (port, topic) in t { - if let Some(s) = topic.as_str() { - topics.insert(port.clone(), s.to_string()); - } - } - } +) -> io::Result<(Topics, C)> { + let topics = parse_topics(json)?; let config_value = json.get("config").ok_or_else(|| { io::Error::new( @@ -305,7 +382,7 @@ pub trait Module: Sized + Send + 'static { } pub struct Builder { - topics: HashMap, + topics: Topics, // Every port the module asked for a topic, matched against topics after build. requested: BTreeSet, routes: HashMap>>, @@ -315,7 +392,7 @@ pub struct Builder { } impl Builder { - pub(crate) fn new(topics: HashMap) -> Self { + pub(crate) fn new(topics: Topics) -> Self { Self { topics, requested: BTreeSet::new(), @@ -328,15 +405,21 @@ impl Builder { fn topic_for(&mut self, port: &str) -> String { self.requested.insert(port.to_string()); self.topics + .single .get(port) .cloned() .unwrap_or_else(|| format!("/{port}")) } + fn group_for(&mut self, port: &str) -> Vec { + self.requested.insert(port.to_string()); + self.topics.grouped.get(port).cloned().unwrap_or_default() + } + // A mismatch is dead wiring: an unclaimed topic reaches no port, and an // unsent one leaves the port on a fallback name nothing else publishes to. pub(crate) fn enforce_topics_match_ports(&self) -> io::Result<()> { - let provided: BTreeSet<&String> = self.topics.keys().collect(); + let provided = self.topics.ports(); let requested: BTreeSet<&String> = self.requested.iter().collect(); if provided == requested { return Ok(()); @@ -351,22 +434,31 @@ impl Builder { )) } - fn add_route( + fn push_route( &mut self, topic: &str, - decode: fn(&[u8]) -> io::Result, - ) -> mpsc::Receiver { - let (tx, rx) = mpsc::channel(INPUT_CHANNEL_CAPACITY); + decode: Decode, + sender: mpsc::Sender, + ) { self.routes .entry(topic.to_string()) .or_default() .push(Box::new(TypedRoute { topic: topic.to_string(), decode, - sender: tx, + sender, drop_count: AtomicU64::new(0), last_log_ns: AtomicU64::new(0), })); + } + + fn add_route( + &mut self, + topic: &str, + decode: fn(&[u8]) -> io::Result, + ) -> mpsc::Receiver { + let (tx, rx) = mpsc::channel(INPUT_CHANNEL_CAPACITY); + self.push_route(topic, Box::new(decode), tx); rx } @@ -386,6 +478,22 @@ impl Builder { Input { topic, receiver } } + /// A port wired to every topic the launch line lists under `port`, all of + /// one message type, delivered to one handler in arrival order. + pub fn input_group( + &mut self, + port: &str, + decode: fn(&[u8]) -> io::Result, + ) -> InputGroup { + let topics = self.group_for(port); + let (sender, receiver) = mpsc::channel(INPUT_CHANNEL_CAPACITY); + for (index, topic) in topics.iter().enumerate() { + let tag = Box::new(move |bytes: &[u8]| decode(bytes).map(|msg| (index, msg))); + self.push_route(topic, tag, sender.clone()); + } + InputGroup { topics, receiver } + } + pub fn output(&mut self, port: &str, encode: fn(&T) -> Vec) -> Output { let topic = self.topic_for(port); let sender = self.add_publisher(&topic); @@ -514,7 +622,7 @@ where /// and by the baked host, which drives several of these on one transport. pub(crate) async fn run_module_core( transport: Arc, - topics: HashMap, + topics: Topics, config: M::Config, mut shutdown: watch::Receiver, ) -> io::Result<()> @@ -551,10 +659,13 @@ where /// Log the resolved wiring of a module, tagged with whatever the operator sees /// in `ps`: the executable for a lone module, the module id inside a host. -pub(crate) fn log_wiring(exe: &str, topics: &HashMap, config: &C) { - for (port, topic) in topics { +pub(crate) fn log_wiring(exe: &str, topics: &Topics, config: &C) { + for (port, topic) in &topics.single { info!(exe = %exe, port = %port, topic = %topic, "topic mapping"); } + for (port, group) in &topics.grouped { + info!(exe = %exe, port = %port, topics = ?group, "topic group mapping"); + } info!(exe = %exe, config = ?config, "config loaded"); } @@ -599,9 +710,7 @@ mod tests { /// Parse a raw config line the way `run` does, for exercising /// `parse_config_value` from the string form the coordinator sends. - fn parse_config_json( - line: &str, - ) -> io::Result<(HashMap, C)> { + fn parse_config_json(line: &str) -> io::Result<(Topics, C)> { parse_config_value(&parse_launch_config(line)?) } @@ -726,8 +835,8 @@ mod tests { fn parses_topics_and_config() { let json = r#"{"topics": {"data": "/foo/data", "confirm": "/foo/confirm"}, "config": {"value": 42, "name": "hello"}}"#; let (topics, config) = parse_config_json::(json).unwrap(); - assert_eq!(topics["data"], "/foo/data"); - assert_eq!(topics["confirm"], "/foo/confirm"); + assert_eq!(topics.single["data"], "/foo/data"); + assert_eq!(topics.single["confirm"], "/foo/confirm"); assert_eq!( config, TestConfig { @@ -798,7 +907,7 @@ mod tests { fn missing_topics_field_gives_empty_map() { let json = r#"{"config": {"value": 1, "name": "x"}}"#; let (topics, _config) = parse_config_json::(json).unwrap(); - assert!(topics.is_empty()); + assert!(topics.ports().is_empty()); } #[test] @@ -887,11 +996,24 @@ mod tests { // topic_for fallback - fn topics(pairs: &[(&str, &str)]) -> HashMap { - pairs - .iter() - .map(|(p, t)| (p.to_string(), t.to_string())) - .collect() + fn topics(pairs: &[(&str, &str)]) -> Topics { + Topics { + single: pairs + .iter() + .map(|(port, topic)| (port.to_string(), topic.to_string())) + .collect(), + grouped: HashMap::new(), + } + } + + fn grouped_topics(port: &str, group: &[&str]) -> Topics { + Topics { + single: HashMap::new(), + grouped: HashMap::from([( + port.to_string(), + group.iter().map(|t| t.to_string()).collect(), + )]), + } } fn builder_with_topics(pairs: &[(&str, &str)]) -> Builder { @@ -931,6 +1053,63 @@ mod tests { assert_eq!(output.topic, "/robot/cmd_vel"); } + // input groups + + #[test] + fn a_port_given_an_array_of_topics_becomes_a_group() { + let json = r#"{"topics": {"cams": ["/cam0", "/cam1"], "odom": "/odom"}, "config": null}"#; + let (topics, _config) = parse_config_json::<()>(json).unwrap(); + assert_eq!(topics.grouped["cams"], ["/cam0", "/cam1"]); + assert_eq!(topics.single["odom"], "/odom"); + } + + #[test] + fn a_group_entry_that_is_not_a_string_is_rejected() { + let json = r#"{"topics": {"cams": ["/cam0", 7]}, "config": null}"#; + let err = parse_config_json::<()>(json).expect_err("a non-string group entry is invalid"); + assert!(err.to_string().contains("cams"), "{err}"); + } + + #[test] + fn a_group_subscribes_every_topic_it_was_given() { + let mut builder = Builder::new(grouped_topics("cams", &["/cam0", "/cam1"])); + let group = builder.input_group("cams", |b| Ok(b.to_vec())); + assert_eq!(group.topics, ["/cam0", "/cam1"]); + assert_eq!(builder.routes.get("/cam0").map(Vec::len), Some(1)); + assert_eq!(builder.routes.get("/cam1").map(Vec::len), Some(1)); + builder.enforce_topics_match_ports().expect("group claimed"); + } + + /// One handler sees every topic, so the index is the only thing that says + /// which camera of a rig a frame came from. + #[tokio::test] + async fn a_group_tags_each_message_with_its_topic_index() { + let mut builder = Builder::new(grouped_topics("cams", &["/cam0", "/cam1"])); + let mut group = builder.input_group("cams", |b| Ok(b.to_vec())); + + builder.routes["/cam1"][0].try_dispatch(b"second"); + builder.routes["/cam0"][0].try_dispatch(b"first"); + + assert_eq!( + group.recv().await.expect("cam1 frame"), + (1, b"second".to_vec()) + ); + assert_eq!( + group.recv().await.expect("cam0 frame"), + (0, b"first".to_vec()) + ); + assert_eq!(group.topic(1), "/cam1"); + } + + #[test] + fn an_empty_group_still_claims_its_port() { + let mut builder = Builder::new(grouped_topics("cams", &[])); + let group = builder.input_group("cams", |b| Ok(b.to_vec())); + assert!(group.is_empty()); + assert!(builder.routes.is_empty()); + builder.enforce_topics_match_ports().expect("group claimed"); + } + #[test] fn topics_matching_ports_exactly_pass() { let mut builder = builder_with_topics(&[("cmd", "/robot/cmd"), ("odom", "/robot/odom")]); @@ -1241,7 +1420,7 @@ mod tests { let (tx, _rx) = mpsc::channel::>(1); let route = TypedRoute { topic: "/test".to_string(), - decode: |b| Ok(b.to_vec()), + decode: Box::new(|b: &[u8]| Ok(b.to_vec())), sender: tx, drop_count: AtomicU64::new(0), last_log_ns: AtomicU64::new(0), @@ -1284,6 +1463,41 @@ mod tests { } } + #[derive(crate::Module)] + struct Rig { + #[input_group(decode = decode)] + cams: crate::InputGroup, + #[output(encode = encode)] + seen: crate::Output, + } + + impl Rig { + async fn handle_cams(&mut self, index: usize, msg: Msg) { + let mut tagged = vec![index as u8]; + tagged.extend(msg.0); + self.seen.publish(&Msg(tagged)).await.expect("publish"); + } + } + + #[tokio::test] + async fn input_group_field_hands_its_handler_the_topic_index() { + let mut builder = Builder::new(Topics { + single: HashMap::from([("seen".to_string(), "/seen".to_string())]), + grouped: HashMap::from([( + "cams".to_string(), + vec!["/cam0".to_string(), "/cam1".to_string()], + )]), + }); + let mut rig = Rig::build(&mut builder, NoConfig); + + builder.routes["/cam1"][0].try_dispatch(b"frame"); + builder.routes.clear(); + rig.handle().await; + + let (_, rx) = &mut builder.outputs[0]; + assert_eq!(rx.recv().await.expect("handler output"), b"\x01frame"); + } + #[tokio::test] async fn io_field_is_wired_to_its_handler_and_can_publish() { let mut builder = Builder::new(topics(&[("cmd", "/robot/cmd")])); From de946a93872fb1b54f6d4243f7ec937b24250ecd Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 29 Aug 2026 14:45:10 -0700 Subject: [PATCH 02/16] Make module topic groups transport-neutral stream groups Group entries were written as backend topic strings (`/cam0/color`), which reads as LCM-specific and bypasses the stream-name vocabulary the rest of a blueprint uses. They are now stream names as they read after remapping and namespacing, resolved to wire channels the same way a declared stream name is, with a leading `/` rejected outright. --- dimos/core/module.py | 113 ++++++++++++++---- dimos/core/native_module.py | 25 +--- dimos/core/test_async_module_stream_groups.py | 61 ++++++++++ dimos/core/test_native_module.py | 26 ++-- native/rust/README.md | 21 +++- 5 files changed, 191 insertions(+), 55 deletions(-) create mode 100644 dimos/core/test_async_module_stream_groups.py diff --git a/dimos/core/module.py b/dimos/core/module.py index f5aab99421..acce3a8a69 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -31,7 +31,7 @@ get_type_hints, ) -from pydantic import Field +from pydantic import BaseModel, Field, field_validator from reactivex.disposable import CompositeDisposable, Disposable from dimos.core.core import T, rpc @@ -41,7 +41,7 @@ from dimos.core.resource import CompositeResource from dimos.core.rpc_client import RpcCall from dimos.core.stream import IO, In, Out, RemoteOut, Transport -from dimos.core.transport_factory import rpc_backend +from dimos.core.transport_factory import make_transport, rpc_backend from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, DEFAULT_RPC_TIMEOUTS, RPCSpec from dimos.protocol.service.spec import BaseConfig, Configurable from dimos.protocol.tf.tf import TF @@ -103,6 +103,34 @@ def get_loop() -> tuple[asyncio.AbstractEventLoop, threading.Thread | None]: Deployment = Literal["python", "docker"] +class StreamGroup(BaseModel): + """Several same-typed streams that one port fans into a single handler. + + `names` are stream names as they read after remapping and namespacing — + `left_cam`, `robot1/lidar` — never a backend topic string. Each is resolved + to a wire channel the same way a declared stream's name is, so the group + says nothing about which transport is in use. + + No `In` stream is declared for these, so they take no part in blueprint + autoconnection. A python module subscribes them itself; a native module + hands them to its subprocess, which does the subscribing. + """ + + names: list[str] + msg_type: type[Any] | None = None + + @field_validator("names") + @classmethod + def _reject_topic_strings(cls, names: list[str]) -> list[str]: + leading_slash = [n for n in names if n.startswith("/")] + if leading_slash: + raise ValueError( + f"stream group names must be stream names, not topics: {leading_slash} " + "start with '/' (use `left_cam`, not `/left_cam`)" + ) + return names + + class ModuleConfig(BaseConfig): rpc_transport: type[RPCSpec] = Field(default_factory=rpc_backend) default_rpc_timeout: float = DEFAULT_RPC_TIMEOUT @@ -113,6 +141,8 @@ class ModuleConfig(BaseConfig): # once (see BlueprintAtom.instance_name). Changes the RPC topic prefix # from the class name to this name. instance_name: str | None = None + # Port name -> the group of streams that port's single handler receives. + stream_groups: dict[str, StreamGroup] = Field(default_factory=dict) g: GlobalConfig = global_config @@ -190,6 +220,7 @@ def build(self) -> None: def start(self) -> None: self._start_main() self._auto_bind_handlers() + self._bind_stream_groups() @rpc def stop(self) -> None: @@ -666,6 +697,30 @@ def _auto_bind_handlers(self) -> None: # backpressure. self.process_observable(in_stream.pure_observable(), handler) + def _bind_stream_groups(self) -> None: + """For each `stream_groups` port with an `async def handle_`, subscribe + every stream in the group into that one handler. + + A native module defines no such method — its subprocess subscribes instead — + so this is a no-op there. + """ + for port, group in self.config.stream_groups.items(): + handler = getattr(self, f"handle_{port}", None) + if handler is None: + continue + if hasattr(handler, "aio"): + handler = handler.aio.__get__(self, type(self)) + if not inspect.iscoroutinefunction(handler): + raise TypeError( + f"{type(self).__name__}.handle_{port} must be `async def` " + "(stream groups have no sync path)" + ) + on_msg, dispatcher_disp = self._make_keyed_dispatch(handler) + self.register_disposable(dispatcher_disp) + for index, name in enumerate(group.names): + transport = make_transport(name, group.msg_type) + self.register_disposable(Disposable(transport.subscribe(partial(on_msg, index)))) + def _make_async_dispatch( self, async_handler: Callable[[Any], Any] ) -> tuple[Callable[[Any], None], "DisposableBase"]: @@ -680,45 +735,63 @@ def _make_async_dispatch( message is kept (LATEST policy). - The returned Disposable cancels the dispatcher task. """ + + async def single(_: int, msg: Any) -> None: + await async_handler(msg) + + on_msg, disposable = self._make_keyed_dispatch(single) + return partial(on_msg, 0), disposable + + def _make_keyed_dispatch( + self, async_handler: Callable[[int, Any], Any] + ) -> tuple[Callable[[int, Any], None], "DisposableBase"]: + """`_make_async_dispatch` generalized to a mailbox keyed by sender. + + The mailbox keeps the latest unprocessed message *per key* and drains them + oldest-waiting first, so one chatty topic in a group cannot starve its + siblings the way a single shared slot would. One key is the degenerate case + and behaves exactly like a single LATEST slot. + """ loop = self._loop if loop is None or not loop.is_running(): raise RuntimeError(f"{type(self).__name__}._loop is not running") - async def _bootstrap() -> tuple[asyncio.Event, dict[str, Any], asyncio.Task[None]]: + async def _bootstrap() -> tuple[asyncio.Event, dict[int, Any], asyncio.Task[None]]: event = asyncio.Event() - slot: dict[str, Any] = {"value": None, "has_value": False} + # Insertion-ordered, and re-assigning an existing key keeps its + # position, so a waiting topic holds its place while its value ages up. + pending: dict[int, Any] = {} async def dispatcher() -> None: try: while True: await event.wait() event.clear() - if not slot["has_value"]: - continue - msg = slot["value"] - slot["value"] = None - slot["has_value"] = False - try: - await async_handler(msg) - except asyncio.CancelledError: - raise - except BaseException as e: - self._log_async_handler_exception(e) + while pending: + key = next(iter(pending)) + msg = pending.pop(key) + try: + await async_handler(key, msg) + except asyncio.CancelledError: + raise + except BaseException as e: + self._log_async_handler_exception(e) except asyncio.CancelledError: return - return event, slot, asyncio.create_task(dispatcher()) + return event, pending, asyncio.create_task(dispatcher()) - event, slot, task = asyncio.run_coroutine_threadsafe(_bootstrap(), loop).result(timeout=5.0) + event, pending, task = asyncio.run_coroutine_threadsafe(_bootstrap(), loop).result( + timeout=5.0 + ) - def on_msg(msg: Any) -> None: + def on_msg(key: int, msg: Any) -> None: loop_now = self._loop if loop_now is None or not loop_now.is_running(): return def _set() -> None: - slot["value"] = msg - slot["has_value"] = True + pending[key] = msg event.set() loop_now.call_soon_threadsafe(_set) diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index ff75597d82..ec728a7d27 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -54,7 +54,7 @@ class MyCppModule(NativeModule): import time from typing import IO, Any -from pydantic import BaseModel, Field, model_validator +from pydantic import Field, model_validator from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc @@ -114,18 +114,6 @@ class LogFormat(enum.Enum): } -class TopicGroup(BaseModel): - """Several same-typed channels a native port fans into one handler. - - The module declares no In stream for these, so nothing here participates in - blueprint autoconnection: the names are resolved to wire channels and handed - to the native process, which subscribes them itself. - """ - - names: list[str] - msg_type: type[Any] | None = None - - class NativeModuleConfig(ModuleConfig): """Configuration for a native subprocess module.""" @@ -143,9 +131,6 @@ class NativeModuleConfig(ModuleConfig): stdin_config: bool = False - # Port name -> the group of channels that port's single handler receives. - topic_groups: dict[str, TopicGroup] = Field(default_factory=dict) - cli_exclude: frozenset[str] = frozenset() cli_name_override: dict[str, str] = Field(default_factory=dict) @@ -216,8 +201,8 @@ class NativeModule(Module): ``DIMOS_TRANSPORT`` env var. With ``stdin_config``, the topics, config, publisher QoS and session settings also arrive as one JSON line on stdin. - A port named in ``topic_groups`` gets a list of channels instead of one, so - several same-typed topics reach a single native handler. + A port named in ``stream_groups`` gets a list of channels instead of one, so + several same-typed streams reach a single native handler. The native process should parse whichever it uses and pub/sub on the given topics directly. On ``stop()``, the process receives SIGTERM. @@ -544,10 +529,10 @@ def _collect_topics(self) -> dict[str, str | list[str]]: channel = getattr(transport, "channel", None) if channel is not None: topics[name] = channel - for port, group in self.config.topic_groups.items(): + for port, group in self.config.stream_groups.items(): if port in topics: raise ValueError( - f"[{self._module_label}] topic group {port!r} collides with the " + f"[{self._module_label}] stream group {port!r} collides with the " "port of the same name declared as a stream" ) topics[port] = [channel_for(n, group.msg_type) for n in group.names] diff --git a/dimos/core/test_async_module_stream_groups.py b/dimos/core/test_async_module_stream_groups.py new file mode 100644 index 0000000000..a0120e5815 --- /dev/null +++ b/dimos/core/test_async_module_stream_groups.py @@ -0,0 +1,61 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from queue import Queue + +import pytest + +from dimos.core.coordination.module_coordinator import ModuleCoordinator +from dimos.core.module import Module, StreamGroup +from dimos.core.stream import Out +from dimos.core.transport_factory import make_transport + + +class FanInModule(Module): + """One handler for two same-typed streams; echoes which one it came from.""" + + tagged: Out[int] + + async def handle_sensors(self, index: int, value: int) -> None: + self.tagged.publish(index * 100 + value) + + +@pytest.fixture +def start_fan_in_module(each_transport): + blueprint = FanInModule.blueprint(stream_groups={"sensors": StreamGroup(names=["s0", "s1"])}) + coordinator = ModuleCoordinator.build(blueprint) + yield + coordinator.stop() + + +@pytest.fixture +def group_transports(each_transport): + transports = [make_transport("s0"), make_transport("s1"), make_transport("tagged")] + for transport in transports: + transport.start() + yield transports + for transport in transports: + transport.stop() + + +def test_stream_group_tags_each_message_with_its_stream(start_fan_in_module, group_transports): + s0, s1, tagged = group_transports + queue: Queue[int] = Queue() + tagged.subscribe(queue.put) + + s0.publish(7) + assert queue.get(timeout=1.0) == 7 + + s1.publish(7) + assert queue.get(timeout=1.0) == 107 diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index fbf30871bb..551388b5f2 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -35,8 +35,8 @@ from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.core import rpc from dimos.core.global_config import GlobalConfig, TransportBackend -from dimos.core.module import Module -from dimos.core.native_module import LogFormat, NativeModule, NativeModuleConfig, TopicGroup +from dimos.core.module import Module, StreamGroup +from dimos.core.native_module import LogFormat, NativeModule, NativeModuleConfig from dimos.core.stream import IO, In, Out from dimos.core.transport import LCMTransport, ZenohTransport from dimos.core.transport_factory import make_transport, transport_topic @@ -211,12 +211,12 @@ def test_tf_topic_comes_from_the_declared_port_only() -> None: transport.stop() -def test_a_topic_group_reaches_the_native_process_as_a_list(monkeypatch) -> None: +def test_a_stream_group_reaches_the_native_process_as_a_list(monkeypatch) -> None: """A group's channels are resolved without the module declaring a stream each.""" monkeypatch.setattr(native_module_mod.global_config, "transport", "lcm") module = StubNativeModule( executable=_ECHO, - topic_groups={"cams": TopicGroup(names=["/cam0/imu", "/cam1/imu"], msg_type=Imu)}, + stream_groups={"cams": StreamGroup(names=["cam0/imu", "cam1/imu"], msg_type=Imu)}, ) try: topics = module._collect_topics() @@ -229,10 +229,16 @@ def test_a_topic_group_reaches_the_native_process_as_a_list(monkeypatch) -> None module.stop() -def test_a_topic_group_cannot_shadow_a_declared_port() -> None: +def test_a_stream_group_rejects_a_backend_topic_string() -> None: + """Names are stream names, so a leading slash is a transport leaking in.""" + with pytest.raises(ValidationError, match="not topics"): + StreamGroup(names=["/cam0/imu"], msg_type=Imu) + + +def test_a_stream_group_cannot_shadow_a_declared_port() -> None: module = StubNativeModule( executable=_ECHO, - topic_groups={"cmd_vel": TopicGroup(names=["/other"], msg_type=Twist)}, + stream_groups={"cmd_vel": StreamGroup(names=["other"], msg_type=Twist)}, ) transport = LCMTransport("/cmd_vel", Twist) try: @@ -245,15 +251,15 @@ def test_a_topic_group_cannot_shadow_a_declared_port() -> None: transport.stop() -def test_a_topic_group_is_not_a_native_config_field() -> None: +def test_a_stream_group_is_not_a_native_config_field() -> None: """The group is wiring, so it belongs in `topics`, not the config struct.""" module = StubNativeModule( executable=_ECHO, - topic_groups={"cams": TopicGroup(names=["/cam0/imu"], msg_type=Imu)}, + stream_groups={"cams": StreamGroup(names=["cam0/imu"], msg_type=Imu)}, ) try: - assert "topic_groups" not in module.config.to_config_dict() - assert "--topic_groups" not in module._argv({}) + assert "stream_groups" not in module.config.to_config_dict() + assert "--stream_groups" not in module._argv({}) finally: module.stop() diff --git a/native/rust/README.md b/native/rust/README.md index 9d6ad5c995..1d2240493d 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -62,7 +62,7 @@ Every transport is compiled into the binary. `run_with_transport` opens the one - `#[derive(Module)]`: on the struct. Required. - `#[module(setup = fn, teardown = fn)]`: on the struct. Both optional. Names methods on `Self`. `setup` runs once before the input dispatch loop starts (use it to spawn background tasks or initialize resources); `teardown` runs once after the loop exits (use it for cleanup). - `#[input(decode = fn, handler = fn)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. -- `#[input_group(decode = fn, handler = fn)]`: on a field of type `InputGroup`, one port fed by several topics of the same message type (see [Topic groups](#topic-groups)). `decode` is required; `handler` defaults to `handle_` and takes `(index, msg)`. +- `#[input_group(decode = fn, handler = fn)]`: on a field of type `InputGroup`, one port fed by several topics of the same message type (see [Stream groups](#stream-groups)). `decode` is required; `handler` defaults to `handle_` and takes `(index, msg)`. - `#[output(encode = fn)]`: on a field of type `Output`. `encode` is required. - `#[io(decode = fn, encode = fn, handler = fn)]`: on a field of type `Io`, a port that publishes to and subscribes on one topic. `decode` and `encode` are required; `handler` defaults to `handle_`. The transports deliver a message back to its own sender, so the handler also sees what the module publishes. Use `#[output]` instead when the module only publishes. - `#[config]`: on one field. The type must be defined with `#[native_config]` (see [Config](#config)). At most one per struct. If absent, `Config` defaults to `dimos_module::NoConfig`. @@ -109,7 +109,7 @@ At runtime `run()` enforces the mapping on the Python payload: deserialization r Field name = port name. Ports map to topics via the stdin JSON; unmapped ports fall back to `/{port}`. -## Topic groups +## Stream groups A rig with N identical sensors would otherwise need N ports and N near-identical handlers. An `InputGroup` is one port wired to a list of topics that all carry `T`, delivered to one handler in arrival order. Each message is tagged with the index of the topic it arrived on, so the handler can tell the sources apart. @@ -127,15 +127,26 @@ impl MultiCam { } ``` -The Python wrapper supplies the topics with `topic_groups`, keyed by port name: +The Python wrapper supplies the sources with `stream_groups`, keyed by port name: ```python MultiCam.blueprint( - topic_groups={"cameras": TopicGroup(names=["/cam0/color", "/cam1/color"], msg_type=Image)}, + stream_groups={"cameras": StreamGroup(names=["left_cam", "right_cam"], msg_type=Image)}, ) ``` -Those channels are not declared as `In` streams, so they take no part in blueprint autoconnection: the names are resolved to wire channels and the native process subscribes them itself. On the launch line the port's value is an array rather than a string. A group configured with no topics still claims its port but never yields. +`names` are stream names as they read after remapping and namespacing, not backend topics — a leading `/` is rejected. Python resolves each to a wire channel the same way it resolves a declared stream's name, so the same blueprint runs unchanged over LCM or zenoh. + +Those streams are not declared as `In` ports, so they take no part in blueprint autoconnection: the names are resolved to wire channels and the native process subscribes them itself. On the launch line the port's value is an array rather than a string. A group configured with no names still claims its port but never yields. + +`stream_groups` lives on `ModuleConfig`, so a plain Python `Module` takes the same field. There the module subscribes the group itself and dispatches to `async def handle_(self, index, msg)`, matching the Rust handler signature: + +```python +class MultiCam(Module): + async def handle_cameras(self, index: int, image: Image) -> None: ... +``` + +The whole group shares one dispatcher, so the handler is never re-entered, and its mailbox holds the latest unprocessed message per stream rather than one slot for the group — a chatty camera cannot starve the others. ## Transforms From b64ae1db65a980d412a9df89a188d1917d6ab27b Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 29 Aug 2026 15:17:13 -0700 Subject: [PATCH 03/16] Wire stream groups through the blueprint machinery Group entries used to be resolved to wire channels directly, so `.remappings()`, `.namespace()` and transport pins silently did not apply to them. Each entry now becomes a synthetic `In` stream: the blueprint atom declares it, the coordinator wires it like a declared port, and the module (python or native) reads the wired transport instead of computing a channel by convention. --- dimos/core/coordination/blueprints.py | 11 +++++ dimos/core/module.py | 49 ++++++++++++++----- dimos/core/native_module.py | 10 +++- dimos/core/test_async_module_stream_groups.py | 39 +++++++++++++++ dimos/core/test_native_module.py | 16 ++++++ native/rust/README.md | 4 +- 6 files changed, 115 insertions(+), 14 deletions(-) diff --git a/dimos/core/coordination/blueprints.py b/dimos/core/coordination/blueprints.py index b02453e7f5..a3a81d098c 100644 --- a/dimos/core/coordination/blueprints.py +++ b/dimos/core/coordination/blueprints.py @@ -140,6 +140,17 @@ def create(cls, module: type[ModuleBase], kwargs: dict[str, Any]) -> Self: elif is_module_type(inner): module_refs.append(ModuleRef(name=name, spec=inner, optional=True)) + # Stream-group entries are synthetic In streams: part of stream wiring + # (autoconnect, remapping, namespacing, transport pins) without being + # ports of their own. + for group in (kwargs.get("stream_groups") or {}).values(): + names = group["names"] if isinstance(group, dict) else group.names + msg_type = group.get("msg_type") if isinstance(group, dict) else group.msg_type + for name in names: + streams.append( + StreamRef(name=name, type=msg_type or Any, direction="in") # type: ignore[arg-type] + ) + instance_name = kwargs.get("instance_name") if instance_name is not None and not isinstance(instance_name, str): raise TypeError("instance_name must be a string or None") diff --git a/dimos/core/module.py b/dimos/core/module.py index acce3a8a69..08ff70af5e 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -106,14 +106,12 @@ def get_loop() -> tuple[asyncio.AbstractEventLoop, threading.Thread | None]: class StreamGroup(BaseModel): """Several same-typed streams that one port fans into a single handler. - `names` are stream names as they read after remapping and namespacing — - `left_cam`, `robot1/lidar` — never a backend topic string. Each is resolved - to a wire channel the same way a declared stream's name is, so the group - says nothing about which transport is in use. - - No `In` stream is declared for these, so they take no part in blueprint - autoconnection. A python module subscribes them itself; a native module - hands them to its subprocess, which does the subscribing. + `names` are stream names — `left_cam`, `robot1/lidar` — never a backend + topic string. Each entry becomes a synthetic `In` stream, so the blueprint + machinery treats it like a declared port: autoconnect matches it against + producers, `.remappings()` and `.namespace()` rewrite it, and transport + pins apply. A python module subscribes the wired streams itself; a native + module hands their channels to its subprocess, which does the subscribing. """ names: list[str] @@ -176,9 +174,11 @@ class ModuleBase(Configurable, CompositeResource): _main_gen: AsyncGenerator[None, None] | None = None _tools: dict[str, Any] _tools_lock: threading.Lock + _group_streams: dict[str, In[Any]] def __init__(self, config_args: dict[str, Any]) -> None: super().__init__(**config_args) + self._group_streams = self._make_group_streams() self._module_closed_lock = threading.Lock() self._tools = {} self._tools_lock = threading.Lock() @@ -697,6 +697,30 @@ def _auto_bind_handlers(self) -> None: # backpressure. self.process_observable(in_stream.pure_observable(), handler) + def _make_group_streams(self) -> "dict[str, In[Any]]": + """One synthetic `In` per stream-group entry, keyed by stream name. + + These are wired like declared ports (`set_transport` finds them, the + blueprint machinery remaps/namespaces/pins them), but they are not + attributes, so `inputs` and the native launch line never see them as + ports of their own. + """ + streams: dict[str, In[Any]] = {} + declared = set(self.inputs) | set(self.outputs) | set(self.ios) + for port, group in self.config.stream_groups.items(): + for name in group.names: + if name in declared: + raise ValueError( + f"stream group {port!r} entry {name!r} collides with a " + f"declared stream of {type(self).__name__}" + ) + if name in streams: + raise ValueError( + f"stream group {port!r} entry {name!r} appears in more than one group" + ) + streams[name] = In(group.msg_type or Any, name, self) + return streams + def _bind_stream_groups(self) -> None: """For each `stream_groups` port with an `async def handle_`, subscribe every stream in the group into that one handler. @@ -718,8 +742,11 @@ def _bind_stream_groups(self) -> None: on_msg, dispatcher_disp = self._make_keyed_dispatch(handler) self.register_disposable(dispatcher_disp) for index, name in enumerate(group.names): - transport = make_transport(name, group.msg_type) - self.register_disposable(Disposable(transport.subscribe(partial(on_msg, index)))) + stream = self._group_streams[name] + if getattr(stream, "_transport", None) is None: + # Not wired by a coordinator (standalone use): default transport. + stream.transport = make_transport(name, group.msg_type) + self.register_disposable(Disposable(stream.subscribe(partial(on_msg, index)))) def _make_async_dispatch( self, async_handler: Callable[[Any], Any] @@ -884,7 +911,7 @@ def __str__(self) -> str: @rpc def set_transport(self, stream_name: str, transport: Transport) -> bool: # type: ignore[type-arg] - stream = getattr(self, stream_name, None) + stream = self._group_streams.get(stream_name) or getattr(self, stream_name, None) if not stream: raise ValueError(f"{stream_name} not found in {self.__class__.__name__}") diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index ec728a7d27..38df3eab26 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -535,7 +535,15 @@ def _collect_topics(self) -> dict[str, str | list[str]]: f"[{self._module_label}] stream group {port!r} collides with the " "port of the same name declared as a stream" ) - topics[port] = [channel_for(n, group.msg_type) for n in group.names] + channels = [] + for name in group.names: + transport = getattr(self._group_streams[name], "_transport", None) + channel = getattr(transport, "channel", None) + # Unwired (standalone use): the channel the default transport lands on. + channels.append( + channel if channel is not None else channel_for(name, group.msg_type) + ) + topics[port] = channels return topics def _collect_output_qos(self) -> dict[str, dict[str, str]]: diff --git a/dimos/core/test_async_module_stream_groups.py b/dimos/core/test_async_module_stream_groups.py index a0120e5815..39dcf57d68 100644 --- a/dimos/core/test_async_module_stream_groups.py +++ b/dimos/core/test_async_module_stream_groups.py @@ -59,3 +59,42 @@ def test_stream_group_tags_each_message_with_its_stream(start_fan_in_module, gro s1.publish(7) assert queue.get(timeout=1.0) == 107 + + +@pytest.fixture +def start_remapped_fan_in_module(each_transport): + blueprint = FanInModule.blueprint( + stream_groups={"sensors": StreamGroup(names=["s0", "s1"])} + ).remappings([(FanInModule, "s0", "alt0")]) + coordinator = ModuleCoordinator.build(blueprint) + yield + coordinator.stop() + + +def test_stream_group_entries_follow_remappings(start_remapped_fan_in_module, each_transport): + """A remapped entry keeps its index but listens on the remapped stream.""" + alt0, tagged = make_transport("alt0"), make_transport("tagged") + for transport in (alt0, tagged): + transport.start() + try: + queue: Queue[int] = Queue() + tagged.subscribe(queue.put) + alt0.publish(7) + assert queue.get(timeout=1.0) == 7 + finally: + for transport in (alt0, tagged): + transport.stop() + + +def test_namespace_prefixes_stream_group_entries(): + blueprint = FanInModule.blueprint( + stream_groups={"sensors": StreamGroup(names=["s0", "s1"])} + ).namespace("bot") + atom = blueprint.blueprints[0] + assert blueprint.remapping_map[atom.name, "s0"] == "bot/s0" + assert blueprint.remapping_map[atom.name, "s1"] == "bot/s1" + + +def test_a_group_entry_cannot_collide_with_a_declared_stream(): + with pytest.raises(ValueError, match="collides"): + FanInModule(stream_groups={"sensors": StreamGroup(names=["tagged"])}) diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index 551388b5f2..6adbc2ecda 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -229,6 +229,22 @@ def test_a_stream_group_reaches_the_native_process_as_a_list(monkeypatch) -> Non module.stop() +def test_a_wired_stream_group_entry_uses_its_transport() -> None: + """Remapping/pins arrive as set_transport on the entry, and the launch line follows.""" + module = StubNativeModule( + executable=_ECHO, + stream_groups={"cams": StreamGroup(names=["cam0/imu"], msg_type=Imu)}, + ) + transport = LCMTransport("/remapped/imu", Imu) + try: + module.set_transport("cam0/imu", transport) + assert module._collect_topics()["cams"] == ["/remapped/imu#sensor_msgs.Imu"] + finally: + module.stop() + with contextlib.suppress(Exception): + transport.stop() + + def test_a_stream_group_rejects_a_backend_topic_string() -> None: """Names are stream names, so a leading slash is a transport leaking in.""" with pytest.raises(ValidationError, match="not topics"): diff --git a/native/rust/README.md b/native/rust/README.md index 1d2240493d..2adc198e07 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -135,9 +135,9 @@ MultiCam.blueprint( ) ``` -`names` are stream names as they read after remapping and namespacing, not backend topics — a leading `/` is rejected. Python resolves each to a wire channel the same way it resolves a declared stream's name, so the same blueprint runs unchanged over LCM or zenoh. +`names` are stream names, not backend topics — a leading `/` is rejected. Each entry becomes a synthetic `In` stream on the python side, so the blueprint machinery treats it like a declared port: autoconnect matches it against producers' `Out` streams, `.remappings()` and `.namespace()` rewrite it, and blueprint transport pins apply. The same blueprint runs unchanged over LCM or zenoh. -Those streams are not declared as `In` ports, so they take no part in blueprint autoconnection: the names are resolved to wire channels and the native process subscribes them itself. On the launch line the port's value is an array rather than a string. A group configured with no names still claims its port but never yields. +The group itself is not a port with a stream of its own: python hands the wired entries' channels to the native process, which subscribes them directly. On the launch line the port's value is an array rather than a string. A group configured with no names still claims its port but never yields. `stream_groups` lives on `ModuleConfig`, so a plain Python `Module` takes the same field. There the module subscribes the group itself and dispatches to `async def handle_(self, index, msg)`, matching the Rust handler signature: From 4ec5eb920cadff2d536feb53948045386e3840ae Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 29 Aug 2026 15:32:28 -0700 Subject: [PATCH 04/16] Rename stream groups to topic funnels `StreamGroup`/`stream_groups` becomes `TopicFunnel`/`topic_funnels` in python, and rust's `InputGroup`/`#[input_group]` becomes `TopicFunnel`/`#[topic_funnel]` to match. Naming only, no behavior change. --- dimos/core/coordination/blueprints.py | 4 +-- dimos/core/module.py | 34 +++++++++---------- dimos/core/native_module.py | 8 ++--- ....py => test_async_module_topic_funnels.py} | 16 ++++----- dimos/core/test_native_module.py | 26 +++++++------- native/rust/README.md | 16 ++++----- native/rust/dimos-module-macros/src/lib.rs | 25 ++++++++------ native/rust/dimos-module/src/lib.rs | 2 +- native/rust/dimos-module/src/module.rs | 22 ++++++------ 9 files changed, 78 insertions(+), 75 deletions(-) rename dimos/core/{test_async_module_stream_groups.py => test_async_module_topic_funnels.py} (84%) diff --git a/dimos/core/coordination/blueprints.py b/dimos/core/coordination/blueprints.py index a3a81d098c..f03e8fc096 100644 --- a/dimos/core/coordination/blueprints.py +++ b/dimos/core/coordination/blueprints.py @@ -140,10 +140,10 @@ def create(cls, module: type[ModuleBase], kwargs: dict[str, Any]) -> Self: elif is_module_type(inner): module_refs.append(ModuleRef(name=name, spec=inner, optional=True)) - # Stream-group entries are synthetic In streams: part of stream wiring + # Topic-funnel entries are synthetic In streams: part of stream wiring # (autoconnect, remapping, namespacing, transport pins) without being # ports of their own. - for group in (kwargs.get("stream_groups") or {}).values(): + for group in (kwargs.get("topic_funnels") or {}).values(): names = group["names"] if isinstance(group, dict) else group.names msg_type = group.get("msg_type") if isinstance(group, dict) else group.msg_type for name in names: diff --git a/dimos/core/module.py b/dimos/core/module.py index 08ff70af5e..740a94d24a 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -103,7 +103,7 @@ def get_loop() -> tuple[asyncio.AbstractEventLoop, threading.Thread | None]: Deployment = Literal["python", "docker"] -class StreamGroup(BaseModel): +class TopicFunnel(BaseModel): """Several same-typed streams that one port fans into a single handler. `names` are stream names — `left_cam`, `robot1/lidar` — never a backend @@ -123,7 +123,7 @@ def _reject_topic_strings(cls, names: list[str]) -> list[str]: leading_slash = [n for n in names if n.startswith("/")] if leading_slash: raise ValueError( - f"stream group names must be stream names, not topics: {leading_slash} " + f"topic funnel names must be stream names, not topics: {leading_slash} " "start with '/' (use `left_cam`, not `/left_cam`)" ) return names @@ -140,7 +140,7 @@ class ModuleConfig(BaseConfig): # from the class name to this name. instance_name: str | None = None # Port name -> the group of streams that port's single handler receives. - stream_groups: dict[str, StreamGroup] = Field(default_factory=dict) + topic_funnels: dict[str, TopicFunnel] = Field(default_factory=dict) g: GlobalConfig = global_config @@ -174,11 +174,11 @@ class ModuleBase(Configurable, CompositeResource): _main_gen: AsyncGenerator[None, None] | None = None _tools: dict[str, Any] _tools_lock: threading.Lock - _group_streams: dict[str, In[Any]] + _funnel_streams: dict[str, In[Any]] def __init__(self, config_args: dict[str, Any]) -> None: super().__init__(**config_args) - self._group_streams = self._make_group_streams() + self._funnel_streams = self._make_funnel_streams() self._module_closed_lock = threading.Lock() self._tools = {} self._tools_lock = threading.Lock() @@ -220,7 +220,7 @@ def build(self) -> None: def start(self) -> None: self._start_main() self._auto_bind_handlers() - self._bind_stream_groups() + self._bind_topic_funnels() @rpc def stop(self) -> None: @@ -697,8 +697,8 @@ def _auto_bind_handlers(self) -> None: # backpressure. self.process_observable(in_stream.pure_observable(), handler) - def _make_group_streams(self) -> "dict[str, In[Any]]": - """One synthetic `In` per stream-group entry, keyed by stream name. + def _make_funnel_streams(self) -> "dict[str, In[Any]]": + """One synthetic `In` per topic-funnel entry, keyed by stream name. These are wired like declared ports (`set_transport` finds them, the blueprint machinery remaps/namespaces/pins them), but they are not @@ -707,28 +707,28 @@ def _make_group_streams(self) -> "dict[str, In[Any]]": """ streams: dict[str, In[Any]] = {} declared = set(self.inputs) | set(self.outputs) | set(self.ios) - for port, group in self.config.stream_groups.items(): + for port, group in self.config.topic_funnels.items(): for name in group.names: if name in declared: raise ValueError( - f"stream group {port!r} entry {name!r} collides with a " + f"topic funnel {port!r} entry {name!r} collides with a " f"declared stream of {type(self).__name__}" ) if name in streams: raise ValueError( - f"stream group {port!r} entry {name!r} appears in more than one group" + f"topic funnel {port!r} entry {name!r} appears in more than one group" ) streams[name] = In(group.msg_type or Any, name, self) return streams - def _bind_stream_groups(self) -> None: - """For each `stream_groups` port with an `async def handle_`, subscribe + def _bind_topic_funnels(self) -> None: + """For each `topic_funnels` port with an `async def handle_`, subscribe every stream in the group into that one handler. A native module defines no such method — its subprocess subscribes instead — so this is a no-op there. """ - for port, group in self.config.stream_groups.items(): + for port, group in self.config.topic_funnels.items(): handler = getattr(self, f"handle_{port}", None) if handler is None: continue @@ -737,12 +737,12 @@ def _bind_stream_groups(self) -> None: if not inspect.iscoroutinefunction(handler): raise TypeError( f"{type(self).__name__}.handle_{port} must be `async def` " - "(stream groups have no sync path)" + "(topic funnels have no sync path)" ) on_msg, dispatcher_disp = self._make_keyed_dispatch(handler) self.register_disposable(dispatcher_disp) for index, name in enumerate(group.names): - stream = self._group_streams[name] + stream = self._funnel_streams[name] if getattr(stream, "_transport", None) is None: # Not wired by a coordinator (standalone use): default transport. stream.transport = make_transport(name, group.msg_type) @@ -911,7 +911,7 @@ def __str__(self) -> str: @rpc def set_transport(self, stream_name: str, transport: Transport) -> bool: # type: ignore[type-arg] - stream = self._group_streams.get(stream_name) or getattr(self, stream_name, None) + stream = self._funnel_streams.get(stream_name) or getattr(self, stream_name, None) if not stream: raise ValueError(f"{stream_name} not found in {self.__class__.__name__}") diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index 38df3eab26..c133dafe1d 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -201,7 +201,7 @@ class NativeModule(Module): ``DIMOS_TRANSPORT`` env var. With ``stdin_config``, the topics, config, publisher QoS and session settings also arrive as one JSON line on stdin. - A port named in ``stream_groups`` gets a list of channels instead of one, so + A port named in ``topic_funnels`` gets a list of channels instead of one, so several same-typed streams reach a single native handler. The native process should parse whichever it uses and pub/sub on the given @@ -529,15 +529,15 @@ def _collect_topics(self) -> dict[str, str | list[str]]: channel = getattr(transport, "channel", None) if channel is not None: topics[name] = channel - for port, group in self.config.stream_groups.items(): + for port, group in self.config.topic_funnels.items(): if port in topics: raise ValueError( - f"[{self._module_label}] stream group {port!r} collides with the " + f"[{self._module_label}] topic funnel {port!r} collides with the " "port of the same name declared as a stream" ) channels = [] for name in group.names: - transport = getattr(self._group_streams[name], "_transport", None) + transport = getattr(self._funnel_streams[name], "_transport", None) channel = getattr(transport, "channel", None) # Unwired (standalone use): the channel the default transport lands on. channels.append( diff --git a/dimos/core/test_async_module_stream_groups.py b/dimos/core/test_async_module_topic_funnels.py similarity index 84% rename from dimos/core/test_async_module_stream_groups.py rename to dimos/core/test_async_module_topic_funnels.py index 39dcf57d68..f867159f59 100644 --- a/dimos/core/test_async_module_stream_groups.py +++ b/dimos/core/test_async_module_topic_funnels.py @@ -17,7 +17,7 @@ import pytest from dimos.core.coordination.module_coordinator import ModuleCoordinator -from dimos.core.module import Module, StreamGroup +from dimos.core.module import Module, TopicFunnel from dimos.core.stream import Out from dimos.core.transport_factory import make_transport @@ -33,7 +33,7 @@ async def handle_sensors(self, index: int, value: int) -> None: @pytest.fixture def start_fan_in_module(each_transport): - blueprint = FanInModule.blueprint(stream_groups={"sensors": StreamGroup(names=["s0", "s1"])}) + blueprint = FanInModule.blueprint(topic_funnels={"sensors": TopicFunnel(names=["s0", "s1"])}) coordinator = ModuleCoordinator.build(blueprint) yield coordinator.stop() @@ -49,7 +49,7 @@ def group_transports(each_transport): transport.stop() -def test_stream_group_tags_each_message_with_its_stream(start_fan_in_module, group_transports): +def test_topic_funnel_tags_each_message_with_its_stream(start_fan_in_module, group_transports): s0, s1, tagged = group_transports queue: Queue[int] = Queue() tagged.subscribe(queue.put) @@ -64,14 +64,14 @@ def test_stream_group_tags_each_message_with_its_stream(start_fan_in_module, gro @pytest.fixture def start_remapped_fan_in_module(each_transport): blueprint = FanInModule.blueprint( - stream_groups={"sensors": StreamGroup(names=["s0", "s1"])} + topic_funnels={"sensors": TopicFunnel(names=["s0", "s1"])} ).remappings([(FanInModule, "s0", "alt0")]) coordinator = ModuleCoordinator.build(blueprint) yield coordinator.stop() -def test_stream_group_entries_follow_remappings(start_remapped_fan_in_module, each_transport): +def test_topic_funnel_entries_follow_remappings(start_remapped_fan_in_module, each_transport): """A remapped entry keeps its index but listens on the remapped stream.""" alt0, tagged = make_transport("alt0"), make_transport("tagged") for transport in (alt0, tagged): @@ -86,9 +86,9 @@ def test_stream_group_entries_follow_remappings(start_remapped_fan_in_module, ea transport.stop() -def test_namespace_prefixes_stream_group_entries(): +def test_namespace_prefixes_topic_funnel_entries(): blueprint = FanInModule.blueprint( - stream_groups={"sensors": StreamGroup(names=["s0", "s1"])} + topic_funnels={"sensors": TopicFunnel(names=["s0", "s1"])} ).namespace("bot") atom = blueprint.blueprints[0] assert blueprint.remapping_map[atom.name, "s0"] == "bot/s0" @@ -97,4 +97,4 @@ def test_namespace_prefixes_stream_group_entries(): def test_a_group_entry_cannot_collide_with_a_declared_stream(): with pytest.raises(ValueError, match="collides"): - FanInModule(stream_groups={"sensors": StreamGroup(names=["tagged"])}) + FanInModule(topic_funnels={"sensors": TopicFunnel(names=["tagged"])}) diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index 6adbc2ecda..e411efdcf0 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -35,7 +35,7 @@ from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.core import rpc from dimos.core.global_config import GlobalConfig, TransportBackend -from dimos.core.module import Module, StreamGroup +from dimos.core.module import Module, TopicFunnel from dimos.core.native_module import LogFormat, NativeModule, NativeModuleConfig from dimos.core.stream import IO, In, Out from dimos.core.transport import LCMTransport, ZenohTransport @@ -211,12 +211,12 @@ def test_tf_topic_comes_from_the_declared_port_only() -> None: transport.stop() -def test_a_stream_group_reaches_the_native_process_as_a_list(monkeypatch) -> None: +def test_a_topic_funnel_reaches_the_native_process_as_a_list(monkeypatch) -> None: """A group's channels are resolved without the module declaring a stream each.""" monkeypatch.setattr(native_module_mod.global_config, "transport", "lcm") module = StubNativeModule( executable=_ECHO, - stream_groups={"cams": StreamGroup(names=["cam0/imu", "cam1/imu"], msg_type=Imu)}, + topic_funnels={"cams": TopicFunnel(names=["cam0/imu", "cam1/imu"], msg_type=Imu)}, ) try: topics = module._collect_topics() @@ -229,11 +229,11 @@ def test_a_stream_group_reaches_the_native_process_as_a_list(monkeypatch) -> Non module.stop() -def test_a_wired_stream_group_entry_uses_its_transport() -> None: +def test_a_wired_topic_funnel_entry_uses_its_transport() -> None: """Remapping/pins arrive as set_transport on the entry, and the launch line follows.""" module = StubNativeModule( executable=_ECHO, - stream_groups={"cams": StreamGroup(names=["cam0/imu"], msg_type=Imu)}, + topic_funnels={"cams": TopicFunnel(names=["cam0/imu"], msg_type=Imu)}, ) transport = LCMTransport("/remapped/imu", Imu) try: @@ -245,16 +245,16 @@ def test_a_wired_stream_group_entry_uses_its_transport() -> None: transport.stop() -def test_a_stream_group_rejects_a_backend_topic_string() -> None: +def test_a_topic_funnel_rejects_a_backend_topic_string() -> None: """Names are stream names, so a leading slash is a transport leaking in.""" with pytest.raises(ValidationError, match="not topics"): - StreamGroup(names=["/cam0/imu"], msg_type=Imu) + TopicFunnel(names=["/cam0/imu"], msg_type=Imu) -def test_a_stream_group_cannot_shadow_a_declared_port() -> None: +def test_a_topic_funnel_cannot_shadow_a_declared_port() -> None: module = StubNativeModule( executable=_ECHO, - stream_groups={"cmd_vel": StreamGroup(names=["other"], msg_type=Twist)}, + topic_funnels={"cmd_vel": TopicFunnel(names=["other"], msg_type=Twist)}, ) transport = LCMTransport("/cmd_vel", Twist) try: @@ -267,15 +267,15 @@ def test_a_stream_group_cannot_shadow_a_declared_port() -> None: transport.stop() -def test_a_stream_group_is_not_a_native_config_field() -> None: +def test_a_topic_funnel_is_not_a_native_config_field() -> None: """The group is wiring, so it belongs in `topics`, not the config struct.""" module = StubNativeModule( executable=_ECHO, - stream_groups={"cams": StreamGroup(names=["cam0/imu"], msg_type=Imu)}, + topic_funnels={"cams": TopicFunnel(names=["cam0/imu"], msg_type=Imu)}, ) try: - assert "stream_groups" not in module.config.to_config_dict() - assert "--stream_groups" not in module._argv({}) + assert "topic_funnels" not in module.config.to_config_dict() + assert "--topic_funnels" not in module._argv({}) finally: module.stop() diff --git a/native/rust/README.md b/native/rust/README.md index 2adc198e07..1663a6d5d7 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -62,7 +62,7 @@ Every transport is compiled into the binary. `run_with_transport` opens the one - `#[derive(Module)]`: on the struct. Required. - `#[module(setup = fn, teardown = fn)]`: on the struct. Both optional. Names methods on `Self`. `setup` runs once before the input dispatch loop starts (use it to spawn background tasks or initialize resources); `teardown` runs once after the loop exits (use it for cleanup). - `#[input(decode = fn, handler = fn)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. -- `#[input_group(decode = fn, handler = fn)]`: on a field of type `InputGroup`, one port fed by several topics of the same message type (see [Stream groups](#stream-groups)). `decode` is required; `handler` defaults to `handle_` and takes `(index, msg)`. +- `#[topic_funnel(decode = fn, handler = fn)]`: on a field of type `TopicFunnel`, one port fed by several topics of the same message type (see [Topic funnels](#topic-funnels)). `decode` is required; `handler` defaults to `handle_` and takes `(index, msg)`. - `#[output(encode = fn)]`: on a field of type `Output`. `encode` is required. - `#[io(decode = fn, encode = fn, handler = fn)]`: on a field of type `Io`, a port that publishes to and subscribes on one topic. `decode` and `encode` are required; `handler` defaults to `handle_`. The transports deliver a message back to its own sender, so the handler also sees what the module publishes. Use `#[output]` instead when the module only publishes. - `#[config]`: on one field. The type must be defined with `#[native_config]` (see [Config](#config)). At most one per struct. If absent, `Config` defaults to `dimos_module::NoConfig`. @@ -109,15 +109,15 @@ At runtime `run()` enforces the mapping on the Python payload: deserialization r Field name = port name. Ports map to topics via the stdin JSON; unmapped ports fall back to `/{port}`. -## Stream groups +## Topic funnels -A rig with N identical sensors would otherwise need N ports and N near-identical handlers. An `InputGroup` is one port wired to a list of topics that all carry `T`, delivered to one handler in arrival order. Each message is tagged with the index of the topic it arrived on, so the handler can tell the sources apart. +A rig with N identical sensors would otherwise need N ports and N near-identical handlers. An `TopicFunnel` is one port wired to a list of topics that all carry `T`, delivered to one handler in arrival order. Each message is tagged with the index of the topic it arrived on, so the handler can tell the sources apart. ```rust #[derive(Module)] struct MultiCam { - #[input_group(decode = Image::decode)] - cameras: InputGroup, + #[topic_funnel(decode = Image::decode)] + cameras: TopicFunnel, } impl MultiCam { @@ -127,11 +127,11 @@ impl MultiCam { } ``` -The Python wrapper supplies the sources with `stream_groups`, keyed by port name: +The Python wrapper supplies the sources with `topic_funnels`, keyed by port name: ```python MultiCam.blueprint( - stream_groups={"cameras": StreamGroup(names=["left_cam", "right_cam"], msg_type=Image)}, + topic_funnels={"cameras": TopicFunnel(names=["left_cam", "right_cam"], msg_type=Image)}, ) ``` @@ -139,7 +139,7 @@ MultiCam.blueprint( The group itself is not a port with a stream of its own: python hands the wired entries' channels to the native process, which subscribes them directly. On the launch line the port's value is an array rather than a string. A group configured with no names still claims its port but never yields. -`stream_groups` lives on `ModuleConfig`, so a plain Python `Module` takes the same field. There the module subscribes the group itself and dispatches to `async def handle_(self, index, msg)`, matching the Rust handler signature: +`topic_funnels` lives on `ModuleConfig`, so a plain Python `Module` takes the same field. There the module subscribes the group itself and dispatches to `async def handle_(self, index, msg)`, matching the Rust handler signature: ```python class MultiCam(Module): diff --git a/native/rust/dimos-module-macros/src/lib.rs b/native/rust/dimos-module-macros/src/lib.rs index da468f0afe..ae5bf4878c 100644 --- a/native/rust/dimos-module-macros/src/lib.rs +++ b/native/rust/dimos-module-macros/src/lib.rs @@ -17,7 +17,10 @@ use proc_macro2::TokenStream as TokenStream2; use quote::{format_ident, quote}; use syn::{parse_macro_input, Data, DeriveInput, Field, Fields, Ident, LitStr, Path, Type}; -#[proc_macro_derive(Module, attributes(input, input_group, output, io, config, tf, module))] +#[proc_macro_derive( + Module, + attributes(input, topic_funnel, output, io, config, tf, module) +)] pub fn derive_module(input: TokenStream) -> TokenStream { let input = parse_macro_input!(input as DeriveInput); match expand(input) { @@ -175,7 +178,7 @@ fn is_option(ty: &Type) -> bool { } const ONE_ATTR_ONLY: &str = "field has multiple module attributes; only one of #[input], \ - #[input_group], #[output], #[io], #[config], #[tf] is allowed"; + #[topic_funnel], #[output], #[io], #[config], #[tf] is allowed"; /// What a `#[tf]` port carries, for the Cargo.toml registry cross-check. const TF_PAYLOAD_TYPE: &str = "TFMessage"; @@ -185,7 +188,7 @@ enum FieldKind { decode: Path, handler: Ident, }, - InputGroup { + TopicFunnel { decode: Path, handler: Ident, }, @@ -327,8 +330,8 @@ fn expand(input: DeriveInput) -> syn::Result { FieldKind::Input { decode, .. } => { quote!(#name: builder.input(#name_str, #decode)) } - FieldKind::InputGroup { decode, .. } => { - quote!(#name: builder.input_group(#name_str, #decode)) + FieldKind::TopicFunnel { decode, .. } => { + quote!(#name: builder.topic_funnel(#name_str, #decode)) } FieldKind::Output { encode } => { quote!(#name: builder.output(#name_str, #encode)) @@ -353,7 +356,7 @@ fn expand(input: DeriveInput) -> syn::Result { self.#handler(msg).await } )), - FieldKind::InputGroup { handler, .. } => Some(quote!( + FieldKind::TopicFunnel { handler, .. } => Some(quote!( ::core::option::Option::Some((index, msg)) = self.#name.recv() => { self.#handler(index, msg).await } @@ -444,7 +447,7 @@ fn port_decls(classified: &[ClassifiedField], want_input: bool) -> Vec classified .iter() .filter(|f| match f.kind { - FieldKind::Input { .. } | FieldKind::InputGroup { .. } => want_input, + FieldKind::Input { .. } | FieldKind::TopicFunnel { .. } => want_input, FieldKind::Output { .. } => !want_input, // A `#[tf]` field is a port like any other as far as the registry // goes: bake has to put the topic in the host's map or the module @@ -587,7 +590,7 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { .ok_or_else(|| syn::Error::new_spanned(attr, "#[input] requires `decode = ...`"))?; let handler = handler.unwrap_or_else(|| format_ident!("handle_{}", name)); found = Some(FieldKind::Input { decode, handler }); - } else if path.is_ident("input_group") { + } else if path.is_ident("topic_funnel") { if found.is_some() { return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); } @@ -602,17 +605,17 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { msg = Some(meta.value()?.parse::()?.value()); } else { return Err(meta.error( - "unrecognized #[input_group] argument; expected `decode = ...`, \ + "unrecognized #[topic_funnel] argument; expected `decode = ...`, \ `handler = ...` or `msg = ...`", )); } Ok(()) })?; let decode = decode.ok_or_else(|| { - syn::Error::new_spanned(attr, "#[input_group] requires `decode = ...`") + syn::Error::new_spanned(attr, "#[topic_funnel] requires `decode = ...`") })?; let handler = handler.unwrap_or_else(|| format_ident!("handle_{}", name)); - found = Some(FieldKind::InputGroup { decode, handler }); + found = Some(FieldKind::TopicFunnel { decode, handler }); } else if path.is_ident("output") { if found.is_some() { return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); diff --git a/native/rust/dimos-module/src/lib.rs b/native/rust/dimos-module/src/lib.rs index 5b108b95ae..73d9fce62b 100644 --- a/native/rust/dimos-module/src/lib.rs +++ b/native/rust/dimos-module/src/lib.rs @@ -30,7 +30,7 @@ pub use dimos_module_macros::{native_config, Module}; pub use host::{host_main, HostSpec, ModuleEntry}; pub use lcm::LcmTransport; pub use module::{ - run, Builder, Input, InputGroup, Io, Module, ModuleConfig, NativeConfig, NoConfig, Output, + run, Builder, Input, Io, Module, ModuleConfig, NativeConfig, NoConfig, Output, TopicFunnel, }; pub use tf::{Lookup, Tf, Transform}; pub use transport::{SharedTransport, Transport}; diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index c526af6a20..9af85ff09f 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -129,12 +129,12 @@ impl Input { /// The launch line gives the port an array of topics instead of one, and every /// message is tagged with the index of the topic it arrived on so a handler can /// tell a rig's cameras apart. A group configured with no topics never yields. -pub struct InputGroup { +pub struct TopicFunnel { pub topics: Vec, receiver: mpsc::Receiver<(usize, T)>, } -impl InputGroup { +impl TopicFunnel { pub async fn recv(&mut self) -> Option<(usize, T)> { self.receiver.recv().await } @@ -480,18 +480,18 @@ impl Builder { /// A port wired to every topic the launch line lists under `port`, all of /// one message type, delivered to one handler in arrival order. - pub fn input_group( + pub fn topic_funnel( &mut self, port: &str, decode: fn(&[u8]) -> io::Result, - ) -> InputGroup { + ) -> TopicFunnel { let topics = self.group_for(port); let (sender, receiver) = mpsc::channel(INPUT_CHANNEL_CAPACITY); for (index, topic) in topics.iter().enumerate() { let tag = Box::new(move |bytes: &[u8]| decode(bytes).map(|msg| (index, msg))); self.push_route(topic, tag, sender.clone()); } - InputGroup { topics, receiver } + TopicFunnel { topics, receiver } } pub fn output(&mut self, port: &str, encode: fn(&T) -> Vec) -> Output { @@ -1073,7 +1073,7 @@ mod tests { #[test] fn a_group_subscribes_every_topic_it_was_given() { let mut builder = Builder::new(grouped_topics("cams", &["/cam0", "/cam1"])); - let group = builder.input_group("cams", |b| Ok(b.to_vec())); + let group = builder.topic_funnel("cams", |b| Ok(b.to_vec())); assert_eq!(group.topics, ["/cam0", "/cam1"]); assert_eq!(builder.routes.get("/cam0").map(Vec::len), Some(1)); assert_eq!(builder.routes.get("/cam1").map(Vec::len), Some(1)); @@ -1085,7 +1085,7 @@ mod tests { #[tokio::test] async fn a_group_tags_each_message_with_its_topic_index() { let mut builder = Builder::new(grouped_topics("cams", &["/cam0", "/cam1"])); - let mut group = builder.input_group("cams", |b| Ok(b.to_vec())); + let mut group = builder.topic_funnel("cams", |b| Ok(b.to_vec())); builder.routes["/cam1"][0].try_dispatch(b"second"); builder.routes["/cam0"][0].try_dispatch(b"first"); @@ -1104,7 +1104,7 @@ mod tests { #[test] fn an_empty_group_still_claims_its_port() { let mut builder = Builder::new(grouped_topics("cams", &[])); - let group = builder.input_group("cams", |b| Ok(b.to_vec())); + let group = builder.topic_funnel("cams", |b| Ok(b.to_vec())); assert!(group.is_empty()); assert!(builder.routes.is_empty()); builder.enforce_topics_match_ports().expect("group claimed"); @@ -1465,8 +1465,8 @@ mod tests { #[derive(crate::Module)] struct Rig { - #[input_group(decode = decode)] - cams: crate::InputGroup, + #[topic_funnel(decode = decode)] + cams: crate::TopicFunnel, #[output(encode = encode)] seen: crate::Output, } @@ -1480,7 +1480,7 @@ mod tests { } #[tokio::test] - async fn input_group_field_hands_its_handler_the_topic_index() { + async fn topic_funnel_field_hands_its_handler_the_topic_index() { let mut builder = Builder::new(Topics { single: HashMap::from([("seen".to_string(), "/seen".to_string())]), grouped: HashMap::from([( From 10404f9fd28fbef490e1ca64d950d4911d33eca9 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 29 Aug 2026 16:02:42 -0700 Subject: [PATCH 05/16] Make topic-funnel metadata opt-in on plain inputs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rust: fold TopicFunnel into Input — every port takes one topic or an array from the launch line. Handlers keep their (msg) signature; a handler that needs the source opts in with #[input(meta)] and receives (msg, Metadata { index, topic }). Python: opt in by signature instead — a two-parameter handle_ or handle_ receives (msg, Metadata(index, name)), where name is the declared stream name pre-remapping; a one-parameter handler just gets the message. --- dimos/core/module.py | 58 +++++++- dimos/core/test_async_module_topic_funnels.py | 61 ++++++-- native/rust/README.md | 19 +-- native/rust/dimos-module-macros/src/lib.rs | 68 +++------ native/rust/dimos-module/src/lib.rs | 2 +- native/rust/dimos-module/src/module.rs | 133 +++++++++--------- 6 files changed, 208 insertions(+), 133 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index 740a94d24a..5c716a6f46 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -129,6 +129,33 @@ def _reject_topic_strings(cls, names: list[str]) -> list[str]: return names +@dataclass(frozen=True) +class Metadata: + """Which stream a message arrived on, for handlers that take `(msg, meta)`. + + `index` is the position in the funnel's `names` (0 for a plain input); + `name` is the stream name as the module declared it, pre-remapping. + """ + + index: int + name: str + + +def _handler_wants_metadata(handler: Callable[..., Any], label: str) -> bool: + """True if the bound handler takes `(msg, meta)` rather than just `(msg)`.""" + positional_kinds = (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) + count = sum( + 1 + for parameter in inspect.signature(handler).parameters.values() + if parameter.kind in positional_kinds + ) + if count == 1: + return False + if count == 2: + return True + raise TypeError(f"{label} must take (msg) or (msg, meta), not {count} positional parameters") + + class ModuleConfig(BaseConfig): rpc_transport: type[RPCSpec] = Field(default_factory=rpc_backend) default_rpc_timeout: float = DEFAULT_RPC_TIMEOUT @@ -687,6 +714,17 @@ def _auto_bind_handlers(self) -> None: f"{type(self).__name__}.handle_{input_name} must be `async def` " "(use a manual self..subscribe(...) for sync handlers)" ) + if _handler_wants_metadata(handler, f"{type(self).__name__}.handle_{input_name}"): + metadata = Metadata(index=0, name=input_name) + + async def with_meta( + msg: Any, + _handler: Callable[[Any, Metadata], Any] = handler, + _metadata: Metadata = metadata, + ) -> None: + await _handler(msg, _metadata) + + handler = with_meta bindings.append((in_stream, handler)) for in_stream, handler in bindings: @@ -739,7 +777,25 @@ def _bind_topic_funnels(self) -> None: f"{type(self).__name__}.handle_{port} must be `async def` " "(topic funnels have no sync path)" ) - on_msg, dispatcher_disp = self._make_keyed_dispatch(handler) + if _handler_wants_metadata(handler, f"{type(self).__name__}.handle_{port}"): + + async def invoke( + index: int, + msg: Any, + _handler: Callable[[Any, Metadata], Any] = handler, + _names: tuple[str, ...] = tuple(group.names), + ) -> None: + await _handler(msg, Metadata(index=index, name=_names[index])) + else: + + async def invoke( # type: ignore[misc] + index: int, + msg: Any, + _handler: Callable[[Any], Any] = handler, + ) -> None: + await _handler(msg) + + on_msg, dispatcher_disp = self._make_keyed_dispatch(invoke) self.register_disposable(dispatcher_disp) for index, name in enumerate(group.names): stream = self._funnel_streams[name] diff --git a/dimos/core/test_async_module_topic_funnels.py b/dimos/core/test_async_module_topic_funnels.py index f867159f59..28fdf1ec05 100644 --- a/dimos/core/test_async_module_topic_funnels.py +++ b/dimos/core/test_async_module_topic_funnels.py @@ -17,7 +17,7 @@ import pytest from dimos.core.coordination.module_coordinator import ModuleCoordinator -from dimos.core.module import Module, TopicFunnel +from dimos.core.module import Metadata, Module, TopicFunnel from dimos.core.stream import Out from dimos.core.transport_factory import make_transport @@ -26,9 +26,11 @@ class FanInModule(Module): """One handler for two same-typed streams; echoes which one it came from.""" tagged: Out[int] + named: Out[str] - async def handle_sensors(self, index: int, value: int) -> None: - self.tagged.publish(index * 100 + value) + async def handle_sensors(self, value: int, meta: Metadata) -> None: + self.tagged.publish(meta.index * 100 + value) + self.named.publish(meta.name) @pytest.fixture @@ -41,7 +43,12 @@ def start_fan_in_module(each_transport): @pytest.fixture def group_transports(each_transport): - transports = [make_transport("s0"), make_transport("s1"), make_transport("tagged")] + transports = [ + make_transport("s0"), + make_transport("s1"), + make_transport("tagged"), + make_transport("named"), + ] for transport in transports: transport.start() yield transports @@ -50,15 +57,19 @@ def group_transports(each_transport): def test_topic_funnel_tags_each_message_with_its_stream(start_fan_in_module, group_transports): - s0, s1, tagged = group_transports + s0, s1, tagged, named = group_transports queue: Queue[int] = Queue() + names: Queue[str] = Queue() tagged.subscribe(queue.put) + named.subscribe(names.put) s0.publish(7) assert queue.get(timeout=1.0) == 7 + assert names.get(timeout=1.0) == "s0" s1.publish(7) assert queue.get(timeout=1.0) == 107 + assert names.get(timeout=1.0) == "s1" @pytest.fixture @@ -72,17 +83,21 @@ def start_remapped_fan_in_module(each_transport): def test_topic_funnel_entries_follow_remappings(start_remapped_fan_in_module, each_transport): - """A remapped entry keeps its index but listens on the remapped stream.""" - alt0, tagged = make_transport("alt0"), make_transport("tagged") - for transport in (alt0, tagged): + """A remapped entry keeps its index and declared name, but listens on the + remapped stream.""" + alt0, tagged, named = make_transport("alt0"), make_transport("tagged"), make_transport("named") + for transport in (alt0, tagged, named): transport.start() try: queue: Queue[int] = Queue() + names: Queue[str] = Queue() tagged.subscribe(queue.put) + named.subscribe(names.put) alt0.publish(7) assert queue.get(timeout=1.0) == 7 + assert names.get(timeout=1.0) == "s0" finally: - for transport in (alt0, tagged): + for transport in (alt0, tagged, named): transport.stop() @@ -98,3 +113,31 @@ def test_namespace_prefixes_topic_funnel_entries(): def test_a_group_entry_cannot_collide_with_a_declared_stream(): with pytest.raises(ValueError, match="collides"): FanInModule(topic_funnels={"sensors": TopicFunnel(names=["tagged"])}) + + +class PlainFanInModule(Module): + """A funnel handler that doesn't ask for metadata just gets the message.""" + + echoed: Out[int] + + async def handle_sensors(self, value: int) -> None: + self.echoed.publish(value) + + +def test_a_funnel_handler_without_meta_gets_just_the_message(each_transport): + blueprint = PlainFanInModule.blueprint( + topic_funnels={"sensors": TopicFunnel(names=["s0", "s1"])} + ) + coordinator = ModuleCoordinator.build(blueprint) + s1, echoed = make_transport("s1"), make_transport("echoed") + for transport in (s1, echoed): + transport.start() + try: + queue: Queue[int] = Queue() + echoed.subscribe(queue.put) + s1.publish(7) + assert queue.get(timeout=1.0) == 7 + finally: + for transport in (s1, echoed): + transport.stop() + coordinator.stop() diff --git a/native/rust/README.md b/native/rust/README.md index 1663a6d5d7..8a0ff55c65 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -61,8 +61,7 @@ Every transport is compiled into the binary. `run_with_transport` opens the one - `#[derive(Module)]`: on the struct. Required. - `#[module(setup = fn, teardown = fn)]`: on the struct. Both optional. Names methods on `Self`. `setup` runs once before the input dispatch loop starts (use it to spawn background tasks or initialize resources); `teardown` runs once after the loop exits (use it for cleanup). -- `#[input(decode = fn, handler = fn)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. -- `#[topic_funnel(decode = fn, handler = fn)]`: on a field of type `TopicFunnel`, one port fed by several topics of the same message type (see [Topic funnels](#topic-funnels)). `decode` is required; `handler` defaults to `handle_` and takes `(index, msg)`. +- `#[input(decode = fn, handler = fn, meta)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. The `meta` flag makes the handler take `(msg, meta: Metadata)`, where `meta` names the topic the message arrived on — useful when the port is a topic funnel (see [Topic funnels](#topic-funnels)). - `#[output(encode = fn)]`: on a field of type `Output`. `encode` is required. - `#[io(decode = fn, encode = fn, handler = fn)]`: on a field of type `Io`, a port that publishes to and subscribes on one topic. `decode` and `encode` are required; `handler` defaults to `handle_`. The transports deliver a message back to its own sender, so the handler also sees what the module publishes. Use `#[output]` instead when the module only publishes. - `#[config]`: on one field. The type must be defined with `#[native_config]` (see [Config](#config)). At most one per struct. If absent, `Config` defaults to `dimos_module::NoConfig`. @@ -111,18 +110,18 @@ Field name = port name. Ports map to topics via the stdin JSON; unmapped ports f ## Topic funnels -A rig with N identical sensors would otherwise need N ports and N near-identical handlers. An `TopicFunnel` is one port wired to a list of topics that all carry `T`, delivered to one handler in arrival order. Each message is tagged with the index of the topic it arrived on, so the handler can tell the sources apart. +A rig with N identical sensors would otherwise need N ports and N near-identical handlers. A topic funnel is one `Input` port wired to a list of topics that all carry `T`, delivered to one handler in arrival order. Any input accepts a funnel — on the launch line the port's value is an array of topics rather than a string, and nothing in the module changes. A handler that needs to tell the sources apart opts in with the `meta` flag and receives a `Metadata` alongside each message: ```rust #[derive(Module)] struct MultiCam { - #[topic_funnel(decode = Image::decode)] - cameras: TopicFunnel, + #[input(decode = Image::decode, meta)] + cameras: Input, } impl MultiCam { - async fn handle_cameras(&mut self, index: usize, image: Image) { - let topic = self.cameras.topic(index); + async fn handle_cameras(&mut self, image: Image, meta: Metadata) { + // meta.index: position in the topic list; meta.topic: the topic itself } } ``` @@ -139,13 +138,15 @@ MultiCam.blueprint( The group itself is not a port with a stream of its own: python hands the wired entries' channels to the native process, which subscribes them directly. On the launch line the port's value is an array rather than a string. A group configured with no names still claims its port but never yields. -`topic_funnels` lives on `ModuleConfig`, so a plain Python `Module` takes the same field. There the module subscribes the group itself and dispatches to `async def handle_(self, index, msg)`, matching the Rust handler signature: +`topic_funnels` lives on `ModuleConfig`, so a plain Python `Module` takes the same field. There the module subscribes the group itself and dispatches to `async def handle_`. Opting into metadata is by signature — a one-parameter handler just gets the message, a two-parameter handler also gets a `Metadata` with `index` (position in `names`) and `name` (the stream name as declared, pre-remapping): ```python class MultiCam(Module): - async def handle_cameras(self, index: int, image: Image) -> None: ... + async def handle_cameras(self, image: Image, meta: Metadata) -> None: ... ``` +The same applies to any `handle_` for a plain `In` port, where the metadata is always `index=0, name=`. + The whole group shares one dispatcher, so the handler is never re-entered, and its mailbox holds the latest unprocessed message per stream rather than one slot for the group — a chatty camera cannot starve the others. ## Transforms diff --git a/native/rust/dimos-module-macros/src/lib.rs b/native/rust/dimos-module-macros/src/lib.rs index ae5bf4878c..b7d8fcda51 100644 --- a/native/rust/dimos-module-macros/src/lib.rs +++ b/native/rust/dimos-module-macros/src/lib.rs @@ -17,10 +17,7 @@ use proc_macro2::TokenStream as TokenStream2; use quote::{format_ident, quote}; use syn::{parse_macro_input, Data, DeriveInput, Field, Fields, Ident, LitStr, Path, Type}; -#[proc_macro_derive( - Module, - attributes(input, topic_funnel, output, io, config, tf, module) -)] +#[proc_macro_derive(Module, attributes(input, output, io, config, tf, module))] pub fn derive_module(input: TokenStream) -> TokenStream { let input = parse_macro_input!(input as DeriveInput); match expand(input) { @@ -178,7 +175,7 @@ fn is_option(ty: &Type) -> bool { } const ONE_ATTR_ONLY: &str = "field has multiple module attributes; only one of #[input], \ - #[topic_funnel], #[output], #[io], #[config], #[tf] is allowed"; + #[output], #[io], #[config], #[tf] is allowed"; /// What a `#[tf]` port carries, for the Cargo.toml registry cross-check. const TF_PAYLOAD_TYPE: &str = "TFMessage"; @@ -187,10 +184,7 @@ enum FieldKind { Input { decode: Path, handler: Ident, - }, - TopicFunnel { - decode: Path, - handler: Ident, + wants_meta: bool, }, Output { encode: Path, @@ -330,9 +324,6 @@ fn expand(input: DeriveInput) -> syn::Result { FieldKind::Input { decode, .. } => { quote!(#name: builder.input(#name_str, #decode)) } - FieldKind::TopicFunnel { decode, .. } => { - quote!(#name: builder.topic_funnel(#name_str, #decode)) - } FieldKind::Output { encode } => { quote!(#name: builder.output(#name_str, #encode)) } @@ -351,16 +342,20 @@ fn expand(input: DeriveInput) -> syn::Result { .filter_map(|f| { let name = f.name; match &f.kind { + FieldKind::Input { + handler, + wants_meta: true, + .. + } => Some(quote!( + ::core::option::Option::Some((msg, meta)) = self.#name.recv_meta() => { + self.#handler(msg, meta).await + } + )), FieldKind::Input { handler, .. } | FieldKind::Io { handler, .. } => Some(quote!( ::core::option::Option::Some(msg) = self.#name.recv() => { self.#handler(msg).await } )), - FieldKind::TopicFunnel { handler, .. } => Some(quote!( - ::core::option::Option::Some((index, msg)) = self.#name.recv() => { - self.#handler(index, msg).await - } - )), _ => None, } }) @@ -447,7 +442,7 @@ fn port_decls(classified: &[ClassifiedField], want_input: bool) -> Vec classified .iter() .filter(|f| match f.kind { - FieldKind::Input { .. } | FieldKind::TopicFunnel { .. } => want_input, + FieldKind::Input { .. } => want_input, FieldKind::Output { .. } => !want_input, // A `#[tf]` field is a port like any other as far as the registry // goes: bake has to put the topic in the host's map or the module @@ -571,6 +566,7 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { } let mut decode: Option = None; let mut handler: Option = None; + let mut wants_meta = false; attr.parse_nested_meta(|meta| { if meta.path.is_ident("decode") { decode = Some(meta.value()?.parse()?); @@ -578,10 +574,12 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { handler = Some(meta.value()?.parse()?); } else if meta.path.is_ident("msg") { msg = Some(meta.value()?.parse::()?.value()); + } else if meta.path.is_ident("meta") { + wants_meta = true; } else { return Err(meta.error( "unrecognized #[input] argument; expected `decode = ...`, \ - `handler = ...` or `msg = ...`", + `handler = ...`, `msg = ...` or `meta`", )); } Ok(()) @@ -589,33 +587,11 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { let decode = decode .ok_or_else(|| syn::Error::new_spanned(attr, "#[input] requires `decode = ...`"))?; let handler = handler.unwrap_or_else(|| format_ident!("handle_{}", name)); - found = Some(FieldKind::Input { decode, handler }); - } else if path.is_ident("topic_funnel") { - if found.is_some() { - return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); - } - let mut decode: Option = None; - let mut handler: Option = None; - attr.parse_nested_meta(|meta| { - if meta.path.is_ident("decode") { - decode = Some(meta.value()?.parse()?); - } else if meta.path.is_ident("handler") { - handler = Some(meta.value()?.parse()?); - } else if meta.path.is_ident("msg") { - msg = Some(meta.value()?.parse::()?.value()); - } else { - return Err(meta.error( - "unrecognized #[topic_funnel] argument; expected `decode = ...`, \ - `handler = ...` or `msg = ...`", - )); - } - Ok(()) - })?; - let decode = decode.ok_or_else(|| { - syn::Error::new_spanned(attr, "#[topic_funnel] requires `decode = ...`") - })?; - let handler = handler.unwrap_or_else(|| format_ident!("handle_{}", name)); - found = Some(FieldKind::TopicFunnel { decode, handler }); + found = Some(FieldKind::Input { + decode, + handler, + wants_meta, + }); } else if path.is_ident("output") { if found.is_some() { return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); diff --git a/native/rust/dimos-module/src/lib.rs b/native/rust/dimos-module/src/lib.rs index 73d9fce62b..1390402dd7 100644 --- a/native/rust/dimos-module/src/lib.rs +++ b/native/rust/dimos-module/src/lib.rs @@ -30,7 +30,7 @@ pub use dimos_module_macros::{native_config, Module}; pub use host::{host_main, HostSpec, ModuleEntry}; pub use lcm::LcmTransport; pub use module::{ - run, Builder, Input, Io, Module, ModuleConfig, NativeConfig, NoConfig, Output, TopicFunnel, + run, Builder, Input, Io, Metadata, Module, ModuleConfig, NativeConfig, NoConfig, Output, }; pub use tf::{Lookup, Tf, Transform}; pub use transport::{SharedTransport, Transport}; diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index 9af85ff09f..58f622eb01 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -74,8 +74,8 @@ pub(crate) trait Route: Send + Sync { fn try_dispatch(&self, data: &[u8]); } -/// Decodes a frame into whatever the port's channel carries. Boxed because a -/// group route tags the message with the index of the topic it arrived on. +/// Decodes a frame into whatever the port's channel carries. Boxed because an +/// input route tags the message with the index of the topic it arrived on. type Decode = Box io::Result + Send + Sync>; struct TypedRoute { @@ -113,30 +113,35 @@ impl Route for TypedRoute { } } } -pub struct Input { +/// Which topic of an input a message arrived on, for handlers that opt in +/// with `#[input(meta)]`. +#[derive(Clone, Debug)] +pub struct Metadata { + pub index: usize, pub topic: String, - receiver: mpsc::Receiver, -} - -impl Input { - pub async fn recv(&mut self) -> Option { - self.receiver.recv().await - } } -/// Several topics of one message type, fanned into a single handler. -/// -/// The launch line gives the port an array of topics instead of one, and every -/// message is tagged with the index of the topic it arrived on so a handler can -/// tell a rig's cameras apart. A group configured with no topics never yields. -pub struct TopicFunnel { +/// A port wired to one topic, or to every topic the launch line lists under +/// it — a topic funnel. Every message carries the index of the topic it +/// arrived on; `recv` drops it, `recv_meta` hands it over. A port given an +/// empty array never yields. +pub struct Input { pub topics: Vec, receiver: mpsc::Receiver<(usize, T)>, } -impl TopicFunnel { - pub async fn recv(&mut self) -> Option<(usize, T)> { - self.receiver.recv().await +impl Input { + pub async fn recv(&mut self) -> Option { + self.receiver.recv().await.map(|(_, msg)| msg) + } + + pub async fn recv_meta(&mut self) -> Option<(T, Metadata)> { + let (index, msg) = self.receiver.recv().await?; + let meta = Metadata { + index, + topic: self.topics[index].clone(), + }; + Some((msg, meta)) } pub fn topic(&self, index: usize) -> &str { @@ -411,9 +416,17 @@ impl Builder { .unwrap_or_else(|| format!("/{port}")) } - fn group_for(&mut self, port: &str) -> Vec { + /// One topic, or every topic an array under `port` lists. An unnamed port + /// falls back to `/{port}`; an empty array is a port with no sources. + fn input_topics_for(&mut self, port: &str) -> Vec { self.requested.insert(port.to_string()); - self.topics.grouped.get(port).cloned().unwrap_or_default() + if let Some(topic) = self.topics.single.get(port) { + vec![topic.clone()] + } else if let Some(group) = self.topics.grouped.get(port) { + group.clone() + } else { + vec![format!("/{port}")] + } } // A mismatch is dead wiring: an unclaimed topic reaches no port, and an @@ -473,25 +486,13 @@ impl Builder { port: &str, decode: fn(&[u8]) -> io::Result, ) -> Input { - let topic = self.topic_for(port); - let receiver = self.add_route(&topic, decode); - Input { topic, receiver } - } - - /// A port wired to every topic the launch line lists under `port`, all of - /// one message type, delivered to one handler in arrival order. - pub fn topic_funnel( - &mut self, - port: &str, - decode: fn(&[u8]) -> io::Result, - ) -> TopicFunnel { - let topics = self.group_for(port); + let topics = self.input_topics_for(port); let (sender, receiver) = mpsc::channel(INPUT_CHANNEL_CAPACITY); for (index, topic) in topics.iter().enumerate() { let tag = Box::new(move |bytes: &[u8]| decode(bytes).map(|msg| (index, msg))); self.push_route(topic, tag, sender.clone()); } - TopicFunnel { topics, receiver } + Input { topics, receiver } } pub fn output(&mut self, port: &str, encode: fn(&T) -> Vec) -> Output { @@ -1036,14 +1037,14 @@ mod tests { fn input_uses_mapped_topic() { let mut builder = builder_with_topics(&[("data", "/test/data")]); let input = builder.input("data", |b| Ok(b.to_vec())); - assert_eq!(input.topic, "/test/data"); + assert_eq!(input.topics, ["/test/data"]); } #[test] fn input_falls_back_to_slash_port_when_unmapped() { let mut builder = builder_with_topics(&[]); let input = builder.input("data", |b| Ok(b.to_vec())); - assert_eq!(input.topic, "/data"); + assert_eq!(input.topics, ["/data"]); } #[test] @@ -1053,10 +1054,10 @@ mod tests { assert_eq!(output.topic, "/robot/cmd_vel"); } - // input groups + // topic funnels: a port wired to an array of topics #[test] - fn a_port_given_an_array_of_topics_becomes_a_group() { + fn a_port_given_an_array_of_topics_becomes_a_funnel() { let json = r#"{"topics": {"cams": ["/cam0", "/cam1"], "odom": "/odom"}, "config": null}"#; let (topics, _config) = parse_config_json::<()>(json).unwrap(); assert_eq!(topics.grouped["cams"], ["/cam0", "/cam1"]); @@ -1064,50 +1065,48 @@ mod tests { } #[test] - fn a_group_entry_that_is_not_a_string_is_rejected() { + fn a_funnel_entry_that_is_not_a_string_is_rejected() { let json = r#"{"topics": {"cams": ["/cam0", 7]}, "config": null}"#; - let err = parse_config_json::<()>(json).expect_err("a non-string group entry is invalid"); + let err = parse_config_json::<()>(json).expect_err("a non-string funnel entry is invalid"); assert!(err.to_string().contains("cams"), "{err}"); } #[test] - fn a_group_subscribes_every_topic_it_was_given() { + fn a_funneled_input_subscribes_every_topic_it_was_given() { let mut builder = Builder::new(grouped_topics("cams", &["/cam0", "/cam1"])); - let group = builder.topic_funnel("cams", |b| Ok(b.to_vec())); - assert_eq!(group.topics, ["/cam0", "/cam1"]); + let input = builder.input("cams", |b| Ok(b.to_vec())); + assert_eq!(input.topics, ["/cam0", "/cam1"]); assert_eq!(builder.routes.get("/cam0").map(Vec::len), Some(1)); assert_eq!(builder.routes.get("/cam1").map(Vec::len), Some(1)); - builder.enforce_topics_match_ports().expect("group claimed"); + builder.enforce_topics_match_ports().expect("port claimed"); } - /// One handler sees every topic, so the index is the only thing that says - /// which camera of a rig a frame came from. + /// One handler sees every topic, so the metadata is the only thing that + /// says which camera of a rig a frame came from. #[tokio::test] - async fn a_group_tags_each_message_with_its_topic_index() { + async fn recv_meta_names_the_topic_a_message_arrived_on() { let mut builder = Builder::new(grouped_topics("cams", &["/cam0", "/cam1"])); - let mut group = builder.topic_funnel("cams", |b| Ok(b.to_vec())); + let mut input = builder.input("cams", |b| Ok(b.to_vec())); builder.routes["/cam1"][0].try_dispatch(b"second"); builder.routes["/cam0"][0].try_dispatch(b"first"); - assert_eq!( - group.recv().await.expect("cam1 frame"), - (1, b"second".to_vec()) - ); - assert_eq!( - group.recv().await.expect("cam0 frame"), - (0, b"first".to_vec()) - ); - assert_eq!(group.topic(1), "/cam1"); + let (msg, meta) = input.recv_meta().await.expect("cam1 frame"); + assert_eq!(msg, b"second"); + assert_eq!(meta.index, 1); + assert_eq!(meta.topic, "/cam1"); + + // recv drops the tag for handlers that treat the funnel as one stream. + assert_eq!(input.recv().await.expect("cam0 frame"), b"first"); } #[test] - fn an_empty_group_still_claims_its_port() { + fn an_empty_funnel_still_claims_its_port() { let mut builder = Builder::new(grouped_topics("cams", &[])); - let group = builder.topic_funnel("cams", |b| Ok(b.to_vec())); - assert!(group.is_empty()); + let input = builder.input("cams", |b| Ok(b.to_vec())); + assert!(input.is_empty()); assert!(builder.routes.is_empty()); - builder.enforce_topics_match_ports().expect("group claimed"); + builder.enforce_topics_match_ports().expect("port claimed"); } #[test] @@ -1465,22 +1464,22 @@ mod tests { #[derive(crate::Module)] struct Rig { - #[topic_funnel(decode = decode)] - cams: crate::TopicFunnel, + #[input(decode = decode, meta)] + cams: crate::Input, #[output(encode = encode)] seen: crate::Output, } impl Rig { - async fn handle_cams(&mut self, index: usize, msg: Msg) { - let mut tagged = vec![index as u8]; + async fn handle_cams(&mut self, msg: Msg, meta: crate::Metadata) { + let mut tagged = vec![meta.index as u8]; tagged.extend(msg.0); self.seen.publish(&Msg(tagged)).await.expect("publish"); } } #[tokio::test] - async fn topic_funnel_field_hands_its_handler_the_topic_index() { + async fn a_meta_input_hands_its_handler_the_topic_index() { let mut builder = Builder::new(Topics { single: HashMap::from([("seen".to_string(), "/seen".to_string())]), grouped: HashMap::from([( From de45dcb3d8cdd8738df7f24a12a5ee6dd35a38a1 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 29 Aug 2026 16:20:36 -0700 Subject: [PATCH 06/16] Attach per-topic info to funnel entries via Metadata TopicFunnel.names may be a dict mapping each stream name to arbitrary info (is this camera rectified, what are the depth units) that the handler receives on Metadata.info. On the launch line such an entry becomes a {"topic", "info"} object instead of a plain string; rust exposes it as serde_json::Value, python as a dict. Entries without info stay plain strings, so existing wiring is untouched. --- dimos/core/module.py | 40 ++++++- dimos/core/native_module.py | 34 ++++-- dimos/core/test_async_module_topic_funnels.py | 18 ++- dimos/core/test_native_module.py | 25 ++++ native/rust/README.md | 5 +- native/rust/dimos-module/src/module.rs | 112 +++++++++++++++--- 6 files changed, 194 insertions(+), 40 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index 5c716a6f46..9548459952 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -13,7 +13,7 @@ # limitations under the License. import asyncio from collections.abc import AsyncGenerator, Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from functools import partial import inspect import json @@ -31,7 +31,7 @@ get_type_hints, ) -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field, field_validator, model_validator from reactivex.disposable import CompositeDisposable, Disposable from dimos.core.core import T, rpc @@ -112,10 +112,27 @@ class TopicFunnel(BaseModel): producers, `.remappings()` and `.namespace()` rewrite it, and transport pins apply. A python module subscribes the wired streams itself; a native module hands their channels to its subprocess, which does the subscribing. + + `names` may also be a dict attaching arbitrary per-stream info that the + handler receives on `Metadata.info`: + `TopicFunnel(names={"depth": {"meters_per_unit": 0.001}, "color": {}})`. """ names: list[str] msg_type: type[Any] | None = None + info: dict[str, dict[str, Any]] = Field(default_factory=dict) + + @model_validator(mode="before") + @classmethod + def _split_names_dict(cls, data: Any) -> Any: + if isinstance(data, dict) and isinstance(data.get("names"), dict): + names_dict = data["names"] + data = { + **data, + "names": list(names_dict), + "info": {name: value for name, value in names_dict.items() if value}, + } + return data @field_validator("names") @classmethod @@ -128,17 +145,26 @@ def _reject_topic_strings(cls, names: list[str]) -> list[str]: ) return names + @model_validator(mode="after") + def _info_keys_must_be_names(self) -> "TopicFunnel": + unknown = [name for name in self.info if name not in self.names] + if unknown: + raise ValueError(f"topic funnel info keys are not in names: {unknown}") + return self + @dataclass(frozen=True) class Metadata: """Which stream a message arrived on, for handlers that take `(msg, meta)`. `index` is the position in the funnel's `names` (0 for a plain input); - `name` is the stream name as the module declared it, pre-remapping. + `name` is the stream name as the module declared it, pre-remapping; + `info` is whatever the funnel's `TopicFunnel.info` attached to that name. """ index: int name: str + info: dict[str, Any] = field(default_factory=dict) def _handler_wants_metadata(handler: Callable[..., Any], label: str) -> bool: @@ -778,14 +804,18 @@ def _bind_topic_funnels(self) -> None: "(topic funnels have no sync path)" ) if _handler_wants_metadata(handler, f"{type(self).__name__}.handle_{port}"): + metas = tuple( + Metadata(index=index, name=name, info=group.info.get(name, {})) + for index, name in enumerate(group.names) + ) async def invoke( index: int, msg: Any, _handler: Callable[[Any, Metadata], Any] = handler, - _names: tuple[str, ...] = tuple(group.names), + _metas: tuple[Metadata, ...] = metas, ) -> None: - await _handler(msg, Metadata(index=index, name=_names[index])) + await _handler(msg, _metas[index]) else: async def invoke( # type: ignore[misc] diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index c133dafe1d..7fa856d026 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -261,17 +261,24 @@ def _session(self) -> SessionConfig: # A blueprint builds its config before global config is settled. return pinned.rebased() - def _argv(self, topics: dict[str, str | list[str]]) -> list[str]: - """The command line the native process is spawned with.""" + def _argv(self, topics: dict[str, str | list[str | dict[str, Any]]]) -> list[str]: + """The command line the native process is spawned with. Per-topic funnel + info travels only on the stdin JSON, so an entry object contributes just + its topic here.""" cmd = [self.config.executable] - for name, topic_str in topics.items(): - joined = ",".join(topic_str) if isinstance(topic_str, list) else topic_str + for name, value in topics.items(): + if isinstance(value, list): + joined = ",".join( + entry["topic"] if isinstance(entry, dict) else entry for entry in value + ) + else: + joined = value cmd.extend([f"--{name}", joined]) cmd.extend(self.config.to_cli_args()) cmd.extend(self.config.extra_args) return cmd - def _stdin_blob(self, topics: dict[str, str | list[str]]) -> bytes: + def _stdin_blob(self, topics: dict[str, str | list[str | dict[str, Any]]]) -> bytes: """The JSON line the native process reads its launch from.""" config_dict = self.config.to_config_dict() blob: dict[str, Any] = { @@ -517,8 +524,8 @@ def _maybe_build(self) -> None: duration_sec=round(build_elapsed, 3), ) - def _collect_topics(self) -> dict[str, str | list[str]]: - topics: dict[str, str | list[str]] = {} + def _collect_topics(self) -> dict[str, str | list[str | dict[str, Any]]]: + topics: dict[str, str | list[str | dict[str, Any]]] = {} for name in list(self.inputs) + list(self.outputs) + list(self.ios): stream = getattr(self, name, None) if stream is None: @@ -535,15 +542,16 @@ def _collect_topics(self) -> dict[str, str | list[str]]: f"[{self._module_label}] topic funnel {port!r} collides with the " "port of the same name declared as a stream" ) - channels = [] + entries: list[str | dict[str, Any]] = [] for name in group.names: transport = getattr(self._funnel_streams[name], "_transport", None) channel = getattr(transport, "channel", None) - # Unwired (standalone use): the channel the default transport lands on. - channels.append( - channel if channel is not None else channel_for(name, group.msg_type) - ) - topics[port] = channels + if channel is None: + # Unwired (standalone use): the channel the default transport lands on. + channel = channel_for(name, group.msg_type) + info = group.info.get(name) + entries.append(channel if info is None else {"topic": channel, "info": info}) + topics[port] = entries return topics def _collect_output_qos(self) -> dict[str, dict[str, str]]: diff --git a/dimos/core/test_async_module_topic_funnels.py b/dimos/core/test_async_module_topic_funnels.py index 28fdf1ec05..aa10f68e27 100644 --- a/dimos/core/test_async_module_topic_funnels.py +++ b/dimos/core/test_async_module_topic_funnels.py @@ -27,15 +27,19 @@ class FanInModule(Module): tagged: Out[int] named: Out[str] + scaled: Out[float] async def handle_sensors(self, value: int, meta: Metadata) -> None: self.tagged.publish(meta.index * 100 + value) self.named.publish(meta.name) + self.scaled.publish(value * meta.info.get("scale", 1.0)) @pytest.fixture def start_fan_in_module(each_transport): - blueprint = FanInModule.blueprint(topic_funnels={"sensors": TopicFunnel(names=["s0", "s1"])}) + blueprint = FanInModule.blueprint( + topic_funnels={"sensors": TopicFunnel(names={"s0": {"scale": 0.5}, "s1": {}})} + ) coordinator = ModuleCoordinator.build(blueprint) yield coordinator.stop() @@ -48,6 +52,7 @@ def group_transports(each_transport): make_transport("s1"), make_transport("tagged"), make_transport("named"), + make_transport("scaled"), ] for transport in transports: transport.start() @@ -57,19 +62,23 @@ def group_transports(each_transport): def test_topic_funnel_tags_each_message_with_its_stream(start_fan_in_module, group_transports): - s0, s1, tagged, named = group_transports + s0, s1, tagged, named, scaled = group_transports queue: Queue[int] = Queue() names: Queue[str] = Queue() + scales: Queue[float] = Queue() tagged.subscribe(queue.put) named.subscribe(names.put) + scaled.subscribe(scales.put) s0.publish(7) assert queue.get(timeout=1.0) == 7 assert names.get(timeout=1.0) == "s0" + assert scales.get(timeout=1.0) == 3.5 s1.publish(7) assert queue.get(timeout=1.0) == 107 assert names.get(timeout=1.0) == "s1" + assert scales.get(timeout=1.0) == 7.0 @pytest.fixture @@ -115,6 +124,11 @@ def test_a_group_entry_cannot_collide_with_a_declared_stream(): FanInModule(topic_funnels={"sensors": TopicFunnel(names=["tagged"])}) +def test_funnel_info_keys_must_be_names(): + with pytest.raises(ValueError, match="not in names"): + TopicFunnel(names=["s0"], info={"s9": {"scale": 2.0}}) + + class PlainFanInModule(Module): """A funnel handler that doesn't ask for metadata just gets the message.""" diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index e411efdcf0..49a4e8b31e 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -229,6 +229,31 @@ def test_a_topic_funnel_reaches_the_native_process_as_a_list(monkeypatch) -> Non module.stop() +def test_funnel_info_rides_the_launch_line_but_not_the_argv(monkeypatch) -> None: + """Per-topic info becomes a {topic, info} entry on stdin; argv keeps only topics.""" + monkeypatch.setattr(native_module_mod.global_config, "transport", "lcm") + module = StubNativeModule( + executable=_ECHO, + topic_funnels={ + "cams": TopicFunnel( + names={"cam0/imu": {"rectified": True}, "cam1/imu": {}}, msg_type=Imu + ) + }, + ) + try: + topics = module._collect_topics() + assert topics["cams"] == [ + {"topic": "/cam0/imu#sensor_msgs.Imu", "info": {"rectified": True}}, + "/cam1/imu#sensor_msgs.Imu", + ] + assert module._argv(topics)[1:3] == [ + "--cams", + "/cam0/imu#sensor_msgs.Imu,/cam1/imu#sensor_msgs.Imu", + ] + finally: + module.stop() + + def test_a_wired_topic_funnel_entry_uses_its_transport() -> None: """Remapping/pins arrive as set_transport on the entry, and the launch line follows.""" module = StubNativeModule( diff --git a/native/rust/README.md b/native/rust/README.md index 8a0ff55c65..9c57f6445b 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -121,7 +121,8 @@ struct MultiCam { impl MultiCam { async fn handle_cameras(&mut self, image: Image, meta: Metadata) { - // meta.index: position in the topic list; meta.topic: the topic itself + // meta.index: position in the topic list; meta.topic: the topic itself; + // meta.info: per-topic JSON the coordinator attached (Null when none) } } ``` @@ -134,6 +135,8 @@ MultiCam.blueprint( ) ``` +`names` may also be a dict attaching arbitrary per-stream info — `names={"left_cam": {"rectified": True}, "right_cam": {}}` — which reaches the handler as `Metadata.info` (a dict in python, `serde_json::Value` in rust; on the launch line such an entry is a `{"topic": ..., "info": ...}` object instead of a plain string). + `names` are stream names, not backend topics — a leading `/` is rejected. Each entry becomes a synthetic `In` stream on the python side, so the blueprint machinery treats it like a declared port: autoconnect matches it against producers' `Out` streams, `.remappings()` and `.namespace()` rewrite it, and blueprint transport pins apply. The same blueprint runs unchanged over LCM or zenoh. The group itself is not a port with a stream of its own: python hands the wired entries' channels to the native process, which subscribes them directly. On the launch line the port's value is an array rather than a string. A group configured with no names still claims its port but never yields. diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index 58f622eb01..45fbb1d3dd 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -114,11 +114,13 @@ impl Route for TypedRoute { } } /// Which topic of an input a message arrived on, for handlers that opt in -/// with `#[input(meta)]`. +/// with `#[input(meta)]`. `info` is whatever the coordinator attached to that +/// topic on the launch line (`Null` when nothing was). #[derive(Clone, Debug)] pub struct Metadata { pub index: usize, pub topic: String, + pub info: serde_json::Value, } /// A port wired to one topic, or to every topic the launch line lists under @@ -127,6 +129,7 @@ pub struct Metadata { /// empty array never yields. pub struct Input { pub topics: Vec, + pub infos: Vec, receiver: mpsc::Receiver<(usize, T)>, } @@ -140,6 +143,7 @@ impl Input { let meta = Metadata { index, topic: self.topics[index].clone(), + info: self.infos[index].clone(), }; Some((msg, meta)) } @@ -148,6 +152,10 @@ impl Input { &self.topics[index] } + pub fn info(&self, index: usize) -> &serde_json::Value { + &self.infos[index] + } + pub fn len(&self) -> usize { self.topics.len() } @@ -201,12 +209,29 @@ pub(crate) async fn publish_encoded( .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "background task gone")) } +/// One source of a funneled input: its topic plus whatever per-topic info the +/// coordinator attached on the launch line. +#[derive(Clone, Debug)] +pub(crate) struct TopicEntry { + topic: String, + info: serde_json::Value, +} + +impl TopicEntry { + fn bare(topic: impl Into) -> Self { + Self { + topic: topic.into(), + info: serde_json::Value::Null, + } + } +} + /// The port-to-topic wiring from the launch line. A port names one topic, or an /// array of same-typed topics that one handler sees. #[derive(Clone, Debug, Default)] pub(crate) struct Topics { single: HashMap, - grouped: HashMap>, + grouped: HashMap>, } impl Topics { @@ -216,7 +241,9 @@ impl Topics { /// Every wire channel the module touches, groups flattened. pub(crate) fn channels(&self) -> impl Iterator { - self.single.values().chain(self.grouped.values().flatten()) + self.single + .values() + .chain(self.grouped.values().flatten().map(|entry| &entry.topic)) } } @@ -239,12 +266,29 @@ fn parse_topics(json: &serde_json::Value) -> io::Result { serde_json::Value::Array(items) => { let group = items .iter() - .map(|item| { - item.as_str() - .map(str::to_string) - .ok_or_else(|| invalid("is a group, so every entry must be a string")) + .map(|item| match item { + serde_json::Value::String(topic) => Ok(TopicEntry::bare(topic)), + serde_json::Value::Object(fields) => { + let topic = + fields + .get("topic") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + invalid("has an entry object without a string 'topic'") + })?; + Ok(TopicEntry { + topic: topic.to_string(), + info: fields + .get("info") + .cloned() + .unwrap_or(serde_json::Value::Null), + }) + } + _ => Err(invalid( + "is a group, so every entry must be a string or a {topic, info} object", + )), }) - .collect::>>()?; + .collect::>>()?; topics.grouped.insert(port.clone(), group); } _ => return Err(invalid("must be a string or an array of strings")), @@ -418,14 +462,14 @@ impl Builder { /// One topic, or every topic an array under `port` lists. An unnamed port /// falls back to `/{port}`; an empty array is a port with no sources. - fn input_topics_for(&mut self, port: &str) -> Vec { + fn input_entries_for(&mut self, port: &str) -> Vec { self.requested.insert(port.to_string()); if let Some(topic) = self.topics.single.get(port) { - vec![topic.clone()] + vec![TopicEntry::bare(topic)] } else if let Some(group) = self.topics.grouped.get(port) { group.clone() } else { - vec![format!("/{port}")] + vec![TopicEntry::bare(format!("/{port}"))] } } @@ -486,13 +530,21 @@ impl Builder { port: &str, decode: fn(&[u8]) -> io::Result, ) -> Input { - let topics = self.input_topics_for(port); + let entries = self.input_entries_for(port); let (sender, receiver) = mpsc::channel(INPUT_CHANNEL_CAPACITY); - for (index, topic) in topics.iter().enumerate() { + for (index, entry) in entries.iter().enumerate() { let tag = Box::new(move |bytes: &[u8]| decode(bytes).map(|msg| (index, msg))); - self.push_route(topic, tag, sender.clone()); + self.push_route(&entry.topic, tag, sender.clone()); + } + let (topics, infos) = entries + .into_iter() + .map(|entry| (entry.topic, entry.info)) + .unzip(); + Input { + topics, + infos, + receiver, } - Input { topics, receiver } } pub fn output(&mut self, port: &str, encode: fn(&T) -> Vec) -> Output { @@ -1012,7 +1064,7 @@ mod tests { single: HashMap::new(), grouped: HashMap::from([( port.to_string(), - group.iter().map(|t| t.to_string()).collect(), + group.iter().map(|t| TopicEntry::bare(*t)).collect(), )]), } } @@ -1060,7 +1112,8 @@ mod tests { fn a_port_given_an_array_of_topics_becomes_a_funnel() { let json = r#"{"topics": {"cams": ["/cam0", "/cam1"], "odom": "/odom"}, "config": null}"#; let (topics, _config) = parse_config_json::<()>(json).unwrap(); - assert_eq!(topics.grouped["cams"], ["/cam0", "/cam1"]); + let cams: Vec<&str> = topics.grouped["cams"].iter().map(|e| &*e.topic).collect(); + assert_eq!(cams, ["/cam0", "/cam1"]); assert_eq!(topics.single["odom"], "/odom"); } @@ -1071,6 +1124,24 @@ mod tests { assert!(err.to_string().contains("cams"), "{err}"); } + #[test] + fn a_funnel_entry_object_carries_per_topic_info() { + let json = r#"{"topics": {"cams": + ["/cam0", {"topic": "/cam1", "info": {"rectified": true}}]}, "config": null}"#; + let (topics, _config) = parse_config_json::<()>(json).unwrap(); + let cams = &topics.grouped["cams"]; + assert_eq!(cams[0].info, serde_json::Value::Null); + assert_eq!(cams[1].topic, "/cam1"); + assert_eq!(cams[1].info["rectified"], true); + } + + #[test] + fn a_funnel_entry_object_without_a_topic_is_rejected() { + let json = r#"{"topics": {"cams": [{"info": {"rectified": true}}]}, "config": null}"#; + let err = parse_config_json::<()>(json).expect_err("an entry object needs a topic"); + assert!(err.to_string().contains("'topic'"), "{err}"); + } + #[test] fn a_funneled_input_subscribes_every_topic_it_was_given() { let mut builder = Builder::new(grouped_topics("cams", &["/cam0", "/cam1"])); @@ -1085,7 +1156,9 @@ mod tests { /// says which camera of a rig a frame came from. #[tokio::test] async fn recv_meta_names_the_topic_a_message_arrived_on() { - let mut builder = Builder::new(grouped_topics("cams", &["/cam0", "/cam1"])); + let mut topics = grouped_topics("cams", &["/cam0", "/cam1"]); + topics.grouped.get_mut("cams").unwrap()[1].info = serde_json::json!({"rectified": true}); + let mut builder = Builder::new(topics); let mut input = builder.input("cams", |b| Ok(b.to_vec())); builder.routes["/cam1"][0].try_dispatch(b"second"); @@ -1095,6 +1168,7 @@ mod tests { assert_eq!(msg, b"second"); assert_eq!(meta.index, 1); assert_eq!(meta.topic, "/cam1"); + assert_eq!(meta.info["rectified"], true); // recv drops the tag for handlers that treat the funnel as one stream. assert_eq!(input.recv().await.expect("cam0 frame"), b"first"); @@ -1484,7 +1558,7 @@ mod tests { single: HashMap::from([("seen".to_string(), "/seen".to_string())]), grouped: HashMap::from([( "cams".to_string(), - vec!["/cam0".to_string(), "/cam1".to_string()], + vec![TopicEntry::bare("/cam0"), TopicEntry::bare("/cam1")], )]), }); let mut rig = Rig::build(&mut builder, NoConfig); From e13b3a36ec68f552ca776c4df9c8b9d7f2bca7f0 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 31 Aug 2026 13:33:37 -0700 Subject: [PATCH 07/16] Fix mypy errors in the topic funnel handler binding The handler captured as a closure default was still typed Any | None from the getattr, and baked_host's _argv/_stdin_blob overrides never widened to the funnel-aware topics type. --- dimos/core/baked_host.py | 10 +++++----- dimos/core/module.py | 22 ++++++++++++---------- dimos/core/native_module.py | 12 ++++++++---- 3 files changed, 25 insertions(+), 19 deletions(-) diff --git a/dimos/core/baked_host.py b/dimos/core/baked_host.py index e058d0ce7e..afa717e9da 100644 --- a/dimos/core/baked_host.py +++ b/dimos/core/baked_host.py @@ -20,11 +20,11 @@ from collections.abc import Mapping import json import sys -from typing import get_args, get_origin, get_type_hints +from typing import Any, get_args, get_origin, get_type_hints from pydantic import Field, create_model -from dimos.core.native_module import NativeModule, NativeModuleConfig +from dimos.core.native_module import NativeModule, NativeModuleConfig, TopicsMap from dimos.core.stream import IO, In, Out @@ -83,7 +83,7 @@ class BakedHost(NativeModule): _members: Mapping[str, type[NativeModule]] = {} _remaps: Mapping[tuple[str, str], str] = {} - def _member_topics(self, instance: str, topics: Mapping[str, str]) -> dict[str, str]: + def _member_topics(self, instance: str, topics: Mapping[str, Any]) -> dict[str, Any]: member = self._members[instance] resolved = {} for port in _member_ports(member): @@ -92,14 +92,14 @@ def _member_topics(self, instance: str, topics: Mapping[str, str]) -> dict[str, resolved[port] = topics[name] return resolved - def _argv(self, topics: dict[str, str]) -> list[str]: + def _argv(self, topics: TopicsMap) -> list[str]: """A baked host takes its whole wiring on the launch line, so no flags. The binary rejects an unknown argument rather than ignoring it. """ return [self.config.executable, *self.config.extra_args] - def _stdin_blob(self, topics: dict[str, str]) -> bytes: + def _stdin_blob(self, topics: TopicsMap) -> bytes: sections: dict[str, object] = {} for instance in self._members: member_config = getattr(self.config, f"{instance}_config") diff --git a/dimos/core/module.py b/dimos/core/module.py index 9548459952..86c6624345 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -727,19 +727,20 @@ def _auto_bind_handlers(self) -> None: # Validate every handler before subscribing any of them. bindings: list[tuple[Any, Callable[[Any], Any]]] = [] for input_name, in_stream in {**self.inputs, **self.ios}.items(): - handler = getattr(self, f"handle_{input_name}", None) - if handler is None: + declared_handler = getattr(self, f"handle_{input_name}", None) + if declared_handler is None: continue # Async @rpc wraps the coroutine fn in a sync dispatcher. Unwrap it # so we subscribe the raw coroutine fn instead of the wrapper (which # would block on run_coroutine_threadsafe from the rx thread). - if hasattr(handler, "aio"): - handler = handler.aio.__get__(self, type(self)) - if not inspect.iscoroutinefunction(handler): + if hasattr(declared_handler, "aio"): + declared_handler = declared_handler.aio.__get__(self, type(self)) + if not inspect.iscoroutinefunction(declared_handler): raise TypeError( f"{type(self).__name__}.handle_{input_name} must be `async def` " "(use a manual self..subscribe(...) for sync handlers)" ) + handler: Callable[..., Any] = declared_handler if _handler_wants_metadata(handler, f"{type(self).__name__}.handle_{input_name}"): metadata = Metadata(index=0, name=input_name) @@ -793,16 +794,17 @@ def _bind_topic_funnels(self) -> None: so this is a no-op there. """ for port, group in self.config.topic_funnels.items(): - handler = getattr(self, f"handle_{port}", None) - if handler is None: + declared_handler = getattr(self, f"handle_{port}", None) + if declared_handler is None: continue - if hasattr(handler, "aio"): - handler = handler.aio.__get__(self, type(self)) - if not inspect.iscoroutinefunction(handler): + if hasattr(declared_handler, "aio"): + declared_handler = declared_handler.aio.__get__(self, type(self)) + if not inspect.iscoroutinefunction(declared_handler): raise TypeError( f"{type(self).__name__}.handle_{port} must be `async def` " "(topic funnels have no sync path)" ) + handler: Callable[..., Any] = declared_handler if _handler_wants_metadata(handler, f"{type(self).__name__}.handle_{port}"): metas = tuple( Metadata(index=index, name=name, info=group.info.get(name, {})) diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index 7fa856d026..0620670e98 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -64,6 +64,10 @@ class MyCppModule(NativeModule): from dimos.protocol.service.spec import SessionConfig from dimos.utils.logging_config import setup_logger +# Port name -> the wire channel(s) it is launched with. A funnel port carries a +# list, and an entry with per-topic info is an object rather than a bare string. +TopicsMap = dict[str, str | list[str | dict[str, Any]]] + if sys.platform.startswith("linux"): import ctypes from ctypes.util import find_library @@ -261,7 +265,7 @@ def _session(self) -> SessionConfig: # A blueprint builds its config before global config is settled. return pinned.rebased() - def _argv(self, topics: dict[str, str | list[str | dict[str, Any]]]) -> list[str]: + def _argv(self, topics: TopicsMap) -> list[str]: """The command line the native process is spawned with. Per-topic funnel info travels only on the stdin JSON, so an entry object contributes just its topic here.""" @@ -278,7 +282,7 @@ def _argv(self, topics: dict[str, str | list[str | dict[str, Any]]]) -> list[str cmd.extend(self.config.extra_args) return cmd - def _stdin_blob(self, topics: dict[str, str | list[str | dict[str, Any]]]) -> bytes: + def _stdin_blob(self, topics: TopicsMap) -> bytes: """The JSON line the native process reads its launch from.""" config_dict = self.config.to_config_dict() blob: dict[str, Any] = { @@ -524,8 +528,8 @@ def _maybe_build(self) -> None: duration_sec=round(build_elapsed, 3), ) - def _collect_topics(self) -> dict[str, str | list[str | dict[str, Any]]]: - topics: dict[str, str | list[str | dict[str, Any]]] = {} + def _collect_topics(self) -> TopicsMap: + topics: TopicsMap = {} for name in list(self.inputs) + list(self.outputs) + list(self.ios): stream = getattr(self, name, None) if stream is None: From 6be212ec5268459a3f6f3c8420ffc1078af8b77e Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 31 Aug 2026 16:07:17 -0700 Subject: [PATCH 08/16] Declare topic funnels through .remappings() Fanning several same-typed streams into one handler is stream routing, so it belongs alongside the other routing rewrites rather than in a module kwarg. A list of names (or a dict of name -> info) as a remapping value now fans that port in; a string still renames a single stream. .remappings() expands the fan-in into the blueprint atom right away: the port's StreamRef is replaced by one In per entry, typed from the declaration. remapping_map stays string-valued, so autoconnect, namespacing and transport pins are untouched. The declared port also supplies the message type, so TopicFunnel.msg_type is gone, and naming a port that does not exist is now an error instead of a silent no-op. --- dimos/core/coordination/blueprints.py | 93 ++++++++++++++++--- dimos/core/module.py | 73 +++++++-------- dimos/core/native_module.py | 20 ++-- dimos/core/test_async_module_topic_funnels.py | 52 +++++++---- dimos/core/test_native_module.py | 54 ++++------- native/rust/README.md | 20 ++-- 6 files changed, 186 insertions(+), 126 deletions(-) diff --git a/dimos/core/coordination/blueprints.py b/dimos/core/coordination/blueprints.py index f03e8fc096..42a4f1d887 100644 --- a/dimos/core/coordination/blueprints.py +++ b/dimos/core/coordination/blueprints.py @@ -27,7 +27,7 @@ if TYPE_CHECKING: from dimos.protocol.service.system_configurator.base import SystemConfigurator -from dimos.core.module import ModuleBase, is_module_type +from dimos.core.module import ModuleBase, TopicFunnel, is_module_type from dimos.core.stream import IO, In, Out, Transport from dimos.spec.utils import Spec, is_spec from dimos.utils.logging_config import setup_logger @@ -140,17 +140,6 @@ def create(cls, module: type[ModuleBase], kwargs: dict[str, Any]) -> Self: elif is_module_type(inner): module_refs.append(ModuleRef(name=name, spec=inner, optional=True)) - # Topic-funnel entries are synthetic In streams: part of stream wiring - # (autoconnect, remapping, namespacing, transport pins) without being - # ports of their own. - for group in (kwargs.get("topic_funnels") or {}).values(): - names = group["names"] if isinstance(group, dict) else group.names - msg_type = group.get("msg_type") if isinstance(group, dict) else group.msg_type - for name in names: - streams.append( - StreamRef(name=name, type=msg_type or Any, direction="in") # type: ignore[arg-type] - ) - instance_name = kwargs.get("instance_name") if instance_name is not None and not isinstance(instance_name, str): raise TypeError("instance_name must be a string or None") @@ -242,13 +231,39 @@ def global_config(self, **kwargs: Any) -> "Blueprint": def remappings( self, remappings: Sequence[ - tuple[type[ModuleBase] | str, str, str | type[ModuleBase] | type[Spec]] + tuple[ + type[ModuleBase] | str, + str, + str + | Sequence[str] + | Mapping[str, Mapping[str, Any]] + | type[ModuleBase] + | type[Spec], + ] ], ) -> "Blueprint": + """Rewrite what a module's streams and module refs connect to. + + A string (or Spec/Module class) renames one stream (or redirects one ref). + A list of names, or a dict of name -> arbitrary info, instead fans several + same-typed streams into that one port, so its handler serves all of them: + + .remappings([ + (Fuser, "scan", ["front_lidar", "rear_lidar"]), + (Fuser, "image", {"depth": {"meters_per_unit": 0.001}, "color": {}}), + ]) + """ remappings_dict = dict(self.remapping_map) + atoms = list(self.blueprints) for module, old, new in remappings: - remappings_dict[(self._instance_key(module), old)] = new - return replace(self, remapping_map=MappingProxyType(remappings_dict)) + instance_key = self._instance_key(module) + if isinstance(new, (str, type)): + remappings_dict[instance_key, old] = new + else: + atoms = _fan_in(atoms, instance_key, old, TopicFunnel.of(new)) + return replace( + self, blueprints=tuple(atoms), remapping_map=MappingProxyType(remappings_dict) + ) def _instance_key(self, module: type[ModuleBase] | str) -> str: if isinstance(module, str): @@ -392,6 +407,54 @@ def autoconnect(*blueprints: Blueprint) -> Blueprint: ) +def _fan_in( + atoms: list[BlueprintAtom], instance_key: str, port: str, funnel: TopicFunnel +) -> list[BlueprintAtom]: + """Point one declared input at several streams instead of one. + + The port stops being a stream of its own; each entry takes its place as an + `In` of the same type, so autoconnect, `.namespace()` and transport pins + treat the entries as ordinary streams. The module keeps the port as the + single place messages from all of them arrive. + """ + updated = [] + matched = False + for atom in atoms: + declared = next( + ( + stream + for stream in atom.streams + if atom.name == instance_key + and stream.name == port + and stream.direction in ("in", "inout") + ), + None, + ) + if declared is None: + updated.append(atom) + continue + matched = True + entries = tuple( + StreamRef(name=name, type=declared.type, direction="in") for name in funnel.names + ) + updated.append( + replace( + atom, + kwargs={ + **atom.kwargs, + "topic_funnels": {**atom.kwargs.get("topic_funnels", {}), port: funnel}, + }, + streams=tuple(s for s in atom.streams if s is not declared) + entries, + ) + ) + if not matched: + raise ValueError( + f"cannot fan {port!r} in: {instance_key} declares no such In/IO stream " + "in this blueprint (already fanned in, or a typo?)" + ) + return updated + + def _eliminate_duplicates(blueprints: list[BlueprintAtom]) -> list[BlueprintAtom]: # The duplicates are eliminated in reverse so that newer blueprints override older ones. seen = set() diff --git a/dimos/core/module.py b/dimos/core/module.py index 86c6624345..1f74c0e7b8 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import asyncio -from collections.abc import AsyncGenerator, Callable +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from dataclasses import dataclass, field from functools import partial import inspect @@ -31,7 +31,7 @@ get_type_hints, ) -from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic import BaseModel, Field, field_validator from reactivex.disposable import CompositeDisposable, Disposable from dimos.core.core import T, rpc @@ -104,35 +104,31 @@ def get_loop() -> tuple[asyncio.AbstractEventLoop, threading.Thread | None]: class TopicFunnel(BaseModel): - """Several same-typed streams that one port fans into a single handler. + """Several same-typed streams that one declared port fans into one handler. + + Built by `Blueprint.remappings()` from a list or dict value; blueprints do + not construct it directly. `names` are stream names — `left_cam`, `robot1/lidar` — never a backend - topic string. Each entry becomes a synthetic `In` stream, so the blueprint - machinery treats it like a declared port: autoconnect matches it against - producers, `.remappings()` and `.namespace()` rewrite it, and transport - pins apply. A python module subscribes the wired streams itself; a native - module hands their channels to its subprocess, which does the subscribing. + topic string. Each entry replaces the port as an `In` stream, so autoconnect + matches it against producers, `.namespace()` rewrites it, and transport pins + apply. A python module subscribes the wired streams itself; a native module + hands their channels to its subprocess, which does the subscribing. `names` may also be a dict attaching arbitrary per-stream info that the - handler receives on `Metadata.info`: - `TopicFunnel(names={"depth": {"meters_per_unit": 0.001}, "color": {}})`. + handler receives on `Metadata.info`. """ names: list[str] - msg_type: type[Any] | None = None info: dict[str, dict[str, Any]] = Field(default_factory=dict) - @model_validator(mode="before") @classmethod - def _split_names_dict(cls, data: Any) -> Any: - if isinstance(data, dict) and isinstance(data.get("names"), dict): - names_dict = data["names"] - data = { - **data, - "names": list(names_dict), - "info": {name: value for name, value in names_dict.items() if value}, - } - return data + def of(cls, entries: "Sequence[str] | Mapping[str, Mapping[str, Any]]") -> "TopicFunnel": + """Build one from a plain list of names, or a dict of name -> info.""" + if isinstance(entries, Mapping): + info = {name: dict(value) for name, value in entries.items() if value} + return cls(names=list(entries), info=info) + return cls(names=list(entries)) @field_validator("names") @classmethod @@ -145,13 +141,6 @@ def _reject_topic_strings(cls, names: list[str]) -> list[str]: ) return names - @model_validator(mode="after") - def _info_keys_must_be_names(self) -> "TopicFunnel": - unknown = [name for name in self.info if name not in self.names] - if unknown: - raise ValueError(f"topic funnel info keys are not in names: {unknown}") - return self - @dataclass(frozen=True) class Metadata: @@ -192,7 +181,8 @@ class ModuleConfig(BaseConfig): # once (see BlueprintAtom.instance_name). Changes the RPC topic prefix # from the class name to this name. instance_name: str | None = None - # Port name -> the group of streams that port's single handler receives. + # Declared In/IO port -> the streams that port's single handler receives. + # Set by `Blueprint.remappings()`, not by blueprint authors. topic_funnels: dict[str, TopicFunnel] = Field(default_factory=dict) g: GlobalConfig = global_config @@ -728,7 +718,8 @@ def _auto_bind_handlers(self) -> None: bindings: list[tuple[Any, Callable[[Any], Any]]] = [] for input_name, in_stream in {**self.inputs, **self.ios}.items(): declared_handler = getattr(self, f"handle_{input_name}", None) - if declared_handler is None: + # A funnelled port's handler is bound to the funnel's streams instead. + if declared_handler is None or input_name in self.config.topic_funnels: continue # Async @rpc wraps the coroutine fn in a sync dispatcher. Unwrap it # so we subscribe the raw coroutine fn instead of the wrapper (which @@ -763,27 +754,31 @@ async def with_meta( self.process_observable(in_stream.pure_observable(), handler) def _make_funnel_streams(self) -> "dict[str, In[Any]]": - """One synthetic `In` per topic-funnel entry, keyed by stream name. + """One `In` per topic-funnel entry, keyed by stream name. These are wired like declared ports (`set_transport` finds them, the blueprint machinery remaps/namespaces/pins them), but they are not attributes, so `inputs` and the native launch line never see them as - ports of their own. + ports of their own — the funnel's port stands in for all of them. """ streams: dict[str, In[Any]] = {} - declared = set(self.inputs) | set(self.outputs) | set(self.ios) - for port, group in self.config.topic_funnels.items(): - for name in group.names: - if name in declared: + for port, funnel in self.config.topic_funnels.items(): + declared = self.inputs.get(port) or self.ios.get(port) + if declared is None: + raise ValueError( + f"topic funnel port {port!r} is not an In or IO stream of {type(self).__name__}" + ) + for name in funnel.names: + if name != port and name in {**self.inputs, **self.outputs, **self.ios}: raise ValueError( f"topic funnel {port!r} entry {name!r} collides with a " f"declared stream of {type(self).__name__}" ) if name in streams: raise ValueError( - f"topic funnel {port!r} entry {name!r} appears in more than one group" + f"topic funnel {port!r} entry {name!r} appears in more than one funnel" ) - streams[name] = In(group.msg_type or Any, name, self) + streams[name] = In(declared.type, name, self) return streams def _bind_topic_funnels(self) -> None: @@ -833,7 +828,7 @@ async def invoke( # type: ignore[misc] stream = self._funnel_streams[name] if getattr(stream, "_transport", None) is None: # Not wired by a coordinator (standalone use): default transport. - stream.transport = make_transport(name, group.msg_type) + stream.transport = make_transport(name, stream.type) self.register_disposable(Disposable(stream.subscribe(partial(on_msg, index)))) def _make_async_dispatch( diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index 0620670e98..3af8a22efe 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -205,8 +205,8 @@ class NativeModule(Module): ``DIMOS_TRANSPORT`` env var. With ``stdin_config``, the topics, config, publisher QoS and session settings also arrive as one JSON line on stdin. - A port named in ``topic_funnels`` gets a list of channels instead of one, so - several same-typed streams reach a single native handler. + A port fanned in by ``.remappings()`` gets a list of channels instead of one, + so several same-typed streams reach a single native handler. The native process should parse whichever it uses and pub/sub on the given topics directly. On ``stop()``, the process receives SIGTERM. @@ -540,20 +540,16 @@ def _collect_topics(self) -> TopicsMap: channel = getattr(transport, "channel", None) if channel is not None: topics[name] = channel - for port, group in self.config.topic_funnels.items(): - if port in topics: - raise ValueError( - f"[{self._module_label}] topic funnel {port!r} collides with the " - "port of the same name declared as a stream" - ) + for port, funnel in self.config.topic_funnels.items(): entries: list[str | dict[str, Any]] = [] - for name in group.names: - transport = getattr(self._funnel_streams[name], "_transport", None) + for name in funnel.names: + stream = self._funnel_streams[name] + transport = getattr(stream, "_transport", None) channel = getattr(transport, "channel", None) if channel is None: # Unwired (standalone use): the channel the default transport lands on. - channel = channel_for(name, group.msg_type) - info = group.info.get(name) + channel = channel_for(name, stream.type) + info = funnel.info.get(name) entries.append(channel if info is None else {"topic": channel, "info": info}) topics[port] = entries return topics diff --git a/dimos/core/test_async_module_topic_funnels.py b/dimos/core/test_async_module_topic_funnels.py index aa10f68e27..5be9039984 100644 --- a/dimos/core/test_async_module_topic_funnels.py +++ b/dimos/core/test_async_module_topic_funnels.py @@ -17,14 +17,15 @@ import pytest from dimos.core.coordination.module_coordinator import ModuleCoordinator -from dimos.core.module import Metadata, Module, TopicFunnel -from dimos.core.stream import Out +from dimos.core.module import Metadata, Module +from dimos.core.stream import In, Out from dimos.core.transport_factory import make_transport class FanInModule(Module): """One handler for two same-typed streams; echoes which one it came from.""" + sensors: In[int] tagged: Out[int] named: Out[str] scaled: Out[float] @@ -37,8 +38,8 @@ async def handle_sensors(self, value: int, meta: Metadata) -> None: @pytest.fixture def start_fan_in_module(each_transport): - blueprint = FanInModule.blueprint( - topic_funnels={"sensors": TopicFunnel(names={"s0": {"scale": 0.5}, "s1": {}})} + blueprint = FanInModule.blueprint().remappings( + [(FanInModule, "sensors", {"s0": {"scale": 0.5}, "s1": {}})] ) coordinator = ModuleCoordinator.build(blueprint) yield @@ -83,9 +84,9 @@ def test_topic_funnel_tags_each_message_with_its_stream(start_fan_in_module, gro @pytest.fixture def start_remapped_fan_in_module(each_transport): - blueprint = FanInModule.blueprint( - topic_funnels={"sensors": TopicFunnel(names=["s0", "s1"])} - ).remappings([(FanInModule, "s0", "alt0")]) + blueprint = FanInModule.blueprint().remappings( + [(FanInModule, "sensors", ["s0", "s1"]), (FanInModule, "s0", "alt0")] + ) coordinator = ModuleCoordinator.build(blueprint) yield coordinator.stop() @@ -111,27 +112,44 @@ def test_topic_funnel_entries_follow_remappings(start_remapped_fan_in_module, ea def test_namespace_prefixes_topic_funnel_entries(): - blueprint = FanInModule.blueprint( - topic_funnels={"sensors": TopicFunnel(names=["s0", "s1"])} - ).namespace("bot") + blueprint = ( + FanInModule.blueprint() + .remappings([(FanInModule, "sensors", ["s0", "s1"])]) + .namespace("bot") + ) atom = blueprint.blueprints[0] assert blueprint.remapping_map[atom.name, "s0"] == "bot/s0" assert blueprint.remapping_map[atom.name, "s1"] == "bot/s1" -def test_a_group_entry_cannot_collide_with_a_declared_stream(): +def test_a_fanned_in_port_no_longer_connects_under_its_own_name(): + """The port stands in for its entries rather than being a stream of its own.""" + blueprint = FanInModule.blueprint().remappings([(FanInModule, "sensors", ["s0", "s1"])]) + names = {stream.name for stream in blueprint.blueprints[0].streams} + assert "sensors" not in names + assert {"s0", "s1"} <= names + + +def test_fanning_in_an_undeclared_port_is_an_error(): + with pytest.raises(ValueError, match="no such In/IO stream"): + FanInModule.blueprint().remappings([(FanInModule, "nope", ["s0"])]) + + +def test_a_funnel_entry_cannot_collide_with_a_declared_stream(): with pytest.raises(ValueError, match="collides"): - FanInModule(topic_funnels={"sensors": TopicFunnel(names=["tagged"])}) + FanInModule(topic_funnels={"sensors": {"names": ["tagged"]}}) -def test_funnel_info_keys_must_be_names(): - with pytest.raises(ValueError, match="not in names"): - TopicFunnel(names=["s0"], info={"s9": {"scale": 2.0}}) +def test_a_funnel_entry_rejects_a_backend_topic_string(): + """Entries are stream names, so a leading slash is a transport leaking in.""" + with pytest.raises(ValueError, match="not topics"): + FanInModule.blueprint().remappings([(FanInModule, "sensors", ["/s0"])]) class PlainFanInModule(Module): """A funnel handler that doesn't ask for metadata just gets the message.""" + sensors: In[int] echoed: Out[int] async def handle_sensors(self, value: int) -> None: @@ -139,8 +157,8 @@ async def handle_sensors(self, value: int) -> None: def test_a_funnel_handler_without_meta_gets_just_the_message(each_transport): - blueprint = PlainFanInModule.blueprint( - topic_funnels={"sensors": TopicFunnel(names=["s0", "s1"])} + blueprint = PlainFanInModule.blueprint().remappings( + [(PlainFanInModule, "sensors", ["s0", "s1"])] ) coordinator = ModuleCoordinator.build(blueprint) s1, echoed = make_transport("s1"), make_transport("echoed") diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index 49a4e8b31e..3ccc97aa0b 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -105,6 +105,12 @@ class StubIoModule(NativeModule): tf: IO[TFMessage] +class StubFunnelModule(NativeModule): + config: StubNativeConfig + cams: In[Imu] + cmd_vel: In[Twist] + + class StubConsumer(Module): pointcloud: In[PointCloud2] imu: In[Imu] @@ -212,11 +218,11 @@ def test_tf_topic_comes_from_the_declared_port_only() -> None: def test_a_topic_funnel_reaches_the_native_process_as_a_list(monkeypatch) -> None: - """A group's channels are resolved without the module declaring a stream each.""" + """The funnelled port carries every entry's channel instead of one of its own.""" monkeypatch.setattr(native_module_mod.global_config, "transport", "lcm") - module = StubNativeModule( + module = StubFunnelModule( executable=_ECHO, - topic_funnels={"cams": TopicFunnel(names=["cam0/imu", "cam1/imu"], msg_type=Imu)}, + topic_funnels={"cams": TopicFunnel(names=["cam0/imu", "cam1/imu"])}, ) try: topics = module._collect_topics() @@ -232,13 +238,9 @@ def test_a_topic_funnel_reaches_the_native_process_as_a_list(monkeypatch) -> Non def test_funnel_info_rides_the_launch_line_but_not_the_argv(monkeypatch) -> None: """Per-topic info becomes a {topic, info} entry on stdin; argv keeps only topics.""" monkeypatch.setattr(native_module_mod.global_config, "transport", "lcm") - module = StubNativeModule( + module = StubFunnelModule( executable=_ECHO, - topic_funnels={ - "cams": TopicFunnel( - names={"cam0/imu": {"rectified": True}, "cam1/imu": {}}, msg_type=Imu - ) - }, + topic_funnels={"cams": TopicFunnel.of({"cam0/imu": {"rectified": True}, "cam1/imu": {}})}, ) try: topics = module._collect_topics() @@ -256,9 +258,9 @@ def test_funnel_info_rides_the_launch_line_but_not_the_argv(monkeypatch) -> None def test_a_wired_topic_funnel_entry_uses_its_transport() -> None: """Remapping/pins arrive as set_transport on the entry, and the launch line follows.""" - module = StubNativeModule( + module = StubFunnelModule( executable=_ECHO, - topic_funnels={"cams": TopicFunnel(names=["cam0/imu"], msg_type=Imu)}, + topic_funnels={"cams": TopicFunnel(names=["cam0/imu"])}, ) transport = LCMTransport("/remapped/imu", Imu) try: @@ -270,33 +272,17 @@ def test_a_wired_topic_funnel_entry_uses_its_transport() -> None: transport.stop() -def test_a_topic_funnel_rejects_a_backend_topic_string() -> None: - """Names are stream names, so a leading slash is a transport leaking in.""" - with pytest.raises(ValidationError, match="not topics"): - TopicFunnel(names=["/cam0/imu"], msg_type=Imu) - - -def test_a_topic_funnel_cannot_shadow_a_declared_port() -> None: - module = StubNativeModule( - executable=_ECHO, - topic_funnels={"cmd_vel": TopicFunnel(names=["other"], msg_type=Twist)}, - ) - transport = LCMTransport("/cmd_vel", Twist) - try: - module.set_transport("cmd_vel", transport) - with pytest.raises(ValueError, match="collides"): - module._collect_topics() - finally: - module.stop() - with contextlib.suppress(Exception): - transport.stop() +def test_a_topic_funnel_needs_a_declared_port() -> None: + """The funnel replaces a port's wiring, so it has to have one to replace.""" + with pytest.raises(ValueError, match="not an In or IO stream"): + StubFunnelModule(executable=_ECHO, topic_funnels={"nope": TopicFunnel(names=["cam0/imu"])}) def test_a_topic_funnel_is_not_a_native_config_field() -> None: - """The group is wiring, so it belongs in `topics`, not the config struct.""" - module = StubNativeModule( + """The funnel is wiring, so it belongs in `topics`, not the config struct.""" + module = StubFunnelModule( executable=_ECHO, - topic_funnels={"cams": TopicFunnel(names=["cam0/imu"], msg_type=Imu)}, + topic_funnels={"cams": TopicFunnel(names=["cam0/imu"])}, ) try: assert "topic_funnels" not in module.config.to_config_dict() diff --git a/native/rust/README.md b/native/rust/README.md index 9c57f6445b..4f5d1c5b4f 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -127,30 +127,32 @@ impl MultiCam { } ``` -The Python wrapper supplies the sources with `topic_funnels`, keyed by port name: +Funnelling happens at the blueprint routing level: `.remappings()` takes a list of stream names where it would otherwise take one. ```python -MultiCam.blueprint( - topic_funnels={"cameras": TopicFunnel(names=["left_cam", "right_cam"], msg_type=Image)}, -) +MultiCam.blueprint().remappings([ + (MultiCam, "cameras", ["left_cam", "right_cam"]), +]) ``` -`names` may also be a dict attaching arbitrary per-stream info — `names={"left_cam": {"rectified": True}, "right_cam": {}}` — which reaches the handler as `Metadata.info` (a dict in python, `serde_json::Value` in rust; on the launch line such an entry is a `{"topic": ..., "info": ...}` object instead of a plain string). +A dict instead of a list attaches arbitrary per-stream info — `{"left_cam": {"rectified": True}, "right_cam": {}}` — which reaches the handler as `Metadata.info` (a dict in python, `serde_json::Value` in rust; on the launch line such an entry is a `{"topic": ..., "info": ...}` object instead of a plain string). -`names` are stream names, not backend topics — a leading `/` is rejected. Each entry becomes a synthetic `In` stream on the python side, so the blueprint machinery treats it like a declared port: autoconnect matches it against producers' `Out` streams, `.remappings()` and `.namespace()` rewrite it, and blueprint transport pins apply. The same blueprint runs unchanged over LCM or zenoh. +The entries are stream names, not backend topics — a leading `/` is rejected. Each one takes the port's place as an `In` stream of the same type, so autoconnect matches it against producers' `Out` streams, `.namespace()` and further `.remappings()` rewrite it, and blueprint transport pins apply. The same blueprint runs unchanged over LCM or zenoh. -The group itself is not a port with a stream of its own: python hands the wired entries' channels to the native process, which subscribes them directly. On the launch line the port's value is an array rather than a string. A group configured with no names still claims its port but never yields. +The port itself stops being a stream of its own: python hands the wired entries' channels to the native process, which subscribes them directly. On the launch line the port's value is an array rather than a string. A funnel with no entries still claims its port but never yields. -`topic_funnels` lives on `ModuleConfig`, so a plain Python `Module` takes the same field. There the module subscribes the group itself and dispatches to `async def handle_`. Opting into metadata is by signature — a one-parameter handler just gets the message, a two-parameter handler also gets a `Metadata` with `index` (position in `names`) and `name` (the stream name as declared, pre-remapping): +Fan-in is a `ModuleConfig` feature, so a plain Python `Module` funnels the same way. There the module subscribes the entries itself and dispatches to `async def handle_`. Opting into metadata is by signature — a one-parameter handler just gets the message, a two-parameter handler also gets a `Metadata` with `index` (position in the list) and `name` (the stream name as declared, pre-remapping): ```python class MultiCam(Module): + cameras: In[Image] + async def handle_cameras(self, image: Image, meta: Metadata) -> None: ... ``` The same applies to any `handle_` for a plain `In` port, where the metadata is always `index=0, name=`. -The whole group shares one dispatcher, so the handler is never re-entered, and its mailbox holds the latest unprocessed message per stream rather than one slot for the group — a chatty camera cannot starve the others. +The whole funnel shares one dispatcher, so the handler is never re-entered, and its mailbox holds the latest unprocessed message per stream rather than one slot for the funnel — a chatty camera cannot starve the others. ## Transforms From 58bc9d962f559691f1e173b4335a03be5e5a1b6a Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 31 Aug 2026 16:41:43 -0700 Subject: [PATCH 09/16] Give funnel launch-line entries a real type The topics map was typed loosely enough that the baked host had to widen its override to Any. Name the info-carrying entry after the native TopicEntry it serializes to, and give the port's value its own alias so the host can state what it passes through. The remaining Any is the info payload itself, which is arbitrary by design. --- dimos/core/baked_host.py | 8 ++++---- dimos/core/native_module.py | 21 +++++++++++++++------ 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/dimos/core/baked_host.py b/dimos/core/baked_host.py index afa717e9da..e4ef793040 100644 --- a/dimos/core/baked_host.py +++ b/dimos/core/baked_host.py @@ -20,11 +20,11 @@ from collections.abc import Mapping import json import sys -from typing import Any, get_args, get_origin, get_type_hints +from typing import get_args, get_origin, get_type_hints from pydantic import Field, create_model -from dimos.core.native_module import NativeModule, NativeModuleConfig, TopicsMap +from dimos.core.native_module import NativeModule, NativeModuleConfig, TopicsMap, TopicValue from dimos.core.stream import IO, In, Out @@ -83,9 +83,9 @@ class BakedHost(NativeModule): _members: Mapping[str, type[NativeModule]] = {} _remaps: Mapping[tuple[str, str], str] = {} - def _member_topics(self, instance: str, topics: Mapping[str, Any]) -> dict[str, Any]: + def _member_topics(self, instance: str, topics: TopicsMap) -> dict[str, TopicValue]: member = self._members[instance] - resolved = {} + resolved: dict[str, TopicValue] = {} for port in _member_ports(member): name = self._remaps.get((instance, port), port) if name in topics: diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index 3af8a22efe..64d70f6caa 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -52,7 +52,7 @@ class MyCppModule(NativeModule): import sys import threading import time -from typing import IO, Any +from typing import IO, Any, TypedDict from pydantic import Field, model_validator @@ -64,9 +64,18 @@ class MyCppModule(NativeModule): from dimos.protocol.service.spec import SessionConfig from dimos.utils.logging_config import setup_logger -# Port name -> the wire channel(s) it is launched with. A funnel port carries a -# list, and an entry with per-topic info is an object rather than a bare string. -TopicsMap = dict[str, str | list[str | dict[str, Any]]] + +class TopicEntry(TypedDict): + """A funnel entry carrying that stream's info, mirroring the native `TopicEntry`.""" + + topic: str + info: dict[str, Any] + + +# The wire channel(s) one port is launched with. A funnel port carries a list, +# and an entry with per-topic info is an object rather than a bare string. +TopicValue = str | list[str | TopicEntry] +TopicsMap = dict[str, TopicValue] if sys.platform.startswith("linux"): import ctypes @@ -541,7 +550,7 @@ def _collect_topics(self) -> TopicsMap: if channel is not None: topics[name] = channel for port, funnel in self.config.topic_funnels.items(): - entries: list[str | dict[str, Any]] = [] + entries: list[str | TopicEntry] = [] for name in funnel.names: stream = self._funnel_streams[name] transport = getattr(stream, "_transport", None) @@ -550,7 +559,7 @@ def _collect_topics(self) -> TopicsMap: # Unwired (standalone use): the channel the default transport lands on. channel = channel_for(name, stream.type) info = funnel.info.get(name) - entries.append(channel if info is None else {"topic": channel, "info": info}) + entries.append(channel if info is None else TopicEntry(topic=channel, info=info)) topics[port] = entries return topics From bc8bc32ee2d3c38099ee0b67421b9b4b8431c1e0 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 31 Aug 2026 16:51:00 -0700 Subject: [PATCH 10/16] Rename Metadata to TopicMetadata Metadata is a word this codebase already uses for several unrelated things, so the bare name says nothing about what the handler is being handed. Name it for what it describes: which topic the message arrived on. --- dimos/core/module.py | 16 ++++++++-------- dimos/core/test_async_module_topic_funnels.py | 4 ++-- native/rust/README.md | 12 ++++++------ native/rust/dimos-module/src/lib.rs | 2 +- native/rust/dimos-module/src/module.rs | 8 ++++---- 5 files changed, 21 insertions(+), 21 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index 1f74c0e7b8..258e966b57 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -116,7 +116,7 @@ class TopicFunnel(BaseModel): hands their channels to its subprocess, which does the subscribing. `names` may also be a dict attaching arbitrary per-stream info that the - handler receives on `Metadata.info`. + handler receives on `TopicMetadata.info`. """ names: list[str] @@ -143,7 +143,7 @@ def _reject_topic_strings(cls, names: list[str]) -> list[str]: @dataclass(frozen=True) -class Metadata: +class TopicMetadata: """Which stream a message arrived on, for handlers that take `(msg, meta)`. `index` is the position in the funnel's `names` (0 for a plain input); @@ -733,12 +733,12 @@ def _auto_bind_handlers(self) -> None: ) handler: Callable[..., Any] = declared_handler if _handler_wants_metadata(handler, f"{type(self).__name__}.handle_{input_name}"): - metadata = Metadata(index=0, name=input_name) + metadata = TopicMetadata(index=0, name=input_name) async def with_meta( msg: Any, - _handler: Callable[[Any, Metadata], Any] = handler, - _metadata: Metadata = metadata, + _handler: Callable[[Any, TopicMetadata], Any] = handler, + _metadata: TopicMetadata = metadata, ) -> None: await _handler(msg, _metadata) @@ -802,15 +802,15 @@ def _bind_topic_funnels(self) -> None: handler: Callable[..., Any] = declared_handler if _handler_wants_metadata(handler, f"{type(self).__name__}.handle_{port}"): metas = tuple( - Metadata(index=index, name=name, info=group.info.get(name, {})) + TopicMetadata(index=index, name=name, info=group.info.get(name, {})) for index, name in enumerate(group.names) ) async def invoke( index: int, msg: Any, - _handler: Callable[[Any, Metadata], Any] = handler, - _metas: tuple[Metadata, ...] = metas, + _handler: Callable[[Any, TopicMetadata], Any] = handler, + _metas: tuple[TopicMetadata, ...] = metas, ) -> None: await _handler(msg, _metas[index]) else: diff --git a/dimos/core/test_async_module_topic_funnels.py b/dimos/core/test_async_module_topic_funnels.py index 5be9039984..3fab39a1cf 100644 --- a/dimos/core/test_async_module_topic_funnels.py +++ b/dimos/core/test_async_module_topic_funnels.py @@ -17,7 +17,7 @@ import pytest from dimos.core.coordination.module_coordinator import ModuleCoordinator -from dimos.core.module import Metadata, Module +from dimos.core.module import Module, TopicMetadata from dimos.core.stream import In, Out from dimos.core.transport_factory import make_transport @@ -30,7 +30,7 @@ class FanInModule(Module): named: Out[str] scaled: Out[float] - async def handle_sensors(self, value: int, meta: Metadata) -> None: + async def handle_sensors(self, value: int, meta: TopicMetadata) -> None: self.tagged.publish(meta.index * 100 + value) self.named.publish(meta.name) self.scaled.publish(value * meta.info.get("scale", 1.0)) diff --git a/native/rust/README.md b/native/rust/README.md index 4f5d1c5b4f..5a9dfddc6c 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -61,7 +61,7 @@ Every transport is compiled into the binary. `run_with_transport` opens the one - `#[derive(Module)]`: on the struct. Required. - `#[module(setup = fn, teardown = fn)]`: on the struct. Both optional. Names methods on `Self`. `setup` runs once before the input dispatch loop starts (use it to spawn background tasks or initialize resources); `teardown` runs once after the loop exits (use it for cleanup). -- `#[input(decode = fn, handler = fn, meta)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. The `meta` flag makes the handler take `(msg, meta: Metadata)`, where `meta` names the topic the message arrived on — useful when the port is a topic funnel (see [Topic funnels](#topic-funnels)). +- `#[input(decode = fn, handler = fn, meta)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. The `meta` flag makes the handler take `(msg, meta: TopicMetadata)`, where `meta` names the topic the message arrived on — useful when the port is a topic funnel (see [Topic funnels](#topic-funnels)). - `#[output(encode = fn)]`: on a field of type `Output`. `encode` is required. - `#[io(decode = fn, encode = fn, handler = fn)]`: on a field of type `Io`, a port that publishes to and subscribes on one topic. `decode` and `encode` are required; `handler` defaults to `handle_`. The transports deliver a message back to its own sender, so the handler also sees what the module publishes. Use `#[output]` instead when the module only publishes. - `#[config]`: on one field. The type must be defined with `#[native_config]` (see [Config](#config)). At most one per struct. If absent, `Config` defaults to `dimos_module::NoConfig`. @@ -110,7 +110,7 @@ Field name = port name. Ports map to topics via the stdin JSON; unmapped ports f ## Topic funnels -A rig with N identical sensors would otherwise need N ports and N near-identical handlers. A topic funnel is one `Input` port wired to a list of topics that all carry `T`, delivered to one handler in arrival order. Any input accepts a funnel — on the launch line the port's value is an array of topics rather than a string, and nothing in the module changes. A handler that needs to tell the sources apart opts in with the `meta` flag and receives a `Metadata` alongside each message: +A rig with N identical sensors would otherwise need N ports and N near-identical handlers. A topic funnel is one `Input` port wired to a list of topics that all carry `T`, delivered to one handler in arrival order. Any input accepts a funnel — on the launch line the port's value is an array of topics rather than a string, and nothing in the module changes. A handler that needs to tell the sources apart opts in with the `meta` flag and receives a `TopicMetadata` alongside each message: ```rust #[derive(Module)] @@ -120,7 +120,7 @@ struct MultiCam { } impl MultiCam { - async fn handle_cameras(&mut self, image: Image, meta: Metadata) { + async fn handle_cameras(&mut self, image: Image, meta: TopicMetadata) { // meta.index: position in the topic list; meta.topic: the topic itself; // meta.info: per-topic JSON the coordinator attached (Null when none) } @@ -135,19 +135,19 @@ MultiCam.blueprint().remappings([ ]) ``` -A dict instead of a list attaches arbitrary per-stream info — `{"left_cam": {"rectified": True}, "right_cam": {}}` — which reaches the handler as `Metadata.info` (a dict in python, `serde_json::Value` in rust; on the launch line such an entry is a `{"topic": ..., "info": ...}` object instead of a plain string). +A dict instead of a list attaches arbitrary per-stream info — `{"left_cam": {"rectified": True}, "right_cam": {}}` — which reaches the handler as `TopicMetadata.info` (a dict in python, `serde_json::Value` in rust; on the launch line such an entry is a `{"topic": ..., "info": ...}` object instead of a plain string). The entries are stream names, not backend topics — a leading `/` is rejected. Each one takes the port's place as an `In` stream of the same type, so autoconnect matches it against producers' `Out` streams, `.namespace()` and further `.remappings()` rewrite it, and blueprint transport pins apply. The same blueprint runs unchanged over LCM or zenoh. The port itself stops being a stream of its own: python hands the wired entries' channels to the native process, which subscribes them directly. On the launch line the port's value is an array rather than a string. A funnel with no entries still claims its port but never yields. -Fan-in is a `ModuleConfig` feature, so a plain Python `Module` funnels the same way. There the module subscribes the entries itself and dispatches to `async def handle_`. Opting into metadata is by signature — a one-parameter handler just gets the message, a two-parameter handler also gets a `Metadata` with `index` (position in the list) and `name` (the stream name as declared, pre-remapping): +Fan-in is a `ModuleConfig` feature, so a plain Python `Module` funnels the same way. There the module subscribes the entries itself and dispatches to `async def handle_`. Opting into metadata is by signature — a one-parameter handler just gets the message, a two-parameter handler also gets a `TopicMetadata` with `index` (position in the list) and `name` (the stream name as declared, pre-remapping): ```python class MultiCam(Module): cameras: In[Image] - async def handle_cameras(self, image: Image, meta: Metadata) -> None: ... + async def handle_cameras(self, image: Image, meta: TopicMetadata) -> None: ... ``` The same applies to any `handle_` for a plain `In` port, where the metadata is always `index=0, name=`. diff --git a/native/rust/dimos-module/src/lib.rs b/native/rust/dimos-module/src/lib.rs index 1390402dd7..7307b1145e 100644 --- a/native/rust/dimos-module/src/lib.rs +++ b/native/rust/dimos-module/src/lib.rs @@ -30,7 +30,7 @@ pub use dimos_module_macros::{native_config, Module}; pub use host::{host_main, HostSpec, ModuleEntry}; pub use lcm::LcmTransport; pub use module::{ - run, Builder, Input, Io, Metadata, Module, ModuleConfig, NativeConfig, NoConfig, Output, + run, Builder, Input, Io, Module, ModuleConfig, NativeConfig, NoConfig, Output, TopicMetadata, }; pub use tf::{Lookup, Tf, Transform}; pub use transport::{SharedTransport, Transport}; diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index 45fbb1d3dd..6e5070d391 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -117,7 +117,7 @@ impl Route for TypedRoute { /// with `#[input(meta)]`. `info` is whatever the coordinator attached to that /// topic on the launch line (`Null` when nothing was). #[derive(Clone, Debug)] -pub struct Metadata { +pub struct TopicMetadata { pub index: usize, pub topic: String, pub info: serde_json::Value, @@ -138,9 +138,9 @@ impl Input { self.receiver.recv().await.map(|(_, msg)| msg) } - pub async fn recv_meta(&mut self) -> Option<(T, Metadata)> { + pub async fn recv_meta(&mut self) -> Option<(T, TopicMetadata)> { let (index, msg) = self.receiver.recv().await?; - let meta = Metadata { + let meta = TopicMetadata { index, topic: self.topics[index].clone(), info: self.infos[index].clone(), @@ -1545,7 +1545,7 @@ mod tests { } impl Rig { - async fn handle_cams(&mut self, msg: Msg, meta: crate::Metadata) { + async fn handle_cams(&mut self, msg: Msg, meta: crate::TopicMetadata) { let mut tagged = vec![meta.index as u8]; tagged.extend(msg.0); self.seen.publish(&Msg(tagged)).await.expect("publish"); From 00efbc0e0d5c10e4204f30f07358698c56306906 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 31 Aug 2026 17:01:03 -0700 Subject: [PATCH 11/16] Drop the funnel tests that assert plumbing rather than behavior Namespacing walks a blueprint's streams uniformly, so prefixing a funnel entry was already covered by the generic namespace tests once the entries became streams. The port's absence from that same list and the leading-slash rejection were each one assertion about an internal structure, not about anything a module does. --- dimos/core/test_async_module_topic_funnels.py | 25 ------------------- 1 file changed, 25 deletions(-) diff --git a/dimos/core/test_async_module_topic_funnels.py b/dimos/core/test_async_module_topic_funnels.py index 3fab39a1cf..4b2fe084a5 100644 --- a/dimos/core/test_async_module_topic_funnels.py +++ b/dimos/core/test_async_module_topic_funnels.py @@ -111,25 +111,6 @@ def test_topic_funnel_entries_follow_remappings(start_remapped_fan_in_module, ea transport.stop() -def test_namespace_prefixes_topic_funnel_entries(): - blueprint = ( - FanInModule.blueprint() - .remappings([(FanInModule, "sensors", ["s0", "s1"])]) - .namespace("bot") - ) - atom = blueprint.blueprints[0] - assert blueprint.remapping_map[atom.name, "s0"] == "bot/s0" - assert blueprint.remapping_map[atom.name, "s1"] == "bot/s1" - - -def test_a_fanned_in_port_no_longer_connects_under_its_own_name(): - """The port stands in for its entries rather than being a stream of its own.""" - blueprint = FanInModule.blueprint().remappings([(FanInModule, "sensors", ["s0", "s1"])]) - names = {stream.name for stream in blueprint.blueprints[0].streams} - assert "sensors" not in names - assert {"s0", "s1"} <= names - - def test_fanning_in_an_undeclared_port_is_an_error(): with pytest.raises(ValueError, match="no such In/IO stream"): FanInModule.blueprint().remappings([(FanInModule, "nope", ["s0"])]) @@ -140,12 +121,6 @@ def test_a_funnel_entry_cannot_collide_with_a_declared_stream(): FanInModule(topic_funnels={"sensors": {"names": ["tagged"]}}) -def test_a_funnel_entry_rejects_a_backend_topic_string(): - """Entries are stream names, so a leading slash is a transport leaking in.""" - with pytest.raises(ValueError, match="not topics"): - FanInModule.blueprint().remappings([(FanInModule, "sensors", ["/s0"])]) - - class PlainFanInModule(Module): """A funnel handler that doesn't ask for metadata just gets the message.""" From ebf03db7f7d1e07ad7f74a705ac6043cb4cacfa5 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 31 Aug 2026 17:08:14 -0700 Subject: [PATCH 12/16] Simplify the handler metadata check Filtering the signature down to positional parameters made a keyword-only or variadic handler quietly count as one-argument. A handler takes the message and optionally the metadata, so count what it declares and reject anything else. --- dimos/core/module.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index 258e966b57..e23579fa79 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -158,17 +158,10 @@ class TopicMetadata: def _handler_wants_metadata(handler: Callable[..., Any], label: str) -> bool: """True if the bound handler takes `(msg, meta)` rather than just `(msg)`.""" - positional_kinds = (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) - count = sum( - 1 - for parameter in inspect.signature(handler).parameters.values() - if parameter.kind in positional_kinds - ) - if count == 1: - return False - if count == 2: - return True - raise TypeError(f"{label} must take (msg) or (msg, meta), not {count} positional parameters") + count = len(inspect.signature(handler).parameters) + if count not in (1, 2): + raise TypeError(f"{label} must take (msg) or (msg, meta), not {count} parameters") + return count == 2 class ModuleConfig(BaseConfig): From 0eee80896ff1c69d7904f1eca78c2b6665998bb2 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 31 Aug 2026 17:19:03 -0700 Subject: [PATCH 13/16] Trim the topic funnel docstrings Review feedback: remappings needs no prose, TopicFunnel is better shown than described, and _argv's existing one-liner was fine as it was. --- dimos/core/coordination/blueprints.py | 11 ----------- dimos/core/module.py | 15 +++++---------- dimos/core/native_module.py | 4 +--- 3 files changed, 6 insertions(+), 24 deletions(-) diff --git a/dimos/core/coordination/blueprints.py b/dimos/core/coordination/blueprints.py index 42a4f1d887..4062980826 100644 --- a/dimos/core/coordination/blueprints.py +++ b/dimos/core/coordination/blueprints.py @@ -242,17 +242,6 @@ def remappings( ] ], ) -> "Blueprint": - """Rewrite what a module's streams and module refs connect to. - - A string (or Spec/Module class) renames one stream (or redirects one ref). - A list of names, or a dict of name -> arbitrary info, instead fans several - same-typed streams into that one port, so its handler serves all of them: - - .remappings([ - (Fuser, "scan", ["front_lidar", "rear_lidar"]), - (Fuser, "image", {"depth": {"meters_per_unit": 0.001}, "color": {}}), - ]) - """ remappings_dict = dict(self.remapping_map) atoms = list(self.blueprints) for module, old, new in remappings: diff --git a/dimos/core/module.py b/dimos/core/module.py index e23579fa79..943ac394f9 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -106,17 +106,12 @@ def get_loop() -> tuple[asyncio.AbstractEventLoop, threading.Thread | None]: class TopicFunnel(BaseModel): """Several same-typed streams that one declared port fans into one handler. - Built by `Blueprint.remappings()` from a list or dict value; blueprints do - not construct it directly. + Built by `Blueprint.remappings()`, not by blueprint authors:: - `names` are stream names — `left_cam`, `robot1/lidar` — never a backend - topic string. Each entry replaces the port as an `In` stream, so autoconnect - matches it against producers, `.namespace()` rewrites it, and transport pins - apply. A python module subscribes the wired streams itself; a native module - hands their channels to its subprocess, which does the subscribing. - - `names` may also be a dict attaching arbitrary per-stream info that the - handler receives on `TopicMetadata.info`. + .remappings([ + (Fuser, "scan", ["front_lidar", "rear_lidar"]), + (Fuser, "image", {"depth": {"meters_per_unit": 0.001}, "color": {}}), + ]) """ names: list[str] diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index 64d70f6caa..e538a9bb32 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -275,9 +275,7 @@ def _session(self) -> SessionConfig: return pinned.rebased() def _argv(self, topics: TopicsMap) -> list[str]: - """The command line the native process is spawned with. Per-topic funnel - info travels only on the stdin JSON, so an entry object contributes just - its topic here.""" + """The command line the native process is spawned with.""" cmd = [self.config.executable] for name, value in topics.items(): if isinstance(value, list): From 4320f899893081d1cf44db61681ef5bd4f62bfcb Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 31 Aug 2026 17:53:09 -0700 Subject: [PATCH 14/16] Construct a TopicFunnel directly instead of via .of A pydantic model is keyword-only, so building one meant naming both stored fields at every call site even though authoring takes a single list-or-dict. A plain dataclass takes that one argument positionally and splits it itself. Also name the funnel tests' delivery timeout rather than repeating the literal. --- dimos/core/coordination/blueprints.py | 2 +- dimos/core/module.py | 30 ++++++++----------- dimos/core/test_async_module_topic_funnels.py | 24 ++++++++------- dimos/core/test_native_module.py | 10 +++---- 4 files changed, 32 insertions(+), 34 deletions(-) diff --git a/dimos/core/coordination/blueprints.py b/dimos/core/coordination/blueprints.py index 4062980826..945819f5db 100644 --- a/dimos/core/coordination/blueprints.py +++ b/dimos/core/coordination/blueprints.py @@ -249,7 +249,7 @@ def remappings( if isinstance(new, (str, type)): remappings_dict[instance_key, old] = new else: - atoms = _fan_in(atoms, instance_key, old, TopicFunnel.of(new)) + atoms = _fan_in(atoms, instance_key, old, TopicFunnel(new)) return replace( self, blueprints=tuple(atoms), remapping_map=MappingProxyType(remappings_dict) ) diff --git a/dimos/core/module.py b/dimos/core/module.py index 943ac394f9..d59a0411ac 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -31,7 +31,7 @@ get_type_hints, ) -from pydantic import BaseModel, Field, field_validator +from pydantic import Field from reactivex.disposable import CompositeDisposable, Disposable from dimos.core.core import T, rpc @@ -103,7 +103,8 @@ def get_loop() -> tuple[asyncio.AbstractEventLoop, threading.Thread | None]: Deployment = Literal["python", "docker"] -class TopicFunnel(BaseModel): +@dataclass(init=False) +class TopicFunnel: """Several same-typed streams that one declared port fans into one handler. Built by `Blueprint.remappings()`, not by blueprint authors:: @@ -115,26 +116,21 @@ class TopicFunnel(BaseModel): """ names: list[str] - info: dict[str, dict[str, Any]] = Field(default_factory=dict) - - @classmethod - def of(cls, entries: "Sequence[str] | Mapping[str, Mapping[str, Any]]") -> "TopicFunnel": - """Build one from a plain list of names, or a dict of name -> info.""" - if isinstance(entries, Mapping): - info = {name: dict(value) for name, value in entries.items() if value} - return cls(names=list(entries), info=info) - return cls(names=list(entries)) - - @field_validator("names") - @classmethod - def _reject_topic_strings(cls, names: list[str]) -> list[str]: - leading_slash = [n for n in names if n.startswith("/")] + info: dict[str, dict[str, Any]] + + def __init__(self, entries: "Sequence[str] | Mapping[str, Mapping[str, Any]]") -> None: + self.names = list(entries) + self.info = ( + {name: dict(value) for name, value in entries.items() if value} + if isinstance(entries, Mapping) + else {} + ) + leading_slash = [name for name in self.names if name.startswith("/")] if leading_slash: raise ValueError( f"topic funnel names must be stream names, not topics: {leading_slash} " "start with '/' (use `left_cam`, not `/left_cam`)" ) - return names @dataclass(frozen=True) diff --git a/dimos/core/test_async_module_topic_funnels.py b/dimos/core/test_async_module_topic_funnels.py index 4b2fe084a5..a0fa520dcb 100644 --- a/dimos/core/test_async_module_topic_funnels.py +++ b/dimos/core/test_async_module_topic_funnels.py @@ -17,10 +17,12 @@ import pytest from dimos.core.coordination.module_coordinator import ModuleCoordinator -from dimos.core.module import Module, TopicMetadata +from dimos.core.module import Module, TopicFunnel, TopicMetadata from dimos.core.stream import In, Out from dimos.core.transport_factory import make_transport +DELIVERY_TIMEOUT = 1.0 + class FanInModule(Module): """One handler for two same-typed streams; echoes which one it came from.""" @@ -72,14 +74,14 @@ def test_topic_funnel_tags_each_message_with_its_stream(start_fan_in_module, gro scaled.subscribe(scales.put) s0.publish(7) - assert queue.get(timeout=1.0) == 7 - assert names.get(timeout=1.0) == "s0" - assert scales.get(timeout=1.0) == 3.5 + assert queue.get(timeout=DELIVERY_TIMEOUT) == 7 + assert names.get(timeout=DELIVERY_TIMEOUT) == "s0" + assert scales.get(timeout=DELIVERY_TIMEOUT) == 3.5 s1.publish(7) - assert queue.get(timeout=1.0) == 107 - assert names.get(timeout=1.0) == "s1" - assert scales.get(timeout=1.0) == 7.0 + assert queue.get(timeout=DELIVERY_TIMEOUT) == 107 + assert names.get(timeout=DELIVERY_TIMEOUT) == "s1" + assert scales.get(timeout=DELIVERY_TIMEOUT) == 7.0 @pytest.fixture @@ -104,8 +106,8 @@ def test_topic_funnel_entries_follow_remappings(start_remapped_fan_in_module, ea tagged.subscribe(queue.put) named.subscribe(names.put) alt0.publish(7) - assert queue.get(timeout=1.0) == 7 - assert names.get(timeout=1.0) == "s0" + assert queue.get(timeout=DELIVERY_TIMEOUT) == 7 + assert names.get(timeout=DELIVERY_TIMEOUT) == "s0" finally: for transport in (alt0, tagged, named): transport.stop() @@ -118,7 +120,7 @@ def test_fanning_in_an_undeclared_port_is_an_error(): def test_a_funnel_entry_cannot_collide_with_a_declared_stream(): with pytest.raises(ValueError, match="collides"): - FanInModule(topic_funnels={"sensors": {"names": ["tagged"]}}) + FanInModule(topic_funnels={"sensors": TopicFunnel(["tagged"])}) class PlainFanInModule(Module): @@ -143,7 +145,7 @@ def test_a_funnel_handler_without_meta_gets_just_the_message(each_transport): queue: Queue[int] = Queue() echoed.subscribe(queue.put) s1.publish(7) - assert queue.get(timeout=1.0) == 7 + assert queue.get(timeout=DELIVERY_TIMEOUT) == 7 finally: for transport in (s1, echoed): transport.stop() diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index 3ccc97aa0b..9490140dea 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -222,7 +222,7 @@ def test_a_topic_funnel_reaches_the_native_process_as_a_list(monkeypatch) -> Non monkeypatch.setattr(native_module_mod.global_config, "transport", "lcm") module = StubFunnelModule( executable=_ECHO, - topic_funnels={"cams": TopicFunnel(names=["cam0/imu", "cam1/imu"])}, + topic_funnels={"cams": TopicFunnel(["cam0/imu", "cam1/imu"])}, ) try: topics = module._collect_topics() @@ -240,7 +240,7 @@ def test_funnel_info_rides_the_launch_line_but_not_the_argv(monkeypatch) -> None monkeypatch.setattr(native_module_mod.global_config, "transport", "lcm") module = StubFunnelModule( executable=_ECHO, - topic_funnels={"cams": TopicFunnel.of({"cam0/imu": {"rectified": True}, "cam1/imu": {}})}, + topic_funnels={"cams": TopicFunnel({"cam0/imu": {"rectified": True}, "cam1/imu": {}})}, ) try: topics = module._collect_topics() @@ -260,7 +260,7 @@ def test_a_wired_topic_funnel_entry_uses_its_transport() -> None: """Remapping/pins arrive as set_transport on the entry, and the launch line follows.""" module = StubFunnelModule( executable=_ECHO, - topic_funnels={"cams": TopicFunnel(names=["cam0/imu"])}, + topic_funnels={"cams": TopicFunnel(["cam0/imu"])}, ) transport = LCMTransport("/remapped/imu", Imu) try: @@ -275,14 +275,14 @@ def test_a_wired_topic_funnel_entry_uses_its_transport() -> None: def test_a_topic_funnel_needs_a_declared_port() -> None: """The funnel replaces a port's wiring, so it has to have one to replace.""" with pytest.raises(ValueError, match="not an In or IO stream"): - StubFunnelModule(executable=_ECHO, topic_funnels={"nope": TopicFunnel(names=["cam0/imu"])}) + StubFunnelModule(executable=_ECHO, topic_funnels={"nope": TopicFunnel(["cam0/imu"])}) def test_a_topic_funnel_is_not_a_native_config_field() -> None: """The funnel is wiring, so it belongs in `topics`, not the config struct.""" module = StubFunnelModule( executable=_ECHO, - topic_funnels={"cams": TopicFunnel(names=["cam0/imu"])}, + topic_funnels={"cams": TopicFunnel(["cam0/imu"])}, ) try: assert "topic_funnels" not in module.config.to_config_dict() From 306fd053a9069552dc3a72a4bf9ed71dc0621b11 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 31 Aug 2026 17:55:39 -0700 Subject: [PATCH 15/16] Say what each language's funnel dispatcher does under load --- native/rust/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/native/rust/README.md b/native/rust/README.md index 5a9dfddc6c..efc9004b88 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -152,7 +152,7 @@ class MultiCam(Module): The same applies to any `handle_` for a plain `In` port, where the metadata is always `index=0, name=`. -The whole funnel shares one dispatcher, so the handler is never re-entered, and its mailbox holds the latest unprocessed message per stream rather than one slot for the funnel — a chatty camera cannot starve the others. +Either way the whole funnel shares one dispatcher, so the handler is never re-entered. What they do under load differs: python's mailbox holds the latest unprocessed message per stream and drains them oldest-waiting first, so a chatty camera cannot starve the others, while rust feeds every topic of a funnel into one bounded channel that drops the arriving message once it is full. ## Transforms From d89d7fa6dca7e592e895efc326e2976b0af1b4dd Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 31 Aug 2026 18:25:35 -0700 Subject: [PATCH 16/16] Bind plain ports and funnel ports on one path A plain port is a funnel of one, so the two bind routines were the same thirty lines twice. Treating a declared port as a single-stream funnel folds them together and drops the second one, along with the extra rx layer the funnel path was going through to reach the same dispatcher. Funnel streams now get their default transport when they are built, which is what makes the merged path safe standalone; a coordinator overwrites it as before. That also leaves channel_for with no callers. --- dimos/core/module.py | 149 +++++++++++---------------- dimos/core/native_module.py | 12 +-- dimos/core/test_transport_factory.py | 9 -- dimos/core/transport_factory.py | 17 --- 4 files changed, 66 insertions(+), 121 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index d59a0411ac..88d8a1666e 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -155,6 +155,24 @@ def _handler_wants_metadata(handler: Callable[..., Any], label: str) -> bool: return count == 2 +def _dispatch_to( + handler: Callable[..., Any], metas: "list[TopicMetadata] | None" +) -> Callable[[int, Any], Any]: + """Adapt a handler to the `(index, msg)` the dispatcher calls with. + + `metas` is None when the handler takes just the message; otherwise it holds + one `TopicMetadata` per stream feeding the port, indexed the same way. + """ + + async def invoke(index: int, msg: Any) -> None: + if metas is None: + await handler(msg) + else: + await handler(msg, metas[index]) + + return invoke + + class ModuleConfig(BaseConfig): rpc_transport: type[RPCSpec] = Field(default_factory=rpc_backend) default_rpc_timeout: float = DEFAULT_RPC_TIMEOUT @@ -247,7 +265,6 @@ def build(self) -> None: def start(self) -> None: self._start_main() self._auto_bind_handlers() - self._bind_topic_funnels() @rpc def stop(self) -> None: @@ -696,46 +713,52 @@ def _stop_main(self) -> None: def _auto_bind_handlers(self) -> None: """ For each declared `x: In[T]` or `x: IO[T]`, if `async def handle_x` exists, - subscribe it via process_observable so it runs on self._loop. + subscribe it to that port's stream so it runs on self._loop. A port that + `.remappings()` fanned in is subscribed to every stream of its funnel + instead, all reaching the one handler. + + A native module defines no such method — its subprocess subscribes + instead — so this is a no-op there. """ # Validate every handler before subscribing any of them. - bindings: list[tuple[Any, Callable[[Any], Any]]] = [] - for input_name, in_stream in {**self.inputs, **self.ios}.items(): - declared_handler = getattr(self, f"handle_{input_name}", None) - # A funnelled port's handler is bound to the funnel's streams instead. - if declared_handler is None or input_name in self.config.topic_funnels: + bindings: list[tuple[list[In[Any] | IO[Any]], Callable[[int, Any], Any]]] = [] + ports: dict[str, In[Any] | IO[Any]] = {**self.inputs, **self.ios} + for port, declared in ports.items(): + handler = getattr(self, f"handle_{port}", None) + if handler is None: continue # Async @rpc wraps the coroutine fn in a sync dispatcher. Unwrap it # so we subscribe the raw coroutine fn instead of the wrapper (which # would block on run_coroutine_threadsafe from the rx thread). - if hasattr(declared_handler, "aio"): - declared_handler = declared_handler.aio.__get__(self, type(self)) - if not inspect.iscoroutinefunction(declared_handler): + if hasattr(handler, "aio"): + handler = handler.aio.__get__(self, type(self)) + if not inspect.iscoroutinefunction(handler): raise TypeError( - f"{type(self).__name__}.handle_{input_name} must be `async def` " + f"{type(self).__name__}.handle_{port} must be `async def` " "(use a manual self..subscribe(...) for sync handlers)" ) - handler: Callable[..., Any] = declared_handler - if _handler_wants_metadata(handler, f"{type(self).__name__}.handle_{input_name}"): - metadata = TopicMetadata(index=0, name=input_name) - - async def with_meta( - msg: Any, - _handler: Callable[[Any, TopicMetadata], Any] = handler, - _metadata: TopicMetadata = metadata, - ) -> None: - await _handler(msg, _metadata) - - handler = with_meta - bindings.append((in_stream, handler)) - - for in_stream, handler in bindings: - # process_observable runs each handler through a per-subscription - # dispatcher task on self._loop that serializes invocations and - # keeps only the latest unprocessed message. We subscribe to - # pure_observable() because the dispatcher already provides - # backpressure. - self.process_observable(in_stream.pure_observable(), handler) + funnel = self.config.topic_funnels.get(port) + names = funnel.names if funnel else [port] + streams: list[In[Any] | IO[Any]] = ( + [self._funnel_streams[name] for name in names] if funnel else [declared] + ) + metas = None + if _handler_wants_metadata(handler, f"{type(self).__name__}.handle_{port}"): + metas = [ + TopicMetadata( + index=index, name=name, info=funnel.info.get(name, {}) if funnel else {} + ) + for index, name in enumerate(names) + ] + bindings.append((streams, _dispatch_to(handler, metas))) + + for streams, invoke in bindings: + # One dispatcher task per port, on self._loop: it serializes the + # handler and holds only the latest unprocessed message per stream. + on_msg, dispatcher = self._make_keyed_dispatch(invoke) + self.register_disposable(dispatcher) + for index, stream in enumerate(streams): + self.register_disposable(Disposable(stream.subscribe(partial(on_msg, index)))) def _make_funnel_streams(self) -> "dict[str, In[Any]]": """One `In` per topic-funnel entry, keyed by stream name. @@ -744,6 +767,9 @@ def _make_funnel_streams(self) -> "dict[str, In[Any]]": blueprint machinery remaps/namespaces/pins them), but they are not attributes, so `inputs` and the native launch line never see them as ports of their own — the funnel's port stands in for all of them. + + Each starts on the transport its name alone implies, which a coordinator + overwrites when it wires the blueprint and standalone use keeps. """ streams: dict[str, In[Any]] = {} for port, funnel in self.config.topic_funnels.items(): @@ -762,59 +788,11 @@ def _make_funnel_streams(self) -> "dict[str, In[Any]]": raise ValueError( f"topic funnel {port!r} entry {name!r} appears in more than one funnel" ) - streams[name] = In(declared.type, name, self) + stream: In[Any] = In(declared.type, name, self) + stream.transport = make_transport(name, declared.type) + streams[name] = stream return streams - def _bind_topic_funnels(self) -> None: - """For each `topic_funnels` port with an `async def handle_`, subscribe - every stream in the group into that one handler. - - A native module defines no such method — its subprocess subscribes instead — - so this is a no-op there. - """ - for port, group in self.config.topic_funnels.items(): - declared_handler = getattr(self, f"handle_{port}", None) - if declared_handler is None: - continue - if hasattr(declared_handler, "aio"): - declared_handler = declared_handler.aio.__get__(self, type(self)) - if not inspect.iscoroutinefunction(declared_handler): - raise TypeError( - f"{type(self).__name__}.handle_{port} must be `async def` " - "(topic funnels have no sync path)" - ) - handler: Callable[..., Any] = declared_handler - if _handler_wants_metadata(handler, f"{type(self).__name__}.handle_{port}"): - metas = tuple( - TopicMetadata(index=index, name=name, info=group.info.get(name, {})) - for index, name in enumerate(group.names) - ) - - async def invoke( - index: int, - msg: Any, - _handler: Callable[[Any, TopicMetadata], Any] = handler, - _metas: tuple[TopicMetadata, ...] = metas, - ) -> None: - await _handler(msg, _metas[index]) - else: - - async def invoke( # type: ignore[misc] - index: int, - msg: Any, - _handler: Callable[[Any], Any] = handler, - ) -> None: - await _handler(msg) - - on_msg, dispatcher_disp = self._make_keyed_dispatch(invoke) - self.register_disposable(dispatcher_disp) - for index, name in enumerate(group.names): - stream = self._funnel_streams[name] - if getattr(stream, "_transport", None) is None: - # Not wired by a coordinator (standalone use): default transport. - stream.transport = make_transport(name, stream.type) - self.register_disposable(Disposable(stream.subscribe(partial(on_msg, index)))) - def _make_async_dispatch( self, async_handler: Callable[[Any], Any] ) -> tuple[Callable[[Any], None], "DisposableBase"]: @@ -830,10 +808,7 @@ def _make_async_dispatch( - The returned Disposable cancels the dispatcher task. """ - async def single(_: int, msg: Any) -> None: - await async_handler(msg) - - on_msg, disposable = self._make_keyed_dispatch(single) + on_msg, disposable = self._make_keyed_dispatch(_dispatch_to(async_handler, None)) return partial(on_msg, 0), disposable def _make_keyed_dispatch( diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index e538a9bb32..94153205b5 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -52,7 +52,7 @@ class MyCppModule(NativeModule): import sys import threading import time -from typing import IO, Any, TypedDict +from typing import IO, Any, TypedDict, cast from pydantic import Field, model_validator @@ -60,7 +60,8 @@ 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.core.transport_factory import channel_for, session_config +from dimos.core.transport import PubSubTransport +from dimos.core.transport_factory import session_config from dimos.protocol.service.spec import SessionConfig from dimos.utils.logging_config import setup_logger @@ -550,12 +551,7 @@ def _collect_topics(self) -> TopicsMap: for port, funnel in self.config.topic_funnels.items(): entries: list[str | TopicEntry] = [] for name in funnel.names: - stream = self._funnel_streams[name] - transport = getattr(stream, "_transport", None) - channel = getattr(transport, "channel", None) - if channel is None: - # Unwired (standalone use): the channel the default transport lands on. - channel = channel_for(name, stream.type) + channel = cast("PubSubTransport[Any]", self._funnel_streams[name].transport).channel info = funnel.info.get(name) entries.append(channel if info is None else TopicEntry(topic=channel, info=info)) topics[port] = entries diff --git a/dimos/core/test_transport_factory.py b/dimos/core/test_transport_factory.py index f6491bad09..69e8e7d519 100644 --- a/dimos/core/test_transport_factory.py +++ b/dimos/core/test_transport_factory.py @@ -25,7 +25,6 @@ ) from dimos.core.transport_factory import ( apply_transport_arg, - channel_for, default_zenoh_qos, make_transport, rpc_backend, @@ -81,14 +80,6 @@ def test_make_transport_zenoh_pickled() -> None: assert t.topic == "dimos/human_input" -@pytest.mark.parametrize("g", [LCM, ZENOH]) -@pytest.mark.parametrize(("name", "msg_type"), [("/camera/color", Image), ("/human_input", None)]) -def test_channel_for_matches_the_transport_it_skips_building( - g: GlobalConfig, name: str, msg_type: type | None -) -> None: - assert channel_for(name, msg_type, g=g) == make_transport(name, msg_type, g=g).channel - - def test_default_zenoh_qos_high_rate_sensor_types_drop() -> None: assert default_zenoh_qos("/camera/color", Image) == QOS_LATEST_WINS diff --git a/dimos/core/transport_factory.py b/dimos/core/transport_factory.py index abfb901d2a..597c0c344f 100644 --- a/dimos/core/transport_factory.py +++ b/dimos/core/transport_factory.py @@ -27,7 +27,6 @@ pLCMTransport, pZenohTransport, ) -from dimos.protocol.pubsub.impl.lcmpubsub import Topic as LCMTopic from dimos.protocol.pubsub.impl.zenohpubsub import ( QOS_LATEST_WINS, QOS_NEVER_DROP, @@ -113,22 +112,6 @@ def make_transport( return LCMTransport(topic, msg_type) -def channel_for(name: str, msg_type: type | None = None, *, g: GlobalConfig = global_config) -> str: - """The wire channel `make_transport` would land on, without building one. - - Lets a caller name a channel it never subscribes to itself, such as the extra - topics a native module fans into one handler. - """ - use_pickled = msg_type is None or getattr(msg_type, "lcm_encode", None) is None - topic = transport_topic(name, g) - if g.transport == "zenoh": - return str(ZenohTopic(topic, None if use_pickled else msg_type).key_expr) - if use_pickled: - return topic - assert msg_type is not None - return str(LCMTopic(topic, msg_type)) - - def _transport_arg_error(argv: list[str], message: str) -> NoReturn: """Print an argparse-style CLI error for `--transport` and exit(2).""" prog = os.path.basename(argv[0]) if argv else "dimos"