diff --git a/dimos/core/coordination/blueprints.py b/dimos/core/coordination/blueprints.py index f21ff3fe30..a84755c61c 100644 --- a/dimos/core/coordination/blueprints.py +++ b/dimos/core/coordination/blueprints.py @@ -12,10 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass, field, replace from functools import cached_property, reduce import operator +import re import sys import types as types_mod from types import MappingProxyType @@ -85,6 +86,14 @@ class BlueprintAtom: module: type[ModuleBase] streams: tuple[StreamRef, ...] module_refs: tuple[ModuleRef, ...] + # Set when the same module class appears more than once in a blueprint + # (e.g. one per robot). `None` means just one. + instance_name: str | None = None + + @property + def name(self) -> str: + """The key identifying this module instance within a blueprint.""" + return self.instance_name if self.instance_name is not None else self.module.name @classmethod def create(cls, module: type[ModuleBase], kwargs: dict[str, Any]) -> Self: @@ -133,11 +142,16 @@ 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)) + 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") + return cls( module=module, streams=tuple(streams), module_refs=tuple(module_refs), kwargs=kwargs, + instance_name=instance_name, ) @@ -153,9 +167,12 @@ class Blueprint: default_factory=lambda: MappingProxyType({}) ) global_config_overrides: Mapping[str, Any] = field(default_factory=lambda: MappingProxyType({})) - remapping_map: Mapping[tuple[type[ModuleBase], str], str | type[ModuleBase] | type[Spec]] = ( - field(default_factory=lambda: MappingProxyType({})) + + # Keyed by (instance name, stream/ref name). + remapping_map: Mapping[tuple[str, str], str | type[ModuleBase] | type[Spec]] = field( + default_factory=lambda: MappingProxyType({}) ) + requirement_checks: tuple[Callable[[], str | None], ...] = field(default_factory=tuple) configurator_checks: "tuple[SystemConfigurator, ...]" = field(default_factory=tuple) @@ -180,11 +197,19 @@ def disabled_modules(self, *modules: type[ModuleBase]) -> "Blueprint": return replace(self, disabled_modules_tuple=self.disabled_modules_tuple + modules) def config(self) -> type: - configs = { - b.module.name: (get_type_hints(b.module)["config"] | None, None) - for b in self.blueprints - } + configs = {} + + for b in self.blueprints: + key = config_key(b.name) + if key in configs: + raise ValueError( + f"Config key collision: two module instances map to {key!r}. " + f"Rename one of the instances." + ) + configs[key] = (get_type_hints(b.module)["config"] | None, None) + configs["g"] = (GlobalConfig | None, None) + return create_model("BlueprintConfig", __config__={"extra": "forbid"}, **configs) # type: ignore[call-overload,no-any-return] def transports(self, transports: dict[tuple[str, type], Any]) -> "Blueprint": @@ -198,13 +223,26 @@ def global_config(self, **kwargs: Any) -> "Blueprint": def remappings( self, - remappings: list[tuple[type[ModuleBase], str, str | type[ModuleBase] | type[Spec]]], + remappings: Sequence[ + tuple[type[ModuleBase] | str, str, str | type[ModuleBase] | type[Spec]] + ], ) -> "Blueprint": remappings_dict = dict(self.remapping_map) for module, old, new in remappings: - remappings_dict[(module, old)] = new + remappings_dict[(self._instance_key(module), old)] = new return replace(self, remapping_map=MappingProxyType(remappings_dict)) + def _instance_key(self, module: type[ModuleBase] | str) -> str: + if isinstance(module, str): + return module + names = [b.name for b in self.blueprints if b.module is module] + if len(names) > 1: + raise ValueError( + f"{module.__name__} has multiple instances in this blueprint " + f"({', '.join(sorted(names))}). Pass the instance name instead of the class." + ) + return names[0] if names else module.name + def requirements(self, *checks: Callable[[], str | None]) -> "Blueprint": return replace(self, requirement_checks=self.requirement_checks + tuple(checks)) @@ -251,7 +289,113 @@ def _eliminate_duplicates(blueprints: list[BlueprintAtom]) -> list[BlueprintAtom seen = set() unique_blueprints = [] for bp in reversed(blueprints): - if bp.module not in seen: - seen.add(bp.module) + if bp.name not in seen: + seen.add(bp.name) unique_blueprints.append(bp) return list(reversed(unique_blueprints)) + + +def config_key(instance_name: str) -> str: + """Escape an instance name into a valid config/CLI/env identifier.""" + return instance_name.replace("/", "_") + + +def namespace(prefix: str, *blueprints: Blueprint, expose: Iterable[str] = ()) -> Blueprint: + """Isolate *blueprints* under a name prefix so several copies can coexist. + + Instance names, stream names (and so their topics), TF frame ids, and RPC + topics all get the `{prefix}/` prefix, which disconnects the namespaced + modules from everything outside. + + Stream names listed in *expose* are left unprefixed, so they connect + globally by the usual (name, type) matching. That is how data crosses the + namespace boundary: + + fleet = autoconnect( + AggregateMapper.blueprint(), # shared: sees every robot's pointcloud + *[ + namespace(f"robot{i}", GO2Connection.blueprint(ip=ip), expose={"pointcloud"}) + for i, ip in enumerate(ips) + ], + ) + """ + if not re.fullmatch(r"[A-Za-z0-9_]+", prefix): + raise ValueError( + f"Invalid namespace prefix {prefix!r}; use letters, digits and underscores " + f"(nest namespaces by composition, not with '/')." + ) + + merged = autoconnect(*blueprints) + expose_set = frozenset(expose) + + effective_names = set() + for atom in merged.blueprints: + for stream in atom.streams: + effective = merged.remapping_map.get((atom.name, stream.name), stream.name) + if isinstance(effective, str): + effective_names.add(effective) + + unknown = expose_set - effective_names + if unknown: + raise ValueError( + f"expose names {sorted(unknown)} do not match any stream in the " + f"namespaced blueprint (available: {sorted(effective_names)})" + ) + + new_atoms = [] + new_remap: dict[tuple[str, str], str | type[ModuleBase] | type[Spec]] = {} + + for atom in merged.blueprints: + new_name = f"{prefix}/{atom.name}" + kwargs = dict(atom.kwargs) + kwargs["instance_name"] = new_name + old_namespace = atom.name.rsplit("/", 1)[0] if "/" in atom.name else "" + frame_id_prefix = kwargs.get("frame_id_prefix") + if frame_id_prefix is None: + kwargs["frame_id_prefix"] = prefix + elif old_namespace and frame_id_prefix == old_namespace: + # Auto-set by an inner namespace(); extend it. User-set values are kept. + kwargs["frame_id_prefix"] = f"{prefix}/{frame_id_prefix}" + new_atoms.append(replace(atom, kwargs=kwargs, instance_name=new_name)) + + for stream in atom.streams: + effective = merged.remapping_map.get((atom.name, stream.name), stream.name) + if not isinstance(effective, str): + continue + if effective in expose_set: + if (atom.name, stream.name) in merged.remapping_map: + new_remap[new_name, stream.name] = effective + else: + new_remap[new_name, stream.name] = f"{prefix}/{effective}" + + # Module-ref remappings (values that are classes/Specs) keep their values. + for (instance, ref_name), value in merged.remapping_map.items(): + if not isinstance(value, str): + new_remap[f"{prefix}/{instance}", ref_name] = value + + new_transports = {} + for (name, type_), transport in merged.transport_map.items(): + if name in expose_set: + new_transports[name, type_] = transport + else: + new_transports[f"{prefix}/{name}", type_] = _reprefix_transport(transport, prefix) + + return replace( + merged, + blueprints=tuple(new_atoms), + remapping_map=MappingProxyType(new_remap), + transport_map=MappingProxyType(new_transports), + ) + + +def _reprefix_transport(transport: PubSubTransport[Any], prefix: str) -> PubSubTransport[Any]: + """Clone a pinned transport with its topic moved under the namespace prefix.""" + cls, args = transport.__reduce__() # type: ignore[misc] + if not (isinstance(args, tuple) and args and isinstance(args[0], str)): + raise ValueError( + f"Cannot namespace pinned transport {transport!r}; pin it on the " + f"namespaced blueprint with an explicit topic instead." + ) + topic = args[0] + new_topic = f"/{prefix}{topic}" if topic.startswith("/") else f"{prefix}/{topic}" + return cls(new_topic, *args[1:]) # type: ignore[no-any-return] diff --git a/dimos/core/coordination/module_coordinator.py b/dimos/core/coordination/module_coordinator.py index 849c928681..4950dc134c 100644 --- a/dimos/core/coordination/module_coordinator.py +++ b/dimos/core/coordination/module_coordinator.py @@ -16,6 +16,7 @@ from collections import defaultdict from collections.abc import Callable, Mapping, MutableMapping +import dataclasses import importlib import inspect import shutil @@ -27,7 +28,7 @@ from dimos.core.coordination.worker_manager import WorkerManager from dimos.core.coordination.worker_manager_python import WorkerManagerPython from dimos.core.global_config import GlobalConfig, global_config -from dimos.core.module import ModuleBase, ModuleSpec +from dimos.core.module import ModuleBase, ModuleSpec, is_module_type from dimos.core.resource import Resource from dimos.core.transport import ( LCMTransport, @@ -55,12 +56,16 @@ class ModuleDescriptor(NamedTuple): class_name: str qualified_path: str rpc_names: list[str] + # RPC topic prefix: the class name, or the instance name for modules + # deployed under a non-default name. Defaulted so the tuple stays + # wire-compatible with older daemons. + rpc_name: str = "" class ModuleCoordinator(Resource): _managers: dict[str, WorkerManager] _global_config: GlobalConfig - _deployed_modules: dict[type[ModuleBase], ModuleProxyProtocol] + _deployed_modules: dict[str, ModuleProxyProtocol] def __init__( self, @@ -70,11 +75,12 @@ def __init__( manager_types: list[type[WorkerManager]] = [WorkerManagerPython] self._managers = {cls.deployment_identifier: cls(g=g) for cls in manager_types} self._deployed_modules = {} - self._deployed_atoms: dict[type[ModuleBase], BlueprintAtom] = {} - self._resolved_module_refs: dict[tuple[type[ModuleBase], str], type[ModuleBase]] = {} + self._instance_classes: dict[str, type[ModuleBase]] = {} + self._deployed_atoms: dict[str, BlueprintAtom] = {} + self._resolved_module_refs: dict[tuple[str, str], str] = {} self._transport_registry: dict[tuple[str, type], PubSubTransport[Any]] = {} self._class_aliases: dict[type[ModuleBase], type[ModuleBase]] = {} - self._module_transports: dict[type[ModuleBase], dict[str, PubSubTransport[Any]]] = {} + self._module_transports: dict[str, dict[str, PubSubTransport[Any]]] = {} self._started = False self._modules_lock = threading.RLock() self._coordinator_rpc: CoordinatorRPC | None = None @@ -92,13 +98,13 @@ def stop(self) -> None: self._coordinator_rpc.stop() self._coordinator_rpc = None - for module_class, module in reversed(self._deployed_modules.items()): - logger.info("Stopping module...", module=module_class.__name__) + for name, module in reversed(self._deployed_modules.items()): + logger.info("Stopping module...", module=name) try: module.stop() except Exception: - logger.error("Error stopping module", module=module_class.__name__, exc_info=True) - logger.info("Module stopped.", module=module_class.__name__) + logger.error("Error stopping module", module=name, exc_info=True) + logger.info("Module stopped.", module=name) def _stop_manager(m: WorkerManager) -> None: try: @@ -123,6 +129,7 @@ def rpcs(self) -> dict[str, Callable[..., Any]]: "load_blueprint_by_name": self.load_blueprint_by_name, "load_blueprint": self.load_blueprint, "restart_module_by_class_name": self.restart_module_by_class_name, + "restart_module_by_name": self.restart_module_by_name, } def ping(self) -> str: @@ -132,13 +139,14 @@ def ping(self) -> str: def list_modules(self) -> list[ModuleDescriptor]: with self._modules_lock: descriptors: list[ModuleDescriptor] = [] - for cls in self._deployed_modules: + for name, cls in self._instance_classes.items(): qualified = f"{cls.__module__}.{cls.__name__}" descriptors.append( ModuleDescriptor( class_name=cls.__name__, qualified_path=qualified, rpc_names=list(cls.rpcs.keys()), + rpc_name=_rpc_name(name, cls), ) ) return descriptors @@ -151,7 +159,7 @@ def load_blueprint_by_name(self, name: str) -> None: def list_module_names(self) -> list[str]: with self._modules_lock: - return [cls.__name__ for cls in self._deployed_modules] + return [_rpc_name(name, cls) for name, cls in self._instance_classes.items()] def health_check(self) -> bool: return all(m.health_check() for m in self._managers.values()) @@ -176,8 +184,10 @@ def deploy( deployed_module = self._managers[module_class.deployment].deploy( module_class, global_config, kwargs ) + name = kwargs.get("instance_name") or module_class.name with self._modules_lock: - self._deployed_modules[module_class] = deployed_module + self._deployed_modules[name] = deployed_module + self._instance_classes[name] = module_class return deployed_module # type: ignore[return-value] def deploy_parallel( @@ -209,13 +219,12 @@ def _deploy_group(dep: str) -> None: raise with self._modules_lock: - self._deployed_modules.update( - { - cls: mod - for (cls, _, _), mod in zip(module_specs, results, strict=True) - if mod is not None - } - ) + for (cls, _, kwargs), mod in zip(module_specs, results, strict=True): + if mod is None: + continue + name = kwargs.get("instance_name") or cls.name + self._deployed_modules[name] = mod + self._instance_classes[name] = cls return results def build_all_modules(self) -> None: @@ -247,8 +256,36 @@ def start_all_modules(self) -> None: def _resolve_class(self, cls: type[ModuleBase]) -> type[ModuleBase]: return self._class_aliases.get(cls, cls) - def get_instance(self, module: type[ModuleBase]) -> ModuleProxy: - return self._deployed_modules.get(self._resolve_class(module)) # type: ignore[return-value] + def _instance_keys_of(self, module: type[ModuleBase]) -> list[str]: + cls = self._resolve_class(module) + return [n for n, c in self._instance_classes.items() if self._resolve_class(c) is cls] + + def _resolve_instance_key(self, module: type[ModuleBase] | str) -> str: + """Resolve a module class or instance name to the deployed instance name.""" + if isinstance(module, str): + if module not in self._deployed_modules: + raise ValueError(f"{module!r} is not deployed") + return module + names = self._instance_keys_of(module) + if not names: + raise ValueError(f"{module.__name__} is not deployed") + if len(names) > 1: + raise ValueError( + f"Multiple instances of {module.__name__} are deployed " + f"({', '.join(sorted(names))}); pass the instance name." + ) + return names[0] + + def get_instance(self, module: type[ModuleBase] | str) -> ModuleProxy: + if isinstance(module, str): + return self._deployed_modules.get(module) # type: ignore[return-value] + names = self._instance_keys_of(module) + if len(names) > 1: + raise ValueError( + f"Multiple instances of {module.__name__} are deployed " + f"({', '.join(sorted(names))}); pass the instance name." + ) + return self._deployed_modules.get(names[0]) if names else None # type: ignore[return-value] def _send_on_system_modules(self) -> None: modules = list(self._deployed_modules.values()) @@ -257,13 +294,13 @@ def _send_on_system_modules(self) -> None: module.on_system_modules(modules) def _connect_streams(self, blueprint: Blueprint) -> None: - streams: dict[tuple[str, type], list[tuple[type, str]]] = defaultdict(list) + streams: dict[tuple[str, type], list[tuple[str, str]]] = defaultdict(list) for bp in blueprint.active_blueprints: for conn in bp.streams: - remapped_name = blueprint.remapping_map.get((bp.module, conn.name), conn.name) + remapped_name = blueprint.remapping_map.get((bp.name, conn.name), conn.name) if isinstance(remapped_name, str): - streams[remapped_name, conn.type].append((bp.module, conn.name)) + streams[remapped_name, conn.type].append((bp.name, conn.name)) for remapped_name, stream_type in streams.keys(): key = (remapped_name, stream_type) @@ -272,17 +309,17 @@ def _connect_streams(self, blueprint: Blueprint) -> None: else: transport = _get_transport_for(blueprint, remapped_name, stream_type) self._transport_registry[key] = transport - for module, original_name in streams[key]: - instance = self.get_instance(module) # type: ignore[assignment] + for instance_key, original_name in streams[key]: + instance = self.get_instance(instance_key) # type: ignore[assignment] instance.set_transport(original_name, transport) # type: ignore[union-attr] - self._module_transports.setdefault(module, {})[original_name] = transport + self._module_transports.setdefault(instance_key, {})[original_name] = transport logger.info( "Transport", name=remapped_name, original_name=original_name, topic=str(getattr(transport, "topic", None)), type=f"{stream_type.__module__}.{stream_type.__qualname__}", - module=module.__name__, + module=instance_key, transport=transport.__class__.__name__, ) @@ -360,18 +397,29 @@ def _load_blueprint( # Reject duplicate modules. for bp in blueprint.active_blueprints: - if bp.module in self._deployed_modules: + if bp.name in self._deployed_modules: raise ValueError( - f"{bp.module.__name__} is already deployed; cannot load the same module twice" + f"{bp.name} is already deployed; cannot load the same module twice" ) before = set(self._deployed_modules) + existing_atoms = tuple( + atom for name in before if (atom := self._deployed_atoms.get(name)) is not None + ) + existing_classes = {self._instance_classes[name] for name in before} _deploy_all_modules(blueprint, self, self._global_config, blueprint_args) self._connect_streams(blueprint) - _connect_module_refs(blueprint, self, existing_modules=before) + _connect_module_refs( + blueprint, + self, + existing_atoms=existing_atoms, + existing_modules=existing_classes, + ) - new_modules = [proxy for cls, proxy in self._deployed_modules.items() if cls not in before] + new_modules = [ + proxy for name, proxy in self._deployed_modules.items() if name not in before + ] if new_modules: safe_thread_map(new_modules, lambda m: m.build()) @@ -386,35 +434,35 @@ def load_module( ) -> None: self.load_blueprint(module_class.blueprint(**blueprint_args or {})) - def unload_module(self, module_class: type[ModuleBase]) -> None: + def unload_module(self, module: type[ModuleBase] | str) -> None: """Stop and tear down a single deployed module. - Removes the module from coordinator state, stops its worker-side - instance, and shuts down the worker process if it becomes empty. - Stream transports and other modules' references are left intact — - callers that expect the module to come back (e.g. ``restart_module``) - are responsible for rewiring. + Accepts a module class (must have exactly one deployed instance) or an + instance name. Removes the module from coordinator state, stops its + worker-side instance, and shuts down the worker process if it becomes + empty. Stream transports and other modules' references are left + intact — callers that expect the module to come back (e.g. + ``restart_module``) are responsible for rewiring. """ with self._modules_lock: - self._unload_module(module_class) + self._unload_module(module) - def _unload_module(self, module_class: type[ModuleBase]) -> None: - module_class = self._resolve_class(module_class) - if module_class not in self._deployed_modules: - raise ValueError(f"{module_class.__name__} is not deployed") + def _unload_module(self, module: type[ModuleBase] | str) -> None: + name = self._resolve_instance_key(module) + module_class = self._instance_classes[name] if module_class.deployment != "python": raise NotImplementedError( f"unload_module only supports python deployment, got {module_class.deployment!r}" ) - proxy = self._deployed_modules[module_class] + proxy = self._deployed_modules[name] try: proxy.stop() except Exception: logger.error( "Error stopping module during unload", - module=module_class.__name__, + module=name, exc_info=True, ) @@ -424,20 +472,22 @@ def _unload_module(self, module_class: type[ModuleBase]) -> None: except Exception: logger.error( "Error undeploying module from worker", - module=module_class.__name__, + module=name, exc_info=True, ) - del self._deployed_modules[module_class] - self._deployed_atoms.pop(module_class, None) - self._module_transports.pop(module_class, None) - self._class_aliases = { - k: v for k, v in self._class_aliases.items() if v is not module_class - } + del self._deployed_modules[name] + del self._instance_classes[name] + self._deployed_atoms.pop(name, None) + self._module_transports.pop(name, None) + if not any(c is module_class for c in self._instance_classes.values()): + self._class_aliases = { + k: v for k, v in self._class_aliases.items() if v is not module_class + } self._resolved_module_refs = { key: target for key, target in self._resolved_module_refs.items() - if key[0] is not module_class and target is not module_class + if key[0] != name and target != name } def restart_module_by_class_name( @@ -447,58 +497,72 @@ def restart_module_by_class_name( reload_source: bool = True, ) -> None: with self._modules_lock: - for cls in self._deployed_modules: - if cls.__name__ == class_name: - self._restart_module(cls, reload_source=reload_source) - return + names = [n for n, c in self._instance_classes.items() if c.__name__ == class_name] + if len(names) > 1: + raise ValueError( + f"Multiple instances of {class_name!r} are deployed " + f"({', '.join(sorted(names))}); use restart_module_by_name." + ) + if names: + self._restart_module(names[0], reload_source=reload_source) + return raise ValueError(f"No deployed module with class name {class_name!r}") + def restart_module_by_name( + self, + name: str, + *, + reload_source: bool = True, + ) -> None: + with self._modules_lock: + self._restart_module(name, reload_source=reload_source) + def restart_module( self, - module_class: type[ModuleBase], + module: type[ModuleBase] | str, *, reload_source: bool = True, ) -> ModuleProxyProtocol: """Restart a single deployed module in place. - Unloads *module_class*, optionally reloads its source file via - ``importlib.reload`` so edited code is picked up, then redeploys it + Accepts a module class (must have exactly one deployed instance) or an + instance name. Unloads the module, optionally reloads its source file + via ``importlib.reload`` so edited code is picked up, then redeploys it onto a fresh worker process, reconnects its streams to the existing transports, and re-injects the new proxy into every other module that held a reference to it. """ with self._modules_lock: - return self._restart_module(module_class, reload_source=reload_source) + return self._restart_module(module, reload_source=reload_source) def _restart_module( self, - module_class: type[ModuleBase], + module: type[ModuleBase] | str, *, reload_source: bool = True, ) -> ModuleProxyProtocol: - module_class = self._resolve_class(module_class) - if module_class not in self._deployed_modules: - raise ValueError(f"{module_class.__name__} is not deployed") + name = self._resolve_instance_key(module) + module_class = self._instance_classes[name] if module_class.deployment != "python": raise NotImplementedError( f"restart_module only supports python deployment, got {module_class.deployment!r}" ) - old_atom = self._deployed_atoms[module_class] + old_atom = self._deployed_atoms[name] kwargs = dict(old_atom.kwargs) - saved_transports = dict(self._module_transports.get(module_class, {})) + saved_transports = dict(self._module_transports.get(name, {})) inbound_refs = [ (consumer, ref_name) for (consumer, ref_name), target in self._resolved_module_refs.items() - if target is module_class + if target == name ] outbound_refs = [ (ref_name, target) for (consumer, ref_name), target in self._resolved_module_refs.items() - if consumer is module_class + if consumer == name ] - self.unload_module(module_class) + self.unload_module(name) if reload_source: source_mod = sys.modules.get(module_class.__module__) @@ -517,35 +581,38 @@ def _restart_module( python_wm = cast("WorkerManagerPython", self._managers["python"]) new_proxy = python_wm.deploy_fresh(new_class, self._global_config, kwargs) - self._deployed_modules[new_class] = new_proxy + self._deployed_modules[name] = new_proxy + self._instance_classes[name] = new_class new_bp = new_class.blueprint(**kwargs) new_atom = new_bp.active_blueprints[0] - self._deployed_atoms[new_class] = new_atom + if old_atom.instance_name is not None: + new_atom = dataclasses.replace(new_atom, instance_name=old_atom.instance_name) + self._deployed_atoms[name] = new_atom for stream_ref in new_atom.streams: transport = saved_transports.get(stream_ref.name) if transport is not None: new_proxy.set_transport(stream_ref.name, transport) - self._module_transports[new_class] = { + self._module_transports[name] = { s.name: t for s in new_atom.streams if (t := saved_transports.get(s.name)) is not None } - for consumer_class, ref_name in inbound_refs: - consumer_proxy = self._deployed_modules.get(consumer_class) + for consumer_key, ref_name in inbound_refs: + consumer_proxy = self._deployed_modules.get(consumer_key) if consumer_proxy is None: continue setattr(consumer_proxy, ref_name, new_proxy) consumer_proxy.set_module_ref(ref_name, new_proxy) # type: ignore[attr-defined] - self._resolved_module_refs[consumer_class, ref_name] = new_class + self._resolved_module_refs[consumer_key, ref_name] = name - for ref_name, target_class in outbound_refs: - target_proxy = self._deployed_modules.get(target_class) + for ref_name, target_key in outbound_refs: + target_proxy = self._deployed_modules.get(target_key) if target_proxy is None: continue setattr(new_proxy, ref_name, target_proxy) new_proxy.set_module_ref(ref_name, target_proxy) # type: ignore[attr-defined] - self._resolved_module_refs[new_class, ref_name] = target_class + self._resolved_module_refs[name, ref_name] = target_key new_proxy.build() new_proxy.start() @@ -564,11 +631,16 @@ def loop(self) -> None: self.stop() +def _rpc_name(instance_key: str, cls: type[ModuleBase]) -> str: + """The module's RPC topic prefix: class name unless an instance name is set.""" + return cls.__name__ if instance_key == cls.name else instance_key + + def _all_name_types(blueprint: Blueprint) -> set[tuple[str, type]]: result = set() for bp in blueprint.active_blueprints: for conn in bp.streams: - remapped_name = blueprint.remapping_map.get((bp.module, conn.name), conn.name) + remapped_name = blueprint.remapping_map.get((bp.name, conn.name), conn.name) if isinstance(remapped_name, str): result.add((remapped_name, conn.type)) return result @@ -618,7 +690,7 @@ def _verify_no_name_conflicts(blueprint: Blueprint) -> None: for bp in blueprint.active_blueprints: for conn in bp.streams: - stream_name = blueprint.remapping_map.get((bp.module, conn.name), conn.name) + stream_name = blueprint.remapping_map.get((bp.name, conn.name), conn.name) name_to_types[stream_name].add(conn.type) name_to_modules[stream_name].append((bp.module, conn.type)) @@ -662,7 +734,7 @@ def _verify_no_conflicts_with_existing( for bp in blueprint.active_blueprints: for conn in bp.streams: - remapped_name = blueprint.remapping_map.get((bp.module, conn.name), conn.name) + remapped_name = blueprint.remapping_map.get((bp.name, conn.name), conn.name) if isinstance(remapped_name, str) and remapped_name in existing_names: for existing_type in existing_names[remapped_name]: if existing_type != conn.type: @@ -715,12 +787,15 @@ def _deploy_all_modules( ) -> None: module_specs: list[ModuleSpec] = [] for bp in blueprint.active_blueprints: - module_specs.append((bp.module, gc, bp.kwargs.copy())) + kwargs = bp.kwargs.copy() + if bp.instance_name is not None: + kwargs.setdefault("instance_name", bp.instance_name) + module_specs.append((bp.module, gc, kwargs)) module_coordinator.deploy_parallel(module_specs, blueprint_args) for bp in blueprint.active_blueprints: - module_coordinator._deployed_atoms[bp.module] = bp + module_coordinator._deployed_atoms[bp.name] = bp def _ref_msg(module_name: str, ref: object, spec_name: str, detail: str) -> str: @@ -730,42 +805,82 @@ def _ref_msg(module_name: str, ref: object, spec_name: str, detail: str) -> str: ) +def _atom_namespace(instance_name: str) -> str: + return instance_name.rsplit("/", 1)[0] if "/" in instance_name else "" + + +def _namespace_levels(consumer_namespace: str) -> list[str]: + """Ref search order: the consumer's namespace, enclosing ones, then global.""" + levels = [] + namespace = consumer_namespace + while namespace: + levels.append(namespace) + namespace = _atom_namespace(namespace) + levels.append("") + return levels + + def _resolve_single_ref( bp: Any, module_ref: Any, spec: Any, blueprint: Blueprint, disabled_set: set[type], + existing_atoms: tuple[BlueprintAtom, ...] = (), existing_modules: set[type[ModuleBase]] | None = None, ) -> Any: """Resolve a module ref to its provider. - Returns a module type, a ``DisabledModuleProxy``, or *None* (skip). + Returns an instance name, a module type (provider outside the blueprint), + a ``DisabledModuleProxy``, or *None* (skip). """ - from dimos.core.coordination.blueprints import DisabledModuleProxy + from dimos.core.coordination.blueprints import BlueprintAtom, DisabledModuleProxy m = bp.module.__name__ s = module_ref.spec.__name__ + is_class_ref = is_module_type(spec) + + def satisfies(cls: type) -> bool: + return cls is spec if is_class_ref else spec_structural_compliance(cls, spec) + + def module_of(candidate: Any) -> type[ModuleBase]: + return candidate.module if isinstance(candidate, BlueprintAtom) else candidate + + def result_of(candidate: Any) -> Any: + return candidate.name if isinstance(candidate, BlueprintAtom) else candidate - possible = [ - other.module - for other in blueprint.active_blueprints - if other != bp and spec_structural_compliance(other.module, spec) - ] - if existing_modules: - bp_module_set = {o.module for o in blueprint.active_blueprints} - for mod_cls in existing_modules: - if ( - mod_cls != bp.module - and mod_cls not in bp_module_set - and spec_structural_compliance(mod_cls, spec) - ): - possible.append(mod_cls) - valid = [c for c in possible if spec_annotation_compliance(c, spec)] + def display(candidate: Any) -> str: + return candidate.name if isinstance(candidate, BlueprintAtom) else candidate.__name__ + + active_atoms = tuple(blueprint.active_blueprints) + atom_candidates = active_atoms + existing_atoms + atom_module_set = {candidate.module for candidate in atom_candidates} + + # Search the consumer's own namespace first so a per-robot consumer binds + # to its own robot's provider instead of colliding with the other robots'. + possible: list[Any] = [] + for level in _namespace_levels(_atom_namespace(bp.name)): + possible = [ + other + for other in atom_candidates + if other != bp and _atom_namespace(other.name) == level and satisfies(other.module) + ] + if level == "" and existing_modules: + possible += [ + mod_cls + for mod_cls in existing_modules + if mod_cls != bp.module and mod_cls not in atom_module_set and satisfies(mod_cls) + ] + if possible: + break if not possible: if module_ref.optional: return None + if is_class_ref: + # The provider class is not in the blueprint; keep the class and + # let get_instance resolve it (legacy behavior, possibly None). + return spec disabled = next( ( other.module @@ -785,6 +900,8 @@ def _resolve_single_ref( return DisabledModuleProxy(s) raise Exception(_ref_msg(m, module_ref, s, "No module met that spec.")) + valid = [c for c in possible if is_class_ref or spec_annotation_compliance(module_of(c), spec)] + if len(possible) == 1: if not valid: logger.warning( @@ -792,15 +909,15 @@ def _resolve_single_ref( m, module_ref, s, - f"{possible[0].__name__} met the spec structurally but had " + f"{display(possible[0])} met the spec structurally but had " f"annotation mismatches.\nPlease either change the {s} spec " - f"or the {possible[0].__name__} module.", + f"or the {module_of(possible[0]).__name__} module.", ) ) - return possible[0] + return result_of(possible[0]) if len(valid) == 1: - return valid[0] + return result_of(valid[0]) if len(valid) > 1: raise Exception( @@ -808,14 +925,14 @@ def _resolve_single_ref( m, module_ref, s, - f"Multiple modules met that spec: {valid}.\n" + f"Multiple modules met that spec: {[display(c) for c in valid]}.\n" f"To fix this use .remappings, for example:\n" f" autoconnect(...).remappings([ ({m}, {module_ref.name!r}, " f") ])", ) ) - names = ", ".join(c.__name__ for c in possible) + names = ", ".join(display(c) for c in possible) raise Exception( _ref_msg( m, @@ -829,6 +946,7 @@ def _resolve_single_ref( def _connect_module_refs( blueprint: Blueprint, module_coordinator: ModuleCoordinator, + existing_atoms: tuple[BlueprintAtom, ...] = (), existing_modules: set[type[ModuleBase]] | None = None, ) -> None: from dimos.core.coordination.blueprints import DisabledModuleProxy @@ -836,51 +954,58 @@ def _connect_module_refs( from dimos.core.rpc_client import AsyncSpecProxy mod_and_mod_ref_to_proxy = { - (module, name): replacement - for (module, name), replacement in blueprint.remapping_map.items() + (instance_key, name): replacement + for (instance_key, name), replacement in blueprint.remapping_map.items() if is_spec(replacement) or is_module_type(replacement) } # Track the consumer's declared spec for each ref so we can wrap the proxy # below if the spec contains async-declared methods. - declared_spec: dict[tuple[type[ModuleBase], str], Any] = {} + declared_spec: dict[tuple[str, str], Any] = {} - disabled_ref_proxies: dict[tuple[type[ModuleBase], str], DisabledModuleProxy] = {} + disabled_ref_proxies: dict[tuple[str, str], DisabledModuleProxy] = {} disabled_set = set(blueprint.disabled_modules_tuple) for bp in blueprint.active_blueprints: for module_ref in bp.module_refs: - declared_spec[bp.module, module_ref.name] = module_ref.spec - spec = mod_and_mod_ref_to_proxy.get((bp.module, module_ref.name), module_ref.spec) - - if is_module_type(spec): - mod_and_mod_ref_to_proxy[bp.module, module_ref.name] = spec - continue + declared_spec[bp.name, module_ref.name] = module_ref.spec + spec = mod_and_mod_ref_to_proxy.get((bp.name, module_ref.name), module_ref.spec) result = _resolve_single_ref( - bp, module_ref, spec, blueprint, disabled_set, existing_modules + bp, + module_ref, + spec, + blueprint, + disabled_set, + existing_atoms, + existing_modules, ) if result is None: + mod_and_mod_ref_to_proxy.pop((bp.name, module_ref.name), None) continue if isinstance(result, DisabledModuleProxy): - disabled_ref_proxies[bp.module, module_ref.name] = result + disabled_ref_proxies[bp.name, module_ref.name] = result + mod_and_mod_ref_to_proxy.pop((bp.name, module_ref.name), None) else: - mod_and_mod_ref_to_proxy[bp.module, module_ref.name] = result + mod_and_mod_ref_to_proxy[bp.name, module_ref.name] = result - for (base_module, ref_name), target_module in mod_and_mod_ref_to_proxy.items(): - base_instance = module_coordinator.get_instance(base_module) - target_instance: Any = module_coordinator.get_instance(target_module) # type: ignore[arg-type] - async_methods = _async_methods_of_spec(declared_spec.get((base_module, ref_name))) + for (base_key, ref_name), target in mod_and_mod_ref_to_proxy.items(): + base_instance = module_coordinator.get_instance(base_key) + target_instance: Any = module_coordinator.get_instance(target) # type: ignore[arg-type] + async_methods = _async_methods_of_spec(declared_spec.get((base_key, ref_name))) if async_methods: target_instance = AsyncSpecProxy(target_instance, async_methods) setattr(base_instance, ref_name, target_instance) base_instance.set_module_ref(ref_name, target_instance) - module_coordinator._resolved_module_refs[base_module, ref_name] = cast( - "type[ModuleBase]", target_module - ) + if isinstance(target, str): + module_coordinator._resolved_module_refs[base_key, ref_name] = target + else: + target_keys = module_coordinator._instance_keys_of(cast("type[ModuleBase]", target)) + if target_keys: + module_coordinator._resolved_module_refs[base_key, ref_name] = target_keys[0] - for (base_module, ref_name), proxy in disabled_ref_proxies.items(): - base_instance = module_coordinator.get_instance(base_module) + for (base_key, ref_name), proxy in disabled_ref_proxies.items(): + base_instance = module_coordinator.get_instance(base_key) setattr(base_instance, ref_name, proxy) base_instance.set_module_ref(ref_name, cast("Any", proxy)) diff --git a/dimos/core/coordination/test_blueprints.py b/dimos/core/coordination/test_blueprints.py index d8ada76b20..6dd910fff3 100644 --- a/dimos/core/coordination/test_blueprints.py +++ b/dimos/core/coordination/test_blueprints.py @@ -32,6 +32,7 @@ ModuleRef, StreamRef, autoconnect, + namespace, ) from dimos.core.core import rpc from dimos.core.module import Module @@ -240,7 +241,7 @@ def test_blueprint_pickle_roundtrip() -> None: for name in ("transport_map", "global_config_overrides", "remapping_map"): assert isinstance(getattr(restored, name), MappingProxyType) assert dict(restored.global_config_overrides) == {"option1": True, "option2": 42} - assert restored.remapping_map[(ModuleA, "module_a")] is ModuleB + assert restored.remapping_map[(ModuleA.name, "module_a")] is ModuleB with pytest.raises(TypeError): restored.global_config_overrides["x"] = 1 @@ -251,3 +252,108 @@ def test_active_blueprints_filters_disabled() -> None: active_modules = {bp.module for bp in blueprint.active_blueprints} assert ModuleA not in active_modules assert ModuleB in active_modules + + +def test_namespace_prefixes_names_streams_and_frames() -> None: + blueprint = namespace("robot0", ModuleA.blueprint(), ModuleB.blueprint()) + + atom_a = next(a for a in blueprint.blueprints if a.module is ModuleA) + assert atom_a.name == "robot0/modulea" + assert atom_a.kwargs["instance_name"] == "robot0/modulea" + assert atom_a.kwargs["frame_id_prefix"] == "robot0" + # Every stream is remapped under the prefix. + assert blueprint.remapping_map[("robot0/modulea", "data1")] == "robot0/data1" + assert blueprint.remapping_map[("robot0/moduleb", "data3")] == "robot0/data3" + + +def test_namespace_expose_keeps_names_global() -> None: + blueprint = namespace("robot0", ModuleA.blueprint(), expose={"data1"}) + + assert ("robot0/modulea", "data1") not in blueprint.remapping_map + assert blueprint.remapping_map[("robot0/modulea", "data2")] == "robot0/data2" + + +def test_namespace_expose_typo_raises() -> None: + with pytest.raises(ValueError, match="data_typo"): + namespace("robot0", ModuleA.blueprint(), expose={"data_typo"}) + + +def test_namespace_invalid_prefix_raises() -> None: + with pytest.raises(ValueError, match="Invalid namespace prefix"): + namespace("a/b", ModuleA.blueprint()) + + +def test_namespace_nesting_composes() -> None: + blueprint = namespace("fleet", namespace("robot0", ModuleA.blueprint())) + + atom = blueprint.blueprints[0] + assert atom.name == "fleet/robot0/modulea" + assert atom.kwargs["frame_id_prefix"] == "fleet/robot0" + assert blueprint.remapping_map[("fleet/robot0/modulea", "data1")] == "fleet/robot0/data1" + + +def test_namespace_keeps_user_frame_id_prefix() -> None: + blueprint = namespace("robot0", ModuleA.blueprint(frame_id_prefix="custom")) + + assert blueprint.blueprints[0].kwargs["frame_id_prefix"] == "custom" + + +def test_namespace_allows_duplicate_module_classes() -> None: + blueprint = autoconnect( + namespace("robot0", ModuleA.blueprint(key1="a")), + namespace("robot1", ModuleA.blueprint(key1="b")), + ) + + atoms = [a for a in blueprint.blueprints if a.module is ModuleA] + assert {a.name for a in atoms} == {"robot0/modulea", "robot1/modulea"} + # Later-wins dedup still applies per instance name. + merged = autoconnect(blueprint, namespace("robot1", ModuleA.blueprint(key1="c"))) + robot1 = next(a for a in merged.blueprints if a.name == "robot1/modulea") + assert robot1.kwargs["key1"] == "c" + + +def test_namespace_config_keys_escaped() -> None: + blueprint = autoconnect( + namespace("robot0", ModuleA.blueprint()), + namespace("robot1", ModuleA.blueprint()), + ) + config = blueprint.config() + assert {"robot0_modulea", "robot1_modulea", "g"} <= set(config.model_fields.keys()) + + +def test_explicit_instance_name_sets_blueprint_identity() -> None: + blueprint = ModuleA.blueprint(instance_name="custom/modulea") + + atom = blueprint.blueprints[0] + assert atom.name == "custom/modulea" + assert atom.instance_name == "custom/modulea" + assert set(blueprint.config().model_fields) == {"custom_modulea", "g"} + + +def test_namespace_prefixes_pinned_transports() -> None: + blueprint = namespace( + "robot0", + autoconnect(ModuleA.blueprint(), ModuleB.blueprint()).transports( + {("data1", Data1): LCMTransport("/custom_topic", Data1)} + ), + ) + + transport = blueprint.transport_map[("robot0/data1", Data1)] + assert transport.topic.pattern == "/robot0/custom_topic" + assert ("data1", Data1) not in blueprint.transport_map + + +def test_namespace_remappings_by_instance_name() -> None: + # A remapping added after namespacing can target an instance by name. + blueprint = autoconnect( + namespace("robot0", ModuleA.blueprint()), + namespace("robot1", ModuleA.blueprint()), + ).remappings([("robot0/modulea", "data1", "special_data")]) + + assert blueprint.remapping_map[("robot0/modulea", "data1")] == "special_data" + # Targeting the class is ambiguous with two instances. + with pytest.raises(ValueError, match="multiple instances"): + autoconnect( + namespace("robot0", ModuleA.blueprint()), + namespace("robot1", ModuleA.blueprint()), + ).remappings([(ModuleA, "data1", "other")]) diff --git a/dimos/core/coordination/test_module_coordinator.py b/dimos/core/coordination/test_module_coordinator.py index c1baad17b2..7467f248ab 100644 --- a/dimos/core/coordination/test_module_coordinator.py +++ b/dimos/core/coordination/test_module_coordinator.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pickle from types import MappingProxyType from typing import Protocol @@ -282,9 +283,9 @@ def test_remapping() -> None: ] ) - # Verify remappings are stored correctly - assert (SourceModule, "color_image") in blueprint_set.remapping_map - assert blueprint_set.remapping_map[(SourceModule, "color_image")] == "remapped_data" + # Verify remappings are stored correctly, keyed by instance name. + assert (SourceModule.name, "color_image") in blueprint_set.remapping_map + assert blueprint_set.remapping_map[(SourceModule.name, "color_image")] == "remapped_data" # Verify that remapped names are used in name resolution all_names = _all_name_types(blueprint_set) @@ -575,6 +576,23 @@ def test_optional_module_ref_without_provider(build_coordinator) -> None: assert mod is not None +def test_build_with_explicit_instance_name(build_coordinator) -> None: + """A blueprint-supplied instance_name is used for deployment and wiring.""" + coordinator = build_coordinator( + autoconnect( + ModuleA.blueprint(instance_name="custom/modulea"), + ModuleB.blueprint(), + ) + ) + + module_a = coordinator.get_instance("custom/modulea") + module_b = coordinator.get_instance(ModuleB) + assert module_a is not None + assert module_b is not None + assert module_a.data1.transport.topic == module_b.data1.transport.topic + assert module_b.what_is_as_name() == "A, Module A" + + def test_load_blueprint_auto_scales_empty_pool(dynamic_coordinator) -> None: """A coordinator with 0 initial workers auto-adds workers on load_blueprint.""" dynamic_coordinator.load_blueprint(ModuleA.blueprint()) @@ -790,3 +808,43 @@ def test_list_module_names(dynamic_coordinator) -> None: dynamic_coordinator.load_module(ModuleA) dynamic_coordinator.load_module(ModuleC) assert set(dynamic_coordinator.list_module_names()) == {"ModuleA", "ModuleC"} + + +class NamedModule(Module): + @rpc + def whoami(self) -> str: + return self.config.instance_name or "default" + + +def test_deploy_two_instances_of_same_class(dynamic_coordinator) -> None: + """Two instances of one class get separate RPC topics and coordinator entries.""" + p0 = dynamic_coordinator.deploy(NamedModule, instance_name="robot0/namedmodule") + p1 = dynamic_coordinator.deploy(NamedModule, instance_name="robot1/namedmodule") + + assert dynamic_coordinator.get_instance("robot0/namedmodule") is p0 + assert dynamic_coordinator.get_instance("robot1/namedmodule") is p1 + with pytest.raises(ValueError, match="Multiple instances"): + dynamic_coordinator.get_instance(NamedModule) + + # Each proxy reaches its own instance over the instance-name RPC topic. + assert p0.remote_name == "robot0/namedmodule" + assert p0.whoami() == "robot0/namedmodule" + assert p1.whoami() == "robot1/namedmodule" + + assert set(dynamic_coordinator.list_module_names()) == { + "robot0/namedmodule", + "robot1/namedmodule", + } + descriptors = {d.rpc_name: d for d in dynamic_coordinator.list_modules()} + assert descriptors["robot0/namedmodule"].class_name == "NamedModule" + + +def test_rpc_client_pickle_preserves_remote_name(dynamic_coordinator) -> None: + """Proxies pickled into workers (set_module_ref) must keep the instance RPC name.""" + proxy = dynamic_coordinator.deploy(NamedModule, instance_name="robot0/namedmodule") + restored = pickle.loads(pickle.dumps(proxy)) + try: + assert restored.remote_name == "robot0/namedmodule" + assert restored.whoami() == "robot0/namedmodule" + finally: + restored.stop_rpc_client() diff --git a/dimos/core/coordination/test_namespace.py b/dimos/core/coordination/test_namespace.py new file mode 100644 index 0000000000..7ac4647917 --- /dev/null +++ b/dimos/core/coordination/test_namespace.py @@ -0,0 +1,192 @@ +# 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 types import MappingProxyType + +import pytest + +from dimos.core.coordination.blueprints import autoconnect, namespace +from dimos.core.coordination.module_coordinator import ModuleCoordinator +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out + +_BUILD_WITHOUT_RERUN = MappingProxyType( + { + "g": {"viewer": "none"}, + } +) + + +class Cloud: + pass + + +class Status: + pass + + +class Command: + pass + + +class SensorConfig(ModuleConfig): + sensitivity: float = 1.0 + + +class Sensor(Module): + config: SensorConfig + + pointcloud: Out[Cloud] + local_status: Out[Status] + cmd: In[Command] + + @rpc + def whoami(self) -> str: + return self.config.instance_name or "default" + + @rpc + def get_sensitivity(self) -> float: + return self.config.sensitivity + + @rpc + def get_frame_id(self) -> str: + return self.frame_id + + +class LocalMapper(Module): + local_status: In[Status] + + sensor: Sensor + + @rpc + def sensor_name(self) -> str: + return self.sensor.whoami() + + +class Aggregator(Module): + pointcloud: In[Cloud] + + +class FleetCommander(Module): + cmd: Out[Command] + + +def _fleet_blueprint(): + return autoconnect( + Aggregator.blueprint(), + FleetCommander.blueprint(), + *[ + namespace( + f"robot{i}", + Sensor.blueprint(), + LocalMapper.blueprint(), + expose={"pointcloud"}, + ) + for i in range(2) + ], + ).remappings([(FleetCommander, "cmd", "robot0/cmd")]) + + +@pytest.fixture +def fleet_coordinator(): + args = dict(_BUILD_WITHOUT_RERUN) + args["robot0_sensor"] = {"sensitivity": 2.0} + + coordinator = ModuleCoordinator.build(_fleet_blueprint(), args) + try: + yield coordinator + finally: + coordinator.stop() + + +def test_fleet_blueprint(fleet_coordinator: ModuleCoordinator): + coordinator = fleet_coordinator + sensor0 = coordinator.get_instance("robot0/sensor") + sensor1 = coordinator.get_instance("robot1/sensor") + mapper0 = coordinator.get_instance("robot0/localmapper") + mapper1 = coordinator.get_instance("robot1/localmapper") + aggregator = coordinator.get_instance(Aggregator) + commander = coordinator.get_instance(FleetCommander) + + # A class lookup with two instances is ambiguous. + with pytest.raises(ValueError, match="Multiple instances"): + coordinator.get_instance(Sensor) + + # RPC is served per instance, on the instance-name topic. + assert sensor0.whoami() == "robot0/sensor" + assert sensor1.whoami() == "robot1/sensor" + + # Namespaced streams get separate topics and exposed streams share one. + assert ( + sensor0.local_status.transport.topic + == mapper0.local_status.transport.topic + == "/robot0/local_status" + ) + assert ( + sensor1.local_status.transport.topic + == mapper1.local_status.transport.topic + == "/robot1/local_status" + ) + assert ( + sensor0.pointcloud.transport.topic + == sensor1.pointcloud.transport.topic + == aggregator.pointcloud.transport.topic + == "/pointcloud" + ) + + # Direct-class module refs resolve namespace-locally. + assert mapper0.sensor_name() == "robot0/sensor" + assert mapper1.sensor_name() == "robot1/sensor" + + # Per-instance config args reach only their instance. + assert sensor0.get_sensitivity() == 2.0 + assert sensor1.get_sensitivity() == 1.0 + + # TF frames carry the namespace. + assert sensor0.get_frame_id() == "robot0/Sensor" + + # Directed wiring: the shared commander drives only robot0. + assert commander.cmd.transport.topic == sensor0.cmd.transport.topic == "/robot0/cmd" + assert sensor1.cmd.transport.topic != sensor0.cmd.transport.topic + + +def test_fleet_blueprint_config_keys(): + config = _fleet_blueprint().config() + assert { + "robot0_sensor", + "robot1_sensor", + "robot0_localmapper", + "robot1_localmapper", + "aggregator", + "fleetcommander", + "g", + } == set(config.model_fields.keys()) + + +def test_load_blueprint_resolves_existing_provider_in_same_namespace(): + coordinator = ModuleCoordinator.build( + autoconnect( + namespace("robot0", Sensor.blueprint()), + namespace("robot1", Sensor.blueprint()), + ), + dict(_BUILD_WITHOUT_RERUN), + ) + try: + coordinator.load_blueprint(namespace("robot0", LocalMapper.blueprint())) + + mapper0 = coordinator.get_instance("robot0/localmapper") + assert mapper0.sensor_name() == "robot0/sensor" + finally: + coordinator.stop() diff --git a/dimos/core/coordination/worker_manager_python.py b/dimos/core/coordination/worker_manager_python.py index 4963db2dac..cd4df951b9 100644 --- a/dimos/core/coordination/worker_manager_python.py +++ b/dimos/core/coordination/worker_manager_python.py @@ -17,6 +17,7 @@ from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any +from dimos.core.coordination.blueprints import config_key from dimos.core.coordination.python_worker import PythonWorker from dimos.core.global_config import GlobalConfig from dimos.core.module import ModuleBase, ModuleSpec @@ -85,7 +86,7 @@ def deploy( self._ensure_capacity_for_dedicated([(module_class, global_config, kwargs)]) worker = self._select_worker(dedicated=module_class.dedicated_worker) actor = worker.deploy_module(module_class, global_config, kwargs=kwargs) - return RPCClient(actor, module_class) + return RPCClient(actor, module_class, kwargs.get("instance_name")) def deploy_fresh( self, @@ -111,7 +112,7 @@ def deploy_fresh( if module_class.dedicated_worker: worker.dedicated = True actor = worker.deploy_module(module_class, global_config, kwargs=kwargs) - return RPCClient(actor, module_class) + return RPCClient(actor, module_class, kwargs.get("instance_name")) def undeploy(self, proxy: ModuleProxyProtocol) -> None: """Undeploy a module and shut down its worker if it is now empty.""" @@ -161,7 +162,11 @@ def deploy_parallel( module_class, _, kwargs = specs[i] worker = self._select_worker(dedicated=module_class.dedicated_worker) worker.reserve_slot() - kwargs.update(blueprint_args.get(module_class.name, {})) + instance_key = kwargs.get("instance_name") or module_class.name + args = blueprint_args.get(config_key(instance_key), {}) + # instance_name is assigned by the blueprint; a user-supplied value + # would desync the module's RPC topic from the coordinator's proxy. + kwargs.update({k: v for k, v in args.items() if k != "instance_name"}) workers_by_index[i] = worker assignments = [(workers_by_index[i], specs[i]) for i in range(len(specs))] @@ -169,7 +174,9 @@ def deploy_parallel( def _deploy(item: tuple[PythonWorker, ModuleSpec]) -> ModuleProxyProtocol: worker, (module_class, global_config, kwargs) = item return RPCClient( - worker.deploy_module(module_class, global_config, kwargs), module_class + worker.deploy_module(module_class, global_config, kwargs), + module_class, + kwargs.get("instance_name"), ) try: diff --git a/dimos/core/global_config.py b/dimos/core/global_config.py index 328d6585ab..db627e6ef0 100644 --- a/dimos/core/global_config.py +++ b/dimos/core/global_config.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import platform import re from typing import Literal, TypeAlias @@ -125,5 +126,13 @@ def mujoco_camera_position_float(self) -> tuple[float, ...]: return (-0.906, 0.008, 1.101, 4.931, 89.749, -46.378) return tuple(_get_all_numbers(self.mujoco_camera_position)) + @property + def processed_robot_ips(self) -> tuple[str, ...]: + ips = [x.strip() for x in (self.robot_ips or "").split(",") if x.strip()] + is_running_tests = "PYTEST_CURRENT_TEST" in os.environ + if not ips and not is_running_tests: + raise ValueError("No robot IPs specified. Must have at least one IP.") + return tuple(ips) + global_config = GlobalConfig() diff --git a/dimos/core/introspection/blueprint/dot.py b/dimos/core/introspection/blueprint/dot.py index 3f5240690a..702dc03dc8 100644 --- a/dimos/core/introspection/blueprint/dot.py +++ b/dimos/core/introspection/blueprint/dot.py @@ -91,7 +91,7 @@ def render( module_classes[bp.module.__name__] = bp.module for conn in bp.streams: # Apply remapping - remapped_name = blueprint_set.remapping_map.get((bp.module, conn.name), conn.name) + remapped_name = blueprint_set.remapping_map.get((bp.name, conn.name), conn.name) key = (remapped_name, conn.type) if conn.direction == "out": producers[key].append(bp.module) # type: ignore[index] diff --git a/dimos/core/module.py b/dimos/core/module.py index 4f7aa9f964..8eb3568083 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -109,6 +109,10 @@ class ModuleConfig(BaseConfig): tf_transport: type[TFSpec] = Field(default_factory=tf_backend) # type: ignore[type-arg] frame_id_prefix: str | None = None frame_id: str | None = None + # Set by the coordinator when the same module class is deployed more than + # once (see BlueprintAtom.instance_name). Changes the RPC topic prefix + # from the class name to this name. + instance_name: str | None = None g: GlobalConfig = global_config @@ -157,7 +161,7 @@ def __init__(self, config_args: dict[str, Any]) -> None: # start() before serve_module_rpc(): Zenoh's subscribe needs an open # session (acquired in start()), whereas LCM tolerates either order. self.rpc.start() # type: ignore[attr-defined] - self.rpc.serve_module_rpc(self) + self.rpc.serve_module_rpc(self, name=self.config.instance_name) except ValueError: ... diff --git a/dimos/core/rpc_client.py b/dimos/core/rpc_client.py index cfdcf3dfa0..073c92e9a0 100644 --- a/dimos/core/rpc_client.py +++ b/dimos/core/rpc_client.py @@ -104,6 +104,7 @@ def __init__( self, actor_instance: Actor | None, actor_class: type[ModuleBase], + remote_name: str | None = None, *, rpc: RPCSpec | None = None, ) -> None: @@ -115,7 +116,7 @@ def __init__( self.rpc = rpc self._owns_rpc = False self.actor_class = actor_class - self.remote_name = actor_class.__name__ + self.remote_name = remote_name or actor_class.__name__ self.actor_instance = actor_instance self.rpcs = actor_class.rpcs.keys() self._unsub_fns: list = [] # type: ignore[type-arg] @@ -139,10 +140,12 @@ def stop_rpc_client(self) -> None: self.rpc = None # type: ignore[assignment] def __reduce__(self): # type: ignore[no-untyped-def] - # Return the class and the arguments needed to reconstruct the object + # Return the class and the arguments needed to reconstruct the object. + # remote_name must be included or proxies pickled into workers would + # fall back to class-name RPC topics. return ( self.__class__, - (self.actor_instance, self.actor_class), + (self.actor_instance, self.actor_class, self.remote_name), ) # passthrough diff --git a/dimos/porcelain/local_module_source.py b/dimos/porcelain/local_module_source.py index 8dea71154f..1a4e56d751 100644 --- a/dimos/porcelain/local_module_source.py +++ b/dimos/porcelain/local_module_source.py @@ -42,9 +42,23 @@ def list_module_names(self) -> list[str]: return self._coordinator.list_module_names() def get_module(self, name: str) -> Any: - for cls, proxy in self._coordinator._deployed_modules.items(): + if name in self._coordinator._deployed_modules: + return self._coordinator._deployed_modules[name] + + matches: list[tuple[str, Any]] = [] + for instance_key, proxy in self._coordinator._deployed_modules.items(): + cls = self._coordinator._instance_classes[instance_key] if cls.__name__ == name: - return proxy + matches.append((instance_key, proxy)) + + if len(matches) == 1: + return matches[0][1] + if len(matches) > 1: + instance_names = ", ".join(sorted(instance_key for instance_key, _ in matches)) + raise ValueError( + f"Multiple instances of {name!r} are deployed " + f"({instance_names}); use the instance name." + ) raise KeyError(name) def invalidate(self, name: str) -> None: diff --git a/dimos/porcelain/remote_module_source.py b/dimos/porcelain/remote_module_source.py index e87aec2284..710fe14ce9 100644 --- a/dimos/porcelain/remote_module_source.py +++ b/dimos/porcelain/remote_module_source.py @@ -75,7 +75,8 @@ def __init__(self, *, timeout: float = 5.0) -> None: def _refresh_descriptors(self) -> dict[str, ModuleDescriptor]: descriptors = self._coord.call("list_modules") - self._descriptors = {d.class_name: d for d in descriptors} + # rpc_name is empty when talking to an older daemon. + self._descriptors = {(d.rpc_name or d.class_name): d for d in descriptors} return self._descriptors def _get_descriptor(self, name: str) -> ModuleDescriptor: @@ -103,7 +104,7 @@ def get_module(self, name: str) -> Any: try: module_path, class_name = descriptor.qualified_path.rsplit(".", 1) cls = getattr(importlib.import_module(module_path), class_name) - proxy = RPCClient(None, cls, rpc=self._coord.rpc) + proxy = RPCClient(None, cls, descriptor.rpc_name or None, rpc=self._coord.rpc) except (ImportError, AttributeError): proxy = _RemoteProxy(self._coord.rpc, name, set(descriptor.rpc_names)) self._cache[name] = proxy diff --git a/dimos/porcelain/test_local_module_source.py b/dimos/porcelain/test_local_module_source.py index 1b10f2a460..47ea2c1316 100644 --- a/dimos/porcelain/test_local_module_source.py +++ b/dimos/porcelain/test_local_module_source.py @@ -16,9 +16,31 @@ import pytest +from dimos.core.coordination.module_coordinator import ModuleCoordinator +from dimos.core.core import rpc +from dimos.core.global_config import GlobalConfig +from dimos.core.module import Module from dimos.porcelain.local_module_source import LocalModuleSource +class NamedLocalModule(Module): + @rpc + def ping_name(self) -> str: + return self.config.instance_name or "default" + + +@pytest.fixture +def multi_instance_source(): + coordinator = ModuleCoordinator(g=GlobalConfig(n_workers=0, viewer="none")) + coordinator.start() + try: + coordinator.deploy(NamedLocalModule, instance_name="robot0/namedlocalmodule") + coordinator.deploy(NamedLocalModule, instance_name="robot1/namedlocalmodule") + yield LocalModuleSource(coordinator) + finally: + coordinator.stop() + + def test_is_not_remote(running_app): assert isinstance(running_app._source, LocalModuleSource) assert running_app._source.is_remote is False @@ -46,6 +68,14 @@ def test_get_module_unknown_raises(running_app): running_app._source.get_module("NonexistentModule") +def test_get_module_class_name_raises_when_ambiguous(multi_instance_source): + with pytest.raises(ValueError, match="Multiple instances"): + multi_instance_source.get_module("NamedLocalModule") + + module = multi_instance_source.get_module("robot1/namedlocalmodule") + assert module.ping_name() == "robot1/namedlocalmodule" + + def test_invalidate_is_noop(running_app): # Coordinator owns the proxy; the source has no per-call cache. running_app._source.invalidate("StressTestModule") diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 61220a0fae..37a8e64e31 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -120,6 +120,7 @@ "unitree-go2-markers": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_markers", "unitree-go2-memory": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_memory", "unitree-go2-mid360-record": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_mid360_record:unitree_go2_mid360_record", + "unitree-go2-multi": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_multi:unitree_go2_multi", "unitree-go2-nav-3d": "dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_nav_3d:unitree_go2_nav_3d", "unitree-go2-relocalization": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_relocalization", "unitree-go2-ros": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2_ros:unitree_go2_ros", diff --git a/dimos/robot/cli/dimos.py b/dimos/robot/cli/dimos.py index d7615a2795..a8e0532279 100644 --- a/dimos/robot/cli/dimos.py +++ b/dimos/robot/cli/dimos.py @@ -163,9 +163,12 @@ def arg_help( module: str = "", _atom: BlueprintAtom | None = None, ) -> str: + # Imported here for performance reasons. + from dimos.core.coordination.blueprints import config_key + output = "" for k, info in config.model_fields.items(): - if k == "g": + if k in ("g", "instance_name"): continue t = info.annotation if isinstance(t, types.GenericAlias): @@ -180,7 +183,7 @@ def arg_help( if inspect.isclass(t) and issubclass(t, BaseModel): output += f"{indent}{module}{k}:\n" # Find blueprint atom - bp = next(bp for bp in blueprint.blueprints if bp.module.name == k) + bp = next(bp for bp in blueprint.blueprints if config_key(bp.name) == k) output += arg_help( t, blueprint, indent=indent + " ", module=module + k + ".", _atom=bp ) @@ -212,7 +215,9 @@ def load_config_args(config: type[BaseModel], args: Iterable[str], path: Path) - for arg in args: k, _, v = arg.partition("=") - parts = k.split(".") + # Accept namespaced instance names in both forms: robot0/sensor.ip + # and robot0_sensor.ip (config keys escape "/" to "_"). + parts = [p.replace("/", "_") for p in k.split(".")] d = kwargs for p in parts[:-1]: d = d.setdefault(p, {}) diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_multi.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_multi.py new file mode 100644 index 0000000000..afea11bb17 --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_multi.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 + +# 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. + +"""Blueprint for multiple independent Go2 robots on one coordinator. + +Each robot gets its own namespaced GO2Connection: topics like `/robot0/lidar`, +TF frames like `robot0/...`, and RPC like `robot0/go2connection/move`. + +Configure a single robot with `-o robot0/go2connection.=`. + +Usage: + ROBOT_IPS=10.0.0.102,10.0.0.209 dimos run unitree-go2-multi +""" + +from dimos.core.coordination.blueprints import autoconnect, namespace +from dimos.core.global_config import global_config +from dimos.robot.unitree.go2.connection import GO2Connection + +_ips = global_config.processed_robot_ips + +unitree_go2_multi = autoconnect( + *[namespace(f"robot{i}", GO2Connection.blueprint(ip=ip)) for i, ip in enumerate(_ips)], +).global_config(n_workers=max(2, 2 * len(_ips)), robot_model="unitree_go2") diff --git a/docs/usage/blueprints.md b/docs/usage/blueprints.md index bceb356cd7..3f8b22a2bb 100644 --- a/docs/usage/blueprints.md +++ b/docs/usage/blueprints.md @@ -209,6 +209,79 @@ blueprint.remappings([ }) ``` +## Multi-robot blueprints (namespaces) + +You can use namespaces to control several robot instances. + +Use `namespace(prefix, *blueprints, expose=...)` to isolate a set of modules +under a name prefix so several copies can coexist. Because blueprints are plain +Python values, a variable-size (and mixed-type) fleet is just a loop: + +```python session=blueprint-ns +from dimos.core.coordination.blueprints import autoconnect, namespace +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out + +class SensorConfig(ModuleConfig): + ip: str = "" + +class Sensor(Module): + config: SensorConfig + pointcloud: Out[str] + +class AggregateMapper(Module): + pointcloud: In[str] + +robot_ips = ["10.0.0.1", "10.0.0.2"] + +fleet = autoconnect( + AggregateMapper.blueprint(), # shared: one instance for the whole fleet + *[ + namespace(f"robot{i}", Sensor.blueprint(ip=ip), expose={"pointcloud"}) + for i, ip in enumerate(robot_ips) + ], +) +``` + +Inside a namespace everything is prefixed: + +* instance names (`robot0/go2connection`), +* stream names and topics (`/robot0/lidar`), +* TF frames (`frame_id_prefix`, unless you set one yourself), +* and RPC topics (`robot0/go2connection/move`). + +Prefixed streams only connect within their namespace. + +Stream names listed in `expose` are left unprefixed, so they connect globally. +That is how data crosses the boundary. + +In the above example, both `Sensor` modules can send `pointcloud` data to the `AggregateMapper` module. + +If you want to manually redirect to a namespaced name, you can use the namespaced prefix: `.remappings([(FleetPlanner, "cmd_r0", "robot0/cmd_vel")])` + +Module references (Specs or direct classes) resolve within the consumer's +namespace first, then enclosing namespaces, then globally, so per-robot +consumers bind to their own robot's providers. + +Configuration addresses instances with `/` escaped to `_`: +`-o robot0/go2connection.ip=10.0.0.5` (or `-o robot0_go2connection.ip=...`), +env `ROBOT0_GO2CONNECTION__IP=...`. `.remappings` accepts an instance name +string wherever it accepts a module class. + +The downside of this is that the number of modules is fixed at blueprint creation. But it can be easily overcome by using a global variable. For example you can invoke with this: + +```bash +ROBOT_IPS=10.0.0.1,10.0.0.2 dimos run blueprint-name +``` + +and construct `robot_ips` from the environment variable: + +```python skip +robot_ips = (global_config.robot_ips or "").split(",") +``` + +For convenience `global_config.processed_robot_ips` is available which automatically splits the string and errors if no IPs are present. + ## Overriding global configuration. Each module includes the global config available as `self.config.g`. E.g.: diff --git a/docs/usage/cli.md b/docs/usage/cli.md index dc4de5d3c9..c2823ff3c6 100644 --- a/docs/usage/cli.md +++ b/docs/usage/cli.md @@ -48,13 +48,13 @@ dimos [GLOBAL OPTIONS] COMMAND [ARGS] Values cascade (later overrides earlier): -1. `GlobalConfig` default → `simulation = False` -2. `.env` file → `DIMOS_SIMULATION=true` -3. Environment variable → `export DIMOS_SIMULATION=true` -4. Blueprint definition → `.global_config(simulation=True)` +1. `GlobalConfig` default → `simulation = ""` +2. `.env` file → `SIMULATION=mujoco` +3. Environment variable → `export SIMULATION=mujoco` +4. Blueprint definition → `.global_config(simulation="mujoco")` 5. CLI flag → `dimos --simulation run ...` -Environment variables and `.env` values must be prefixed with `DIMOS_`. +Environment variables and `.env` values use the field name in uppercase, for example `ROBOT_IPS`. --- @@ -201,8 +201,6 @@ Print resolved GlobalConfig values and their sources. dimos show-config ``` ---- - ## Agent & MCP Commands ### `dimos agent-send` @@ -281,8 +279,6 @@ List deployed modules and their skills. dimos mcp modules ``` ---- - ## Standalone Tools These are installed as separate entry points and can be run directly without the `dimos` prefix. @@ -335,12 +331,10 @@ rerun-bridge Also available as `dimos rerun-bridge`. ---- - ## File Locations | Path | Contents | |------|----------| | `~/.local/state/dimos/runs/.json` | Run registry (PID, blueprint, args, ports). Used by `status`/`stop`/`restart`. Cleaned up when processes exit. | | `~/.local/state/dimos/logs//main.jsonl` | Structured logs (main process + all workers) | -| `.env` | Local config overrides (`DIMOS_ROBOT_IP=192.168.123.161`) | +| `.env` | Local config overrides (`ROBOT_IP=192.168.123.161`) |