From e1a9aa9f5ee08d4abd1b9a9397bbd7d1c056c63c Mon Sep 17 00:00:00 2001 From: Paul Nechifor Date: Sat, 24 Oct 2026 21:42:04 +0300 Subject: [PATCH 1/3] refactor(rpc): replace pickle with CBOR in RPC and the worker pipe ZenohRPC, LCMRPC and the forkserver worker pipe now serialize with a CBOR codec (dimos/protocol/encode/cbor.py) instead of pickle. Tuples, NamedTuples, sets, numpy arrays and scalars, enums, dataclasses, pydantic models, Path, message types (class ref plus LCM bytes), classes, generic aliases, module-level functions and functools.partial round-trip as their original types. Anything else raises TypeError at the call site naming the type and its path in the payload; classes opt in with cbor_encode()/cbor_decode() or register_codec(). Transports, RPCClient and its proxies, Actor, streams, Blueprint, system configurators, LCM and Zenoh services, the relay-bridge metaclass, lazy asset paths and grasp candidates gain the hooks; their pickle hooks are removed. StreamDescriptor carries the real type and transport, HostFragment encodes its payload with the codec, and the pickled pubsub transports keep pickle. Decode failures surface at the caller instead of timing out, a worker that cannot decode a request replies with an error, and the worker encodes before taking its lock. Also fixes start/stop on DemoCalculatorSkill and OsmSkill (missing @rpc, so pickle ran them on a parent-side copy) and gives dimos Header its own lcm_decode. --- AGENTS.md | 2 + dimos/agents/mcp/mcp_client.py | 3 - dimos/agents/skills/demo_calculator_skill.py | 3 + dimos/agents/skills/osm.py | 3 + dimos/agents/voice_input.py | 3 - .../connection/test_connection_module.py | 4 +- dimos/control/contract/keys.py | 7 + dimos/control/contract/test_keys.py | 7 +- dimos/control/contract/test_limits.py | 4 +- dimos/control/contract/test_sequence.py | 7 +- dimos/core/baked_host.py | 4 +- dimos/core/coordination/blueprints.py | 67 ++- dimos/core/coordination/module_coordinator.py | 10 +- dimos/core/coordination/python_worker.py | 76 ++- dimos/core/coordination/test_blueprints.py | 38 +- .../coordination/test_module_coordinator.py | 9 +- dimos/core/coordination/test_worker.py | 28 +- dimos/core/introspection/module/info.py | 7 +- dimos/core/module.py | 40 +- dimos/core/o3dpickle.py | 21 - dimos/core/rpc_client.py | 59 ++- dimos/core/stream.py | 62 ++- dimos/core/test_baked_host.py | 10 +- dimos/core/test_rpc_client.py | 12 +- dimos/core/test_stream.py | 21 +- dimos/core/test_transport.py | 21 +- dimos/core/test_zenoh_transport.py | 10 +- dimos/core/transport.py | 141 ++++-- dimos/experimental/robot/bosdyn/spot/rerun.py | 7 +- dimos/hosted/fragment.py | 8 +- dimos/hosted/test_fragment.py | 4 +- dimos/manipulation/planning/spec/models.py | 2 +- .../manipulation/planning/spec/test_config.py | 4 +- .../manipulation/planning/spec/validation.py | 4 +- dimos/manipulation/test_manipulation_spec.py | 4 +- dimos/manipulation/test_manipulation_unit.py | 8 +- dimos/mapping/experimental/mesh.py | 24 +- dimos/models/base.py | 11 +- .../msgs/manipulation_msgs/GraspCandidate.py | 7 + .../manipulation_msgs/GraspCandidateArray.py | 7 + .../manipulation_msgs/test_grasp_candidate.py | 15 + dimos/msgs/std_msgs/Header.py | 6 + dimos/navigation/global_planner/test_viz.py | 16 +- dimos/navigation/global_planner/viz.py | 2 +- dimos/porcelain/dimos.py | 2 +- dimos/porcelain/remote_module_source.py | 3 +- dimos/porcelain/test_remote_module_source.py | 5 +- dimos/protocol/encode/cbor.py | 420 +++++++++++++++++ dimos/protocol/encode/test_cbor.py | 432 ++++++++++++++++++ .../pubsub/benchmark/test_replay_bench.py | 8 +- .../pubsub/benchmark/tool_replay_bench.py | 8 +- dimos/protocol/pubsub/encoders.py | 22 + dimos/protocol/pubsub/impl/lcmpubsub.py | 7 + .../pubsub/impl/webrtc/providers/spec.py | 6 +- .../pubsub/impl/webrtc/test_transport.py | 6 +- dimos/protocol/pubsub/impl/zenohpubsub.py | 21 - dimos/protocol/pubsub/test_encoder.py | 32 ++ dimos/protocol/rpc/pubsubrpc.py | 44 +- dimos/protocol/rpc/rpc_utils.py | 17 +- dimos/protocol/rpc/spec.py | 4 +- dimos/protocol/rpc/test_spec.py | 109 +++++ dimos/protocol/rpc/zenohrpc.py | 48 +- dimos/protocol/service/call_pool.py | 7 - dimos/protocol/service/lcmservice.py | 36 +- .../service/system_configurator/base.py | 11 + dimos/protocol/service/test_call_pool.py | 38 -- dimos/protocol/service/test_lcmservice.py | 46 -- dimos/protocol/service/test_zenohservice.py | 34 -- dimos/protocol/service/zenohservice.py | 27 +- dimos/robot/assets/source.py | 20 +- dimos/robot/assets/test_model.py | 6 +- dimos/robot/assets/test_source.py | 6 +- .../robot/test_blueprints_cross_processes.py | 43 ++ .../go2/blueprints/basic/unitree_go2_basic.py | 2 +- .../basic/unitree_go2_joystick_record.py | 5 +- dimos/utils/data.py | 8 + .../visualization/rerun/test_bridge_config.py | 4 +- .../visualization/rerun/test_scene_package.py | 4 +- dimos/web/cockpit.py | 7 + dimos/web/codecs.py | 18 +- dimos/web/relay_bridge/channels.py | 2 +- dimos/web/relay_bridge/dynamic.py | 59 +-- dimos/web/relay_bridge/test_dynamic.py | 33 +- .../test_relay_bridge_authoring.py | 6 +- dimos/web/test_cockpit.py | 28 +- dimos/web/test_codecs.py | 8 +- docs/coding-agents/code-quality-rules.md | 2 +- docs/usage/distributed/index.md | 6 +- docs/usage/python-api.md | 4 +- docs/usage/transports/index.md | 1 + pyproject.toml | 1 + uv.lock | 34 ++ 92 files changed, 1834 insertions(+), 674 deletions(-) create mode 100644 dimos/protocol/encode/cbor.py create mode 100644 dimos/protocol/encode/test_cbor.py delete mode 100644 dimos/protocol/service/test_call_pool.py create mode 100644 dimos/robot/test_blueprints_cross_processes.py diff --git a/AGENTS.md b/AGENTS.md index f45f0f3ec3d..dc31ad28dc4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -321,6 +321,8 @@ class MySkillContainer(Module): If multiple modules match the spec, use `.remappings()` to resolve. Source: `dimos/spec/utils.py`, `dimos/core/coordination/blueprints.py`. +RPC arguments, results and worker-pipe messages are CBOR (`dimos/protocol/encode/cbor.py`), not pickle. JSON-like values, tuples, sets, numpy arrays, enums, dataclasses, pydantic models, `Path`, message types with `lcm_encode`, classes and module-level functions cross as themselves. Anything else raises `TypeError` at the call site; give such a class `cbor_encode()` and a `cbor_decode()` classmethod. + ### Adding a New Skill 1. Pick the right container (robot-specific or `dimos/agents/skills/`). diff --git a/dimos/agents/mcp/mcp_client.py b/dimos/agents/mcp/mcp_client.py index dd53a2763ac..14d41e03cc0 100644 --- a/dimos/agents/mcp/mcp_client.py +++ b/dimos/agents/mcp/mcp_client.py @@ -137,9 +137,6 @@ def __init__(self, **kwargs: Any) -> None: self._mcp_url: str | None = None Thread(target=_preload_agent_deps, name="McpClient-preload", daemon=True).start() - def __reduce__(self) -> Any: - return (self.__class__, (), {}) - def _server_url(self) -> str: """The MCP endpoint: configured, else the blueprint's McpServer, else the default.""" if self._mcp_url is None: diff --git a/dimos/agents/skills/demo_calculator_skill.py b/dimos/agents/skills/demo_calculator_skill.py index 6c0605bb7c4..13b3f98c960 100644 --- a/dimos/agents/skills/demo_calculator_skill.py +++ b/dimos/agents/skills/demo_calculator_skill.py @@ -13,13 +13,16 @@ # limitations under the License. from dimos.agents.annotation import skill +from dimos.core.core import rpc from dimos.core.module import Module class DemoCalculatorSkill(Module): + @rpc def start(self) -> None: super().start() + @rpc def stop(self) -> None: super().stop() diff --git a/dimos/agents/skills/osm.py b/dimos/agents/skills/osm.py index 888c16b19eb..57aa21c9ac8 100644 --- a/dimos/agents/skills/osm.py +++ b/dimos/agents/skills/osm.py @@ -18,6 +18,7 @@ from reactivex.disposable import Disposable from dimos.agents.annotation import skill +from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In from dimos.mapping.models import LatLon @@ -40,6 +41,7 @@ def __init__(self, **kwargs: Any) -> None: self._latest_location = None self._current_location_map = CurrentLocationMap(QwenVlModel()) + @rpc def start(self) -> None: super().start() if hasattr(self.gps_location, "subscribe"): @@ -49,6 +51,7 @@ def start(self) -> None: "OsmSkill: gps_location stream does not support direct subscribe (RemoteIn)" ) + @rpc def stop(self) -> None: super().stop() diff --git a/dimos/agents/voice_input.py b/dimos/agents/voice_input.py index b084b63daa5..fcac32bfdfc 100644 --- a/dimos/agents/voice_input.py +++ b/dimos/agents/voice_input.py @@ -92,9 +92,6 @@ def __init__(self, **kwargs: Any) -> None: # and start() waits only for whatever is left. self._pipeline = run_in_thread(whisper_pipeline, "whisper-pipeline") - def __reduce__(self) -> Any: - return (self.__class__, (), {}) - @rpc def start(self) -> None: super().start() diff --git a/dimos/control/connection/test_connection_module.py b/dimos/control/connection/test_connection_module.py index 4830ade106d..e4e12d046e9 100644 --- a/dimos/control/connection/test_connection_module.py +++ b/dimos/control/connection/test_connection_module.py @@ -22,7 +22,6 @@ from collections.abc import Callable, Iterator import math -import pickle import threading import time from typing import Any @@ -33,6 +32,7 @@ from dimos.control.connection.connection_module import ConnectionStatus from dimos.control.contract.limits import Limits from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.protocol.encode.cbor import decode, encode JOINT1 = "mock/joint1/position" DEADMAN_S = 0.05 @@ -300,7 +300,7 @@ def test_status_is_an_rpc_that_crosses_processes(rig: Callable[..., Rig]) -> Non r = rig() status = r.module.status() assert isinstance(status, ConnectionStatus) - assert pickle.loads(pickle.dumps(status)) == status + assert decode(encode(status)) == status assert len(status.session_id) == 32 assert status.connected assert status.last_command_time is None diff --git a/dimos/control/contract/keys.py b/dimos/control/contract/keys.py index 8b71766e58e..d4fcd9f2793 100644 --- a/dimos/control/contract/keys.py +++ b/dimos/control/contract/keys.py @@ -133,6 +133,13 @@ def __new__(cls, value: str) -> Key: key._source, key._resource, key._interface = parts return key + def cbor_encode(self) -> str: + return str(self) + + @classmethod + def cbor_decode(cls, value: str) -> Key: + return cls(value) + @classmethod def of(cls, source: str, resource: str, interface: str) -> Key: """Build a name from its three parts. diff --git a/dimos/control/contract/test_keys.py b/dimos/control/contract/test_keys.py index 5bc599e8da4..c9f97da5fce 100644 --- a/dimos/control/contract/test_keys.py +++ b/dimos/control/contract/test_keys.py @@ -17,11 +17,10 @@ from __future__ import annotations -import pickle - import pytest from dimos.control.contract.keys import Interface, Key, is_valid_segment +from dimos.protocol.encode.cbor import decode, encode @pytest.mark.parametrize("segment", ["arm", "joint1", "left_hip_pitch", "A", "_", "9", "a_1_B"]) @@ -107,9 +106,9 @@ def test_error_quotes_the_key() -> None: Key("arm/joint1") -def test_keys_pickle_with_their_parts() -> None: +def test_keys_cross_processes_with_their_parts() -> None: """Descriptions cross a process boundary, so the parts must survive.""" - restored = pickle.loads(pickle.dumps(Key.of("arm", "joint1", Interface.POSITION))) + restored = decode(encode(Key.of("arm", "joint1", Interface.POSITION))) assert restored == "arm/joint1/position" assert restored.source == "arm" diff --git a/dimos/control/contract/test_limits.py b/dimos/control/contract/test_limits.py index 7566ebfc411..d0da63c17d4 100644 --- a/dimos/control/contract/test_limits.py +++ b/dimos/control/contract/test_limits.py @@ -24,7 +24,6 @@ from collections.abc import Sequence import math from pathlib import Path -import pickle from typing import Any import pytest @@ -40,6 +39,7 @@ ) from dimos.msgs.geometry_msgs.Twist import Twist from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.protocol.encode.cbor import decode, encode from dimos.robot.assets.model import JointDescription, RobotModel JOINTS = ("j1", "j2", "j3") @@ -264,7 +264,7 @@ def test_a_config_that_commands_nothing_is_refused(model: RobotModel) -> None: def test_the_table_crosses_processes(model: RobotModel) -> None: table = limits(config(model=model, pd=True, base_range=RANGE)) - assert pickle.loads(pickle.dumps(table)) == table + assert decode(encode(table)) == table def test_limits_refuse_unless_told_to_clamp() -> None: diff --git a/dimos/control/contract/test_sequence.py b/dimos/control/contract/test_sequence.py index bec2ffe850f..c37df81af75 100644 --- a/dimos/control/contract/test_sequence.py +++ b/dimos/control/contract/test_sequence.py @@ -16,11 +16,10 @@ from __future__ import annotations -import pickle - import pytest from dimos.control.contract.sequence import Freshness, is_newer +from dimos.protocol.encode.cbor import decode, encode def test_anything_is_newer_than_nothing() -> None: @@ -78,9 +77,9 @@ def test_marking_again_refreshes() -> None: assert not fresh.is_stale(now=20.0, timeout_s=1.0) -def test_freshness_pickles() -> None: +def test_freshness_crosses_processes() -> None: """Freshness travels with a source snapshot over RPC.""" fresh = Freshness() fresh.mark(42.0) - assert pickle.loads(pickle.dumps(fresh)).last_receipt == 42.0 + assert decode(encode(fresh)).last_receipt == 42.0 diff --git a/dimos/core/baked_host.py b/dimos/core/baked_host.py index b0226f07014..7bb456f8a32 100644 --- a/dimos/core/baked_host.py +++ b/dimos/core/baked_host.py @@ -128,7 +128,7 @@ def baked_host( ) -> type[BakedHost]: """Build the NativeModule subclass that drives a baked host binary. - Assign the result to `name` at module level: deploying to a worker pickles + Assign the result to `name` at module level: deploying to a worker sends the class by that path. """ if not members: @@ -160,6 +160,6 @@ def baked_host( "_remaps": remaps, } host_cls = type(name, (BakedHost,), namespace) - # Nothing names the config class, so put it where pickle will look for it. + # Nothing names the config class, so put it where the codec will look for it. setattr(sys.modules[caller], config_cls.__name__, config_cls) return host_cls diff --git a/dimos/core/coordination/blueprints.py b/dimos/core/coordination/blueprints.py index dbdaca766b3..58de4d95a10 100644 --- a/dimos/core/coordination/blueprints.py +++ b/dimos/core/coordination/blueprints.py @@ -13,14 +13,23 @@ # limitations under the License. from collections.abc import Callable, Iterable, Mapping, Sequence -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, field, fields, replace from functools import cached_property, reduce import operator import re import sys import types as types_mod from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Literal, Union, get_args, get_origin, get_type_hints +from typing import ( + TYPE_CHECKING, + Any, + Literal, + Union, + cast, + get_args, + get_origin, + get_type_hints, +) from pydantic import BaseModel @@ -29,6 +38,7 @@ from dimos.core.module import ModuleBase, is_module_type from dimos.core.stream import IO, In, Out, Transport +from dimos.core.transport import PubSubTransport from dimos.spec.utils import Spec, is_spec from dimos.utils.logging_config import setup_logger @@ -57,8 +67,12 @@ def _noop(*_args: Any, **_kwargs: Any) -> None: return _noop - def __reduce__(self) -> tuple[type, tuple[str]]: - return (DisabledModuleProxy, (self._spec_name,)) + def cbor_encode(self) -> str: + return cast("str", object.__getattribute__(self, "_spec_name")) + + @classmethod + def cbor_decode(cls, spec_name: str) -> "DisabledModuleProxy": + return cls(spec_name) def __repr__(self) -> str: return f"" @@ -199,7 +213,7 @@ def build(self, config: BaseModel | None = None) -> Transport[Any]: return self.cls(*self.args, **self.kwargs, **extra) -# These fields cannot be pickled. +# MappingProxyType fields: sent as plain dicts, rebuilt on decode. _PROXY_FIELDS = ("transport_map", "global_config_overrides", "remapping_map") @@ -234,17 +248,17 @@ def __post_init__(self) -> None: if unknown: raise ValueError(f"Hosted placement references unknown modules: {', '.join(unknown)}") - def __getstate__(self) -> dict[str, Any]: - state = self.__dict__.copy() - state.pop("active_blueprints", None) # recomputable cached_property + def cbor_encode(self) -> dict[str, Any]: + state = {f.name: getattr(self, f.name) for f in fields(self)} for name in _PROXY_FIELDS: state[name] = dict(state[name]) return state - def __setstate__(self, state: dict[str, Any]) -> None: + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> "Blueprint": for name in _PROXY_FIELDS: - state[name] = MappingProxyType(state[name]) - self.__dict__.update(state) + data[name] = MappingProxyType(data[name]) + return cls(**data) @classmethod def create(cls, module: type[ModuleBase], **kwargs: Any) -> "Blueprint": @@ -522,16 +536,21 @@ def _reprefix_transport( ) -> TransportSpec | Transport[Any]: """Clone a pinned transport with its topic moved under the namespace prefix.""" if isinstance(transport, TransportSpec): - cls, args = transport.cls, transport.args - else: - 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}" - if isinstance(transport, TransportSpec): - return replace(transport, args=(new_topic, *args[1:])) - return cls(new_topic, *args[1:]) # type: ignore[no-any-return, call-arg] + args = transport.args + if not (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." + ) + return replace(transport, args=(_prefixed_topic(args[0], prefix), *args[1:])) + if isinstance(transport, PubSubTransport): + topic = transport.cbor_encode()["topic"] + return transport.with_topic(_prefixed_topic(topic, prefix)) + raise ValueError( + f"Cannot namespace pinned transport {transport!r}; pin it on the " + f"namespaced blueprint with an explicit topic instead." + ) + + +def _prefixed_topic(topic: str, prefix: str) -> str: + return f"/{prefix}{topic}" if topic.startswith("/") else f"{prefix}/{topic}" diff --git a/dimos/core/coordination/module_coordinator.py b/dimos/core/coordination/module_coordinator.py index d69cb5c4265..13ba57037e7 100644 --- a/dimos/core/coordination/module_coordinator.py +++ b/dimos/core/coordination/module_coordinator.py @@ -35,7 +35,6 @@ from dimos.core.global_config import GlobalConfig, global_config from dimos.core.introspection.module.info import StreamDescriptor from dimos.core.module import ModuleBase, ModuleSpec, is_module_type -from dimos.core.o3dpickle import register_picklers from dimos.core.resource import Resource from dimos.core.stream import Transport from dimos.core.transport import ( @@ -99,7 +98,6 @@ def __init__( self._shutdown_event = threading.Event() def start(self) -> None: - register_picklers() for m in self._managers.values(): m.start() self._started = True @@ -373,8 +371,10 @@ def observe_transports( def _send_on_system_modules(self) -> None: modules = list(self._deployed_modules.values()) for module in modules: - if hasattr(module, "on_system_modules"): - module.on_system_modules(modules) + # Only an @rpc counts: a plain method would be fetched over the worker + # pipe as a bound method, which cannot cross it. + if "on_system_modules" in getattr(module, "rpcs", ()): + cast("Any", module).on_system_modules(modules) def _connect_streams( self, blueprint: Blueprint, transports: Mapping[tuple[str, type], Transport[Any]] @@ -852,7 +852,7 @@ def _materialize_transports( WebRTC transports get a freshly constructed provider config from the resolved ``transports..*`` overrides; everything else builds from the spec as-is, then gets coerced to the active pubsub backend. Returns - ready-to-use instances pickled into module workers. + ready-to-use instances sent into module workers. """ materialized: dict[tuple[str, type], Transport[Any]] = {} for key, spec in blueprint.transport_map.items(): diff --git a/dimos/core/coordination/python_worker.py b/dimos/core/coordination/python_worker.py index b5c89c2c627..016a969f98f 100644 --- a/dimos/core/coordination/python_worker.py +++ b/dimos/core/coordination/python_worker.py @@ -36,6 +36,7 @@ WorkerResponse, ) from dimos.core.global_config import GlobalConfig, global_config +from dimos.protocol.encode.cbor import decode, encode from dimos.protocol.pubsub.impl.webrtc.providers.spec import shutdown_all_providers from dimos.utils.logging_config import setup_logger from dimos.utils.sequential_ids import SequentialIds @@ -46,6 +47,14 @@ logger = setup_logger() +def _send(conn: Connection, message: Any) -> None: + conn.send_bytes(encode(message)) + + +def _recv(conn: Connection) -> Any: + return decode(conn.recv_bytes()) + + class ActorFuture: """Mimics Dask's ActorFuture - wraps a result with .result() method.""" @@ -66,8 +75,12 @@ class MethodCallProxy: def __init__(self, actor: Actor) -> None: self._actor = actor - def __reduce__(self) -> tuple[type, tuple[Actor]]: - return (MethodCallProxy, (self._actor,)) + def cbor_encode(self) -> Actor: + return self._actor + + @classmethod + def cbor_decode(cls, actor: Actor) -> MethodCallProxy: + return cls(actor) def __getattr__(self, name: str) -> Any: # Don't intercept private/dunder attributes - they must follow normal lookup. @@ -105,25 +118,32 @@ def __init__( self._module_id = module_id self._lock = lock - def __reduce__(self) -> tuple[type, tuple[None, None, int, int, None]]: - """Exclude the connection, the lock and the module class when pickling. + def cbor_encode(self) -> dict[str, int]: + """Only the ids cross a process boundary. - The class would make every receiving process import the module and - its dependencies; nothing reads it off an actor that crossed a - process boundary. + The connection and the lock cannot, and the class would make every + receiving process import the module and its dependencies; nothing reads + it off an actor that crossed a process boundary. """ - return (Actor, (None, None, self._worker_id, self._module_id, None)) + return {"worker_id": self._worker_id, "module_id": self._module_id} + + @classmethod + def cbor_decode(cls, data: dict[str, int]) -> Actor: + return cls(None, None, data["worker_id"], data["module_id"], None) def _send_request_to_worker(self, request: WorkerRequest) -> Any: if self._conn is None: raise RuntimeError("Actor connection not available - cannot send requests") + # Encode before taking the lock: an unencodable argument raises to the + # caller without holding up the other modules on this worker. + payload = encode(request) if self._lock is not None: with self._lock: - self._conn.send(request) - response: WorkerResponse = self._conn.recv() + self._conn.send_bytes(payload) + response: WorkerResponse = _recv(self._conn) else: - self._conn.send(request) - response = self._conn.recv() + self._conn.send_bytes(payload) + response = _recv(self._conn) if response.error: if "AttributeError" in response.error: # TODO: better error handling raise AttributeError(response.error) @@ -202,7 +222,8 @@ def worker_id(self) -> int: def module_names(self) -> list[str]: names = [] for actor in self._modules.values(): - # Parent-side actors always carry their class; only pickled copies drop it. + # Parent-side actors always carry their class; only copies that crossed a + # process boundary drop it. assert actor._cls is not None names.append(actor._cls.__name__) return names @@ -239,9 +260,10 @@ def deploy_module( request = DeployModuleRequest(module_id=module_id, module_class=module_class, kwargs=kwargs) try: + payload = encode(request) with self._lock: - self._conn.send(request) - response: WorkerResponse = self._conn.recv() + self._conn.send_bytes(payload) + response: WorkerResponse = _recv(self._conn) if response.error: raise RuntimeError(f"Failed to deploy module: {response.error}") @@ -266,8 +288,8 @@ def undeploy_module(self, module_id: int) -> None: raise RuntimeError("Worker process not started") with self._lock: - self._conn.send(UndeployModuleRequest(module_id=module_id)) - response: WorkerResponse = self._conn.recv() + _send(self._conn, UndeployModuleRequest(module_id=module_id)) + response: WorkerResponse = _recv(self._conn) if response.error: raise RuntimeError(f"Failed to undeploy module: {response.error}") @@ -279,8 +301,8 @@ def suppress_console(self) -> None: return try: with self._lock: - self._conn.send(SuppressConsoleRequest()) - self._conn.recv() + _send(self._conn, SuppressConsoleRequest()) + _recv(self._conn) except (BrokenPipeError, EOFError, ConnectionResetError): pass @@ -288,9 +310,9 @@ def shutdown(self) -> None: if self._conn is not None: try: with self._lock: - self._conn.send(ShutdownRequest()) + _send(self._conn, ShutdownRequest()) if self._conn.poll(timeout=5): - self._conn.recv() + _recv(self._conn) else: logger.warning( "Worker did not respond to shutdown within 5s, closing pipe.", @@ -476,11 +498,13 @@ def _worker_loop(conn: Connection, state: _WorkerState) -> None: try: if not conn.poll(timeout=0.1): continue - request = conn.recv() + payload = conn.recv_bytes() except EOFError: break + request: Any = None try: + request = decode(payload) response = _handle_request(request, state) except Exception as e: response = WorkerResponse( @@ -488,18 +512,20 @@ def _worker_loop(conn: Connection, state: _WorkerState) -> None: ) try: - conn.send(response) + _send(conn, response) except (BrokenPipeError, EOFError): break except Exception as e: - # A result that cannot be pickled must not take the worker with it. + # A result that cannot be encoded must not take the worker with it. logger.error( "Worker response could not be sent", worker_id=state.worker_id, + request=type(request).__name__, + name=getattr(request, "name", None), error_repr=repr(e), ) try: - conn.send(WorkerResponse(error=f"{e.__class__.__name__}: {e}")) + _send(conn, WorkerResponse(error=f"{e.__class__.__name__}: {e}")) except (BrokenPipeError, EOFError): break diff --git a/dimos/core/coordination/test_blueprints.py b/dimos/core/coordination/test_blueprints.py index e7960bfc6b8..c7937867a0d 100644 --- a/dimos/core/coordination/test_blueprints.py +++ b/dimos/core/coordination/test_blueprints.py @@ -13,7 +13,6 @@ # limitations under the License. -import pickle from types import MappingProxyType from typing import Any, Protocol @@ -39,6 +38,8 @@ from dimos.core.module import Module from dimos.core.stream import In, Out from dimos.core.transport import LCMTransport +from dimos.protocol.encode.cbor import decode, encode +from dimos.protocol.service.system_configurator.clock_sync import ClockSyncConfigurator from dimos.spec.utils import Spec @@ -327,22 +328,22 @@ def test_autoconnect_eliminates_duplicates_keeps_newer() -> None: assert module_a_atoms[0].kwargs == {"key1": "new"} -def test_disabled_module_proxy_pickle_roundtrip() -> None: +def test_disabled_module_proxy_roundtrip() -> None: proxy = DisabledModuleProxy("SomeSpec") - restored = pickle.loads(pickle.dumps(proxy)) + restored = decode(encode(proxy)) assert repr(restored) == "" assert restored.any_method(1, 2, 3) is None -def test_blueprint_pickle_roundtrip() -> None: +def test_blueprint_roundtrip() -> None: blueprint = ( autoconnect(ModuleA.blueprint(), ModuleB.blueprint()) .global_config(option1=True, option2=42) .remappings([(ModuleA, "module_a", ModuleB)]) ) - restored = pickle.loads(pickle.dumps(blueprint)) + restored = decode(encode(blueprint)) assert restored == blueprint for name in ("transport_map", "global_config_overrides", "remapping_map"): @@ -464,3 +465,30 @@ def test_namespace_remappings_by_instance_name() -> None: ModuleA.blueprint().namespace("robot0"), ModuleA.blueprint().namespace("robot1"), ).remappings([(ModuleA, "data1", "other")]) + + +def _always_ok() -> str | None: + return None + + +def test_blueprint_roundtrip_keeps_checks_and_pinned_transports() -> None: + pinned = LCMTransport("/pinned", Data1) + try: + blueprint = ( + autoconnect(ModuleA.blueprint(), ModuleB.blueprint()) + .transports({("module_a", Data1): pinned}) + .requirements(_always_ok) + .configurators(ClockSyncConfigurator()) + ) + + restored = decode(encode(blueprint)) + + assert restored.requirement_checks == (_always_ok,) + assert type(restored.configurator_checks[0]) is ClockSyncConfigurator + transport = restored.transport_map[("module_a", Data1)] + assert type(transport) is LCMTransport + assert transport.topic.topic == "/pinned" + assert transport.topic.lcm_type is Data1 + assert [atom.module for atom in restored.blueprints] == [ModuleA, ModuleB] + finally: + pinned.stop() diff --git a/dimos/core/coordination/test_module_coordinator.py b/dimos/core/coordination/test_module_coordinator.py index 603f59e574c..8ab3edd3ccc 100644 --- a/dimos/core/coordination/test_module_coordinator.py +++ b/dimos/core/coordination/test_module_coordinator.py @@ -13,7 +13,6 @@ # limitations under the License. from collections.abc import Callable -import pickle import threading from types import MappingProxyType from typing import Any, Protocol @@ -52,6 +51,7 @@ from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.encode.cbor import decode, encode import dimos.robot.get_all_blueprints as resolver from dimos.spec.utils import Spec @@ -1151,11 +1151,12 @@ def test_deploy_two_instances_of_same_class(dynamic_coordinator) -> None: 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.""" +def test_rpc_client_roundtrip_preserves_remote_name(dynamic_coordinator) -> None: + """Proxies sent 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)) + restored = decode(encode(proxy)) try: + assert restored.actor_instance is None assert restored.remote_name == "robot0/namedmodule" assert restored.whoami() == "robot0/namedmodule" finally: diff --git a/dimos/core/coordination/test_worker.py b/dimos/core/coordination/test_worker.py index 3f7677b0827..8155dde8d8e 100644 --- a/dimos/core/coordination/test_worker.py +++ b/dimos/core/coordination/test_worker.py @@ -19,13 +19,14 @@ import pytest from dimos.agents.annotation import skill -from dimos.core.coordination.python_worker import PythonWorker +from dimos.core.coordination.python_worker import Actor, MethodCallProxy, PythonWorker from dimos.core.coordination.worker_manager_python import WorkerManagerPython from dimos.core.core import rpc from dimos.core.global_config import GlobalConfig, global_config from dimos.core.module import Module from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.protocol.encode.cbor import decode, encode if TYPE_CHECKING: from dimos.core.resource_monitor.stats import WorkerStats @@ -55,7 +56,7 @@ def read_global_robot_ip(self) -> str | None: return global_config.robot_ip @rpc - def unpicklable(self) -> threading.Lock: + def unencodable(self) -> threading.Lock: return threading.Lock() @rpc @@ -188,29 +189,29 @@ def test_worker_shutdown_kills_and_reaps_a_process_that_ignores_terminate(mocker assert worker.pid is None -def test_unpicklable_rpc_result_raises_without_killing_the_worker(create_worker_manager): - """An RPC result that cannot pickle comes back as the error, not a dead worker.""" +def test_unencodable_rpc_result_raises_without_killing_the_worker(create_worker_manager): + """An RPC result that cannot be encoded comes back as the error, not a dead worker.""" worker_manager = create_worker_manager(n_workers=1) module = worker_manager.deploy(SimpleModule, global_config, {}) module.start() - with pytest.raises(TypeError, match="cannot pickle"): - module.unpicklable() + with pytest.raises(TypeError, match="Cannot CBOR-encode"): + module.unencodable() assert module.increment() == 1 module.stop() -def test_unpicklable_pipe_response_raises_without_killing_the_worker(create_worker_manager): - """A response the worker pipe cannot pickle errors out; the worker lives on.""" +def test_unencodable_pipe_response_raises_without_killing_the_worker(create_worker_manager): + """A response the worker pipe cannot encode errors out; the worker lives on.""" worker_manager = create_worker_manager(n_workers=1) module = worker_manager.deploy(SimpleModule, global_config, {}) module.start() module.make_lock() # Attribute access rides the pipe, unlike rpc calls, which ride the transport. - with pytest.raises(RuntimeError, match="(?i)pickle"): + with pytest.raises(RuntimeError, match="Cannot CBOR-encode"): module.actor_instance.lock # noqa: B018 assert module.increment() == 1 @@ -510,3 +511,12 @@ def test_module_can_opt_out_of_input_warm_up(create_worker_manager): assert _wait_for_warm_up(module) == (None, "warmup-UnwarmedInputsModule") module.stop() + + +def test_actor_and_method_call_proxy_cross_processes_as_ids() -> None: + actor = Actor(None, SimpleModule, worker_id=3, module_id=7, lock=threading.Lock()) + restored = decode(encode(MethodCallProxy(actor))) + assert type(restored) is MethodCallProxy + inner = restored._actor + assert (inner._worker_id, inner._module_id) == (3, 7) + assert inner._conn is None and inner._cls is None and inner._lock is None diff --git a/dimos/core/introspection/module/info.py b/dimos/core/introspection/module/info.py index 960c87b3b33..1254038320a 100644 --- a/dimos/core/introspection/module/info.py +++ b/dimos/core/introspection/module/info.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any, NamedTuple if TYPE_CHECKING: - from dimos.core.stream import StreamDirection + from dimos.core.stream import StreamDirection, Transport # Internal RPCs to hide from io() output # Every module has these; io() shows what a module adds. @@ -42,14 +42,15 @@ class StreamDescriptor(NamedTuple): - """Wire metadata with the message type and transport in an opaque pickle.""" + """Wire metadata plus the message type and transport of one stream.""" module: str name: str direction: StreamDirection type_name: str channel: str - transport: bytes + type: type + transport: Transport[Any] @dataclass diff --git a/dimos/core/module.py b/dimos/core/module.py index c2efe3f8ae0..6b4b29d4bfe 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -17,7 +17,6 @@ from functools import partial import inspect import json -import pickle import sys import threading from typing import ( @@ -43,7 +42,6 @@ from dimos.core.introspection.module.info import StreamDescriptor, extract_module_info from dimos.core.introspection.module.render import render_module_io 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 import PubSubTransport from dimos.core.transport_factory import rpc_backend @@ -85,14 +83,13 @@ class SkillInfo: _skills_lock = threading.Lock() +@dataclass(frozen=True, slots=True) class PeekNotFound: """Sentinel returned by `Module.peek_stream` when the named stream is - not present on a module. A class instance survives pickle round-trips so + not present on a module. A class instance survives the RPC round trip so `Dimos.peek_stream` can `isinstance(result, PeekNotFound)`-test the reply. """ - __slots__ = () - def __repr__(self) -> str: return "" @@ -159,7 +156,6 @@ class ModuleBase(Configurable, CompositeResource): _tf: TF | None = None _loop: asyncio.AbstractEventLoop | None = None _loop_thread: threading.Thread | None - _bound_rpc_calls: dict[str, RpcCall] = {} _module_closed: bool = False _module_closed_lock: threading.Lock _loop_lock: threading.Lock @@ -270,35 +266,6 @@ def _close_rpc(self) -> None: self.rpc.stop() # type: ignore[attr-defined] self.rpc = None # type: ignore[assignment] - def __getstate__(self): # type: ignore[no-untyped-def] - """Exclude unpicklable runtime attributes when serializing.""" - state = self.__dict__.copy() - # Remove unpicklable attributes - state.pop("_disposables", None) - state.pop("_module_closed_lock", None) - state.pop("_loop_lock", None) - state.pop("_loop", None) - state.pop("_loop_thread", None) - state.pop("_rpc", None) - state.pop("_tf", None) # conventional slot for a TF view over the tf port - state.pop("_main_gen", None) - state.pop("_tools", None) - state.pop("_tools_lock", None) - return state - - def __setstate__(self, state) -> None: # type: ignore[no-untyped-def] - """Restore object from pickled state.""" - self.__dict__.update(state) - # Reinitialize runtime attributes - self._module_closed_lock = threading.Lock() - self._loop_lock = threading.Lock() - self._loop = None - self._loop_thread = None - self._rpc = None - self._main_gen = None - self._tools = {} - self._tools_lock = threading.Lock() - @property def outputs(self) -> dict[str, Out]: # type: ignore[type-arg] return { @@ -394,7 +361,8 @@ def get_stream_descriptors(self) -> list[StreamDescriptor]: channel=transport.channel if isinstance(transport, PubSubTransport) else type(transport).__name__, - transport=pickle.dumps((stream.type, transport)), + type=stream.type, + transport=transport, ) ) return descriptors diff --git a/dimos/core/o3dpickle.py b/dimos/core/o3dpickle.py index 02d2cbc1983..af76effe9d3 100644 --- a/dimos/core/o3dpickle.py +++ b/dimos/core/o3dpickle.py @@ -12,16 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copyreg -import sys - -import numpy as np - - -def reduce_external(obj): # type: ignore[no-untyped-def] - return (reconstruct_pointcloud, (np.asarray(obj.points),)) - - # Recorded lidar pickles under data/ reference this function by module path. def reconstruct_pointcloud(points_array): # type: ignore[no-untyped-def] import open3d as o3d # type: ignore[import-untyped] @@ -29,14 +19,3 @@ def reconstruct_pointcloud(points_array): # type: ignore[no-untyped-def] pc = o3d.geometry.PointCloud() pc.points = o3d.utility.Vector3dVector(points_array) return pc - - -def register_picklers() -> None: - """Teach pickle about open3d point clouds if open3d is loaded. - - Only a process that imported open3d can hold a point cloud, so skipping the - 2 s import when it is absent loses nothing. - """ - o3d = sys.modules.get("open3d") - if o3d is not None: - copyreg.pickle(o3d.geometry.PointCloud, reduce_external) diff --git a/dimos/core/rpc_client.py b/dimos/core/rpc_client.py index 05f2840ad45..528a89cb8b4 100644 --- a/dimos/core/rpc_client.py +++ b/dimos/core/rpc_client.py @@ -86,27 +86,15 @@ def __call__(self, *args, **kwargs): # type: ignore[no-untyped-def] # For stop, use call_nowait to avoid deadlock # (the remote side stops its RPC service before responding) if self._name == "stop": - self._rpc.call_nowait(f"{self._remote_name}/{self._name}", (args, kwargs)) # type: ignore[arg-type] + self._rpc.call_nowait(f"{self._remote_name}/{self._name}", (args, kwargs)) if self._stop_rpc_client: self._stop_rpc_client() return None - result, unsub_fn = self._rpc.call_sync( - f"{self._remote_name}/{self._name}", - (args, kwargs), # type: ignore[arg-type] - ) + result, unsub_fn = self._rpc.call_sync(f"{self._remote_name}/{self._name}", (args, kwargs)) self._unsub_fns.append(unsub_fn) return result - def __getstate__(self): # type: ignore[no-untyped-def] - return (self._name, self._remote_name) - - def __setstate__(self, state) -> None: # type: ignore[no-untyped-def] - self._name, self._remote_name = state - self._unsub_fns = [] - self._rpc = None - self._stop_rpc_client = None - class ModuleProxyProtocol(Protocol): """Protocol for host-side handles to remote modules (worker or Docker).""" @@ -142,7 +130,7 @@ def __init__( self.rpc = rpc self._owns_rpc = False if actor_class is None: - # Unpickled in another process; the class stays out of the pickle (see __reduce__). + # Decoded in another process; the class stays out of the payload (see cbor_encode). assert remote_name is not None and rpcs is not None else: remote_name = remote_name or actor_class.__name__ @@ -179,14 +167,19 @@ def stop_rpc_client(self) -> None: self.rpc.stop() self.rpc = None # type: ignore[assignment] - def __reduce__(self): # type: ignore[no-untyped-def] - # The module class stays out of the pickle: unpickling it in another + def cbor_encode(self) -> dict[str, Any]: + # The module class stays out of the payload: resolving it in another # worker would import the module and everything under it (a # SpatialMemory proxy alone pulls in transformers and torch). A proxy # only needs the rpc names; the class merely gives RpcCall a local - # signature. remote_name must be included or proxies pickled into - # workers would fall back to class-name RPC topics. - return (self.__class__, (self.actor_instance, None, self.remote_name, self.rpcs)) + # signature. remote_name must be included or proxies sent into workers + # would fall back to class-name RPC topics. The parent-side Actor stays + # behind too: receivers only call @rpc methods. + return {"remote_name": self.remote_name, "rpcs": sorted(self.rpcs)} + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> RPCClient: + return cls(None, None, data["remote_name"], data["rpcs"]) def __dir__(self) -> list[str]: return sorted(set(super().__dir__()) | set(self.rpcs)) @@ -234,9 +227,10 @@ def __getattr__(self, name: str): # type: ignore[no-untyped-def] # Try to avoid recursion by directly accessing attributes that are known result = self.actor_instance.__getattr__(name) - # When streams are returned from the worker, their owner is a pickled - # Actor with no connection. Replace it with a MethodCallProxy that can - # talk to the worker through the parent-side Actor's pipe. + # When streams are returned from the worker, their owner is an Actor + # that crossed the pipe without its connection. Replace it with a + # MethodCallProxy that can talk to the worker through the parent-side + # Actor's pipe. if isinstance(result, RemoteStream): result.owner = MethodCallProxy(self.actor_instance) @@ -252,7 +246,7 @@ class AsyncSpecProxy: is still synchronous over the wire. The caller's event loop stays unblocked while the response round-trips. - It's picklable so `set_module_ref`` can ship it across to the worker process. + It is CBOR-encodable so `set_module_ref` can ship it across to the worker process. """ def __init__(self, inner: Any, async_methods: frozenset[str]) -> None: @@ -277,14 +271,15 @@ async def _run() -> Any: return async_call - def __reduce__(self) -> Any: - return ( - AsyncSpecProxy, - ( - object.__getattribute__(self, "_inner"), - object.__getattribute__(self, "_async_methods"), - ), - ) + def cbor_encode(self) -> dict[str, Any]: + return { + "inner": object.__getattribute__(self, "_inner"), + "async_methods": sorted(object.__getattribute__(self, "_async_methods")), + } + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> AsyncSpecProxy: + return cls(data["inner"], frozenset(data["async_methods"])) if TYPE_CHECKING: diff --git a/dimos/core/stream.py b/dimos/core/stream.py index 9359b6d6cc7..cfe629c9a1e 100644 --- a/dimos/core/stream.py +++ b/dimos/core/stream.py @@ -175,18 +175,12 @@ def transport(self, value: Transport[T]) -> None: def state(self) -> State: return State.UNBOUND if self.owner is None else State.READY - def __reduce__(self): # type: ignore[no-untyped-def] - if self.owner is None or not hasattr(self.owner, "ref"): - raise ValueError("Cannot serialise Out without an owner ref") - return ( - RemoteOut, - ( - self.type, - self.name, - self.owner.ref, - self._transport, - ), - ) + def cbor_encode(self) -> dict[str, Any]: + return _remote_payload(self) + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> RemoteOut[Any]: + return RemoteOut(**data) def publish(self, msg: T) -> None: if hasattr(self, "_transport") and self._transport is not None: @@ -205,11 +199,35 @@ def unsubscribe() -> None: return unsubscribe +def _remote_payload(stream: Stream[Any]) -> dict[str, Any]: + """A module-side stream as seen from another process: a handle on the owner's Actor.""" + if stream.owner is None or not hasattr(stream.owner, "ref"): + raise ValueError(f"Cannot serialise {type(stream).__name__} without an owner ref") + return { + "type": stream.type, + "name": stream.name, + "owner": stream.owner.ref, + "transport": stream._transport, + } + + class RemoteStream(Stream[T]): @property def state(self) -> State: return State.UNBOUND if self.owner is None else State.READY + def cbor_encode(self) -> dict[str, Any]: + return { + "type": self.type, + "name": self.name, + "owner": self.owner, + "transport": self._transport, + } + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> RemoteStream[Any]: + return cls(**data) + # this won't work but nvm @property def transport(self) -> Transport[T]: @@ -250,10 +268,12 @@ def transport(self, value: Transport[T]) -> None: def state(self) -> State: return State.UNBOUND if self.owner is None else State.READY - def __reduce__(self): # type: ignore[no-untyped-def] - if self.owner is None or not hasattr(self.owner, "ref"): - raise ValueError("Cannot serialise IO without an owner ref") - return (RemoteIO, (self.type, self.name, self.owner.ref, self._transport)) + def cbor_encode(self) -> dict[str, Any]: + return _remote_payload(self) + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> RemoteIO[Any]: + return RemoteIO(**data) def publish(self, msg: T) -> None: self._transport.broadcast(self, msg) @@ -286,10 +306,12 @@ def __str__(self) -> str: return (mystr + " ◀─").ljust(60, "─") + f" {self.connection}" - def __reduce__(self): # type: ignore[no-untyped-def] - if self.owner is None or not hasattr(self.owner, "ref"): - raise ValueError("Cannot serialise Out without an owner ref") - return (RemoteIn, (self.type, self.name, self.owner.ref, self._transport)) + def cbor_encode(self) -> dict[str, Any]: + return _remote_payload(self) + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> RemoteIn[Any]: + return RemoteIn(**data) @property def transport(self) -> Transport[T]: diff --git a/dimos/core/test_baked_host.py b/dimos/core/test_baked_host.py index 119cb1b4045..0a9e2f37b25 100644 --- a/dimos/core/test_baked_host.py +++ b/dimos/core/test_baked_host.py @@ -17,7 +17,6 @@ from collections.abc import Mapping import json import pathlib -import pickle from typing import get_args, get_origin, get_type_hints import pytest @@ -27,6 +26,7 @@ from dimos.core.stream import IO, In, Out from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.protocol.encode.cbor import decode, encode class MapperConfig(NativeModuleConfig): @@ -181,9 +181,9 @@ def test_an_empty_member_map_is_refused() -> None: Pickled = baked_host("Pickled", executable="dist/x", members=MEMBERS) -def test_the_generated_classes_pickle() -> None: - """Deploying to a worker pickles both, so both must resolve by module path.""" - assert pickle.loads(pickle.dumps(Pickled)) is Pickled +def test_the_generated_classes_cross_the_worker_pipe() -> None: + """Deploying to a worker sends both by reference, so both must resolve by module path.""" + assert decode(encode(Pickled)) is Pickled host = Pickled() - assert pickle.loads(pickle.dumps(host.config)).mapper_config.voxel_size == 0.1 + assert decode(encode(host.config)).mapper_config.voxel_size == 0.1 host.stop() diff --git a/dimos/core/test_rpc_client.py b/dimos/core/test_rpc_client.py index 3b19ca0fc4f..dff5f397b63 100644 --- a/dimos/core/test_rpc_client.py +++ b/dimos/core/test_rpc_client.py @@ -19,8 +19,10 @@ import pytest from traitlets.config import Config +from dimos.core.coordination.blueprints import DisabledModuleProxy from dimos.core.demos.stress_test_module import StressTestModule -from dimos.core.rpc_client import RPCClient +from dimos.core.rpc_client import AsyncSpecProxy, RPCClient +from dimos.protocol.encode.cbor import decode, encode from dimos.protocol.rpc.spec import RPCSpec @@ -64,3 +66,11 @@ def test_ipython_completes_rpc_without_transport_calls(ipython, proxy): assert "echo" in {completion.text for completion in completions} assert transport.mock_calls == [] + + +def test_async_spec_proxy_roundtrip_keeps_its_async_methods() -> None: + proxy = AsyncSpecProxy(DisabledModuleProxy("SomeSpec"), frozenset({"foo"})) + restored = decode(encode(proxy)) + assert type(restored) is AsyncSpecProxy + assert object.__getattribute__(restored, "_async_methods") == frozenset({"foo"}) + assert isinstance(object.__getattribute__(restored, "_inner"), DisabledModuleProxy) diff --git a/dimos/core/test_stream.py b/dimos/core/test_stream.py index 41f7fb95628..bc49f0ec35a 100644 --- a/dimos/core/test_stream.py +++ b/dimos/core/test_stream.py @@ -15,16 +15,19 @@ from collections.abc import Callable import threading import time +from types import SimpleNamespace from typing import Any import pytest +from dimos.core.coordination.python_worker import Actor from dimos.core.core import rpc from dimos.core.module import Module -from dimos.core.stream import In, Out +from dimos.core.stream import In, Out, RemoteOut from dimos.core.testing import MockRobotClient from dimos.core.transport import LCMTransport, pLCMTransport from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.protocol.encode.cbor import decode, encode from dimos.robot.unitree.type.odometry import Odometry @@ -141,9 +144,6 @@ def get_hot(self): class SpyLCMTransport(LCMTransport): active_subscribers: int = 0 - def __reduce__(self): - return (SpyLCMTransport, (self.topic.topic, self.topic.lcm_type)) - def __init__(self, topic: str, type: type, **kwargs) -> None: super().__init__(topic, type, **kwargs) self._subscriber_map = {} # Maps unsubscribe functions to track active subs @@ -295,3 +295,16 @@ def unsubscribe_on_first(_: int) -> None: out.publish(2) assert received == [1, 2] + + +def test_module_stream_crosses_processes_as_a_remote_handle() -> None: + owner = SimpleNamespace(ref=Actor(None, None, worker_id=1, module_id=2)) + out = Out(Odometry, "odom", owner=owner) + remote = decode(encode(out)) + assert type(remote) is RemoteOut + assert remote.type is Odometry and remote.name == "odom" + assert (remote.owner._worker_id, remote.owner._module_id) == (1, 2) + assert remote.transport is None + + with pytest.raises(ValueError, match="owner ref"): + encode(In(Odometry, "odom")) diff --git a/dimos/core/test_transport.py b/dimos/core/test_transport.py index b25292a3b2b..43fb0b5ed21 100644 --- a/dimos/core/test_transport.py +++ b/dimos/core/test_transport.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pickle import pytest @@ -22,14 +21,15 @@ pLCMTransport, ) from dimos.msgs.sensor_msgs.Image import Image +from dimos.protocol.encode.cbor import decode, encode @pytest.mark.parametrize("transport_class", [pLCMTransport, LCMTransport, JpegLcmTransport]) -def test_lcm_transport_pickle_preserves_connection_settings(transport_class): +def test_lcm_transport_roundtrip_preserves_connection_settings(transport_class): args = () if transport_class is pLCMTransport else (Image,) transport = transport_class("/roundtrip", *args, url="udpm://239.255.76.68:7798?ttl=2", ttl=2) try: - restored = pickle.loads(pickle.dumps(transport)) + restored = decode(encode(transport)) try: assert type(restored) is transport_class assert restored.topic == transport.topic @@ -38,3 +38,18 @@ def test_lcm_transport_pickle_preserves_connection_settings(transport_class): restored.stop() finally: transport.stop() + + +def test_with_topic_keeps_every_other_setting() -> None: + transport = LCMTransport("/a", Image, url="udpm://239.255.76.68:7798?ttl=2", ttl=2) + try: + moved = transport.with_topic("/b") + try: + assert type(moved) is LCMTransport + assert moved.topic.topic == "/b" + assert moved.topic.lcm_type is Image + assert moved.lcm.config == transport.lcm.config + finally: + moved.stop() + finally: + transport.stop() diff --git a/dimos/core/test_zenoh_transport.py b/dimos/core/test_zenoh_transport.py index 8231035f826..4776455fe7f 100644 --- a/dimos/core/test_zenoh_transport.py +++ b/dimos/core/test_zenoh_transport.py @@ -14,7 +14,6 @@ from __future__ import annotations -import pickle import threading from types import SimpleNamespace from typing import cast @@ -40,6 +39,7 @@ pZenohTransport, ) from dimos.msgs.sensor_msgs.Image import Image +from dimos.protocol.encode.cbor import decode, encode from dimos.protocol.pubsub.impl.zenohpubsub import ( QOS_LATEST_WINS, QOS_NEVER_DROP, @@ -224,16 +224,16 @@ def test_stop_and_restart(session_pool) -> None: t.stop() -def test_zenoh_transport_pickle_preserves_topic_qos() -> None: +def test_zenoh_transport_roundtrip_preserves_topic_qos() -> None: t = ZenohTransport(ZenohTopic("dimos/camera/color", Image, qos=QOS_LATEST_WINS)) - t2 = pickle.loads(pickle.dumps(t)) + t2 = decode(encode(t)) assert type(t2) is ZenohTransport assert t2.topic == t.topic # topic, lcm_type and qos all round-trip -def test_pzenoh_transport_pickle_preserves_topic_qos() -> None: +def test_pzenoh_transport_roundtrip_preserves_topic_qos() -> None: t = pZenohTransport(ZenohTopic("dimos/human_input", qos=QOS_NEVER_DROP)) - t2 = pickle.loads(pickle.dumps(t)) + t2 = decode(encode(t)) assert type(t2) is pZenohTransport assert t2.topic == "dimos/human_input" assert t2._zenoh_topic == t._zenoh_topic diff --git a/dimos/core/transport.py b/dimos/core/transport.py index ffef78418dd..3e3a242d507 100644 --- a/dimos/core/transport.py +++ b/dimos/core/transport.py @@ -14,7 +14,7 @@ from __future__ import annotations -from functools import partial +import copy import importlib.util import threading import time @@ -98,6 +98,29 @@ def spec(cls, *args: Any, **kwargs: Any) -> TransportSpec: return TransportSpec(cls, args, kwargs) + # Transports cross process boundaries as their constructor settings (never the + # live pubsub): a dict with the plain "topic" string plus whatever the subclass + # needs to rebuild itself. + def cbor_encode(self) -> dict[str, Any]: + raise NotImplementedError + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> PubSubTransport[Any]: + raise NotImplementedError + + def __deepcopy__(self, memo: dict[int, Any]) -> PubSubTransport[T]: + # A fresh instance from the same settings, never a copy of the live pubsub: + # blueprint config values are deep-copied before deployment. + return cast( + "PubSubTransport[T]", type(self).cbor_decode(copy.deepcopy(self.cbor_encode(), memo)) + ) + + def with_topic(self, topic: str) -> PubSubTransport[T]: + """A copy of this transport on another topic, keeping every other setting.""" + return cast( + "PubSubTransport[T]", type(self).cbor_decode({**self.cbor_encode(), "topic": topic}) + ) + class pLCMTransport(PubSubTransport[T]): _started: bool = False @@ -106,11 +129,13 @@ def __init__(self, topic: str, **kwargs) -> None: # type: ignore[no-untyped-def super().__init__(topic) self.lcm = PickleLCM(**kwargs) - def __reduce__(self): # type: ignore[no-untyped-def] - return ( - partial(pLCMTransport, **self.lcm.config.model_dump(exclude={"lcm"})), - (self.topic,), - ) + def cbor_encode(self) -> dict[str, Any]: + return {"topic": self.topic, **self.lcm.config.model_dump(exclude={"lcm"})} + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> pLCMTransport[Any]: + data = dict(data) + return cls(data.pop("topic"), **data) def broadcast(self, _: Stream[T] | None, msg: T) -> None: if not self._started: @@ -150,11 +175,17 @@ def stop(self) -> None: self.lcm.stop() self._started = False - def __reduce__(self): # type: ignore[no-untyped-def] - return ( - partial(LCMTransport, **self.lcm.config.model_dump(exclude={"lcm"})), - (self.topic.topic, self.topic.lcm_type), - ) + def cbor_encode(self) -> dict[str, Any]: + return { + "topic": self.topic.topic, + "type": self.topic.lcm_type, + **self.lcm.config.model_dump(exclude={"lcm"}), + } + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> LCMTransport[Any]: + data = dict(data) + return cls(data.pop("topic"), data.pop("type"), **data) def broadcast(self, _, msg) -> None: # type: ignore[no-untyped-def] if not self._started: @@ -179,12 +210,6 @@ def __init__(self, topic: str, type: type, **kwargs) -> None: # type: ignore[no self.lcm = JpegLCM(**kwargs) # type: ignore[assignment] super().__init__(topic, type) - def __reduce__(self): # type: ignore[no-untyped-def] - return ( - partial(JpegLcmTransport, **self.lcm.config.model_dump(exclude={"lcm"})), - (self.topic.topic, self.topic.lcm_type), - ) - def start(self) -> None: self.lcm.start() self._started = True @@ -204,8 +229,12 @@ def __init__(self, topic: str, msg_type: type[DimosMsg], **kwargs: Any) -> None: super().__init__(ROSTopic(topic, msg_type)) self._kwargs = kwargs - def __reduce__(self) -> tuple[Any, ...]: - return (ROSTransport, (self.topic.topic, self.topic.msg_type)) + def cbor_encode(self) -> dict[str, Any]: + return {"topic": self.topic.topic, "msg_type": self.topic.msg_type} + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> ROSTransport: + return cls(data["topic"], data["msg_type"]) def broadcast(self, _: Stream[DimosMsg] | None, msg: DimosMsg) -> None: if self._ros is None: @@ -249,10 +278,26 @@ def __init__(self, topic: str, type: type, **kwargs) -> None: # type: ignore[no self._started: bool = False self._start_lock = threading.RLock() - def __reduce__(self) -> tuple[Any, ...]: - return ( - partial(DDSTransport, qos=self.dds.qos, **self.dds.config.model_dump()), - (self.topic.name, self.topic.data_type), + def cbor_encode(self) -> dict[str, Any]: + qos = self.dds.qos + return { + "topic": self.topic.name, + "type": self.topic.data_type, + "qos": qos.asdict() if qos is not None else None, # type: ignore[no-untyped-call] + **self.dds.config.model_dump(), + } + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> DDSTransport[Any]: + from cyclonedds.qos import Qos # cyclonedds: see DDS_AVAILABLE. + + data = dict(data) + qos = data.pop("qos") + return cls( + data.pop("topic"), + data.pop("type"), + qos=Qos.fromdict(qos) if qos is not None else None, + **data, ) def start(self) -> None: @@ -283,12 +328,6 @@ def subscribe( M = TypeVar("M", bound=DimosMsg) -def _rebuild_webrtc_transport( - cls: type[WebRTCTransport[M]], topic: str, msg_type: type[M] | None, config: ProviderConfig -) -> WebRTCTransport[M]: - return cls(topic, msg_type, config=config) - - class WebRTCTransport(PubSubTransport[M]): """Transport over WebRTC DataChannels. @@ -301,11 +340,12 @@ class WebRTCTransport(PubSubTransport[M]): transports sharing one multiplexed DataChannel each receive only their own message type. - The transport itself holds no connection: the picklable ``config`` - resolves to a per-process singleton provider on first use, so transports - survive being pickled into module worker processes and all transports in - a process share one session. ``stop()`` intentionally leaves the shared - provider running (it is process-scoped). + The transport itself holds no connection: the ``config`` (a pydantic model + that crosses process boundaries as data) resolves to a per-process + singleton provider on first use, so transports survive the trip into module + worker processes and all transports in a process share one session. + ``stop()`` intentionally leaves the shared provider running (it is + process-scoped). """ _config_cls: type[ProviderConfig] @@ -326,12 +366,16 @@ def __init__( self._pubsub: WebRTCPubSub | None = None # Guards first-use init: concurrent subscribe()/broadcast() must not # construct two WebRTCPubSub wrappers (one would silently orphan any - # subscribe_all state). Never pickled — __reduce__ rebuilds via the + # subscribe_all state). Never serialized: cbor_decode rebuilds via the # constructor. self._init_lock = threading.Lock() - def __reduce__(self): # type: ignore[no-untyped-def] - return (_rebuild_webrtc_transport, (type(self), self.topic, self._msg_type, self._config)) + def cbor_encode(self) -> dict[str, Any]: + return {"topic": self.topic, "msg_type": self._msg_type, "config": self._config} + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> WebRTCTransport[Any]: + return cls(data["topic"], data["msg_type"], config=data["config"]) def broadcast(self, _: Stream[M] | None, msg: M) -> None: if not self._started: @@ -405,6 +449,13 @@ class _ProviderTransport(Transport[T]): def __init__(self, *, config: ProviderConfig | None = None, **config_kwargs: Any) -> None: self._config = config or self._config_cls(**config_kwargs) + def cbor_encode(self) -> ProviderConfig: + return self._config + + @classmethod + def cbor_decode(cls, config: ProviderConfig) -> _ProviderTransport[Any]: + return cls(config=config) + @classmethod def spec(cls, *args: Any, **kwargs: Any) -> TransportSpec: """Defer construction: capture ctor args for the coordinator to build later.""" @@ -528,8 +579,12 @@ def publish_qos(self) -> dict[str, str] | None: qos = self.topic.qos return qos.to_wire() if qos is not None else None - def __reduce__(self) -> tuple[Any, ...]: - return (ZenohTransport, (self.topic,)) + def cbor_encode(self) -> dict[str, Any]: + return {"topic": self.topic.topic, "type": self.topic.lcm_type, "qos": self.topic.qos} + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> ZenohTransport[Any]: + return cls(ZenohTopic(data["topic"], data["type"], qos=data["qos"])) def start(self) -> None: with self._start_lock: @@ -580,8 +635,12 @@ def publish_qos(self) -> dict[str, str] | None: qos = self._zenoh_topic.qos return qos.to_wire() if qos is not None else None - def __reduce__(self) -> tuple[Any, ...]: - return (pZenohTransport, (self._zenoh_topic,)) + def cbor_encode(self) -> dict[str, Any]: + return {"topic": self._zenoh_topic.topic, "qos": self._zenoh_topic.qos} + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> pZenohTransport[Any]: + return cls(ZenohTopic(data["topic"], qos=data["qos"])) def start(self) -> None: with self._start_lock: diff --git a/dimos/experimental/robot/bosdyn/spot/rerun.py b/dimos/experimental/robot/bosdyn/spot/rerun.py index 750bde20dfb..fc128193cc4 100644 --- a/dimos/experimental/robot/bosdyn/spot/rerun.py +++ b/dimos/experimental/robot/bosdyn/spot/rerun.py @@ -16,6 +16,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import TYPE_CHECKING, Any from dimos.protocol.tf.tf import MultiTBuffer @@ -49,7 +50,7 @@ # Mirrors the tf stream so a camera pose can be looked up at the moment its # image was captured. Module-level, not an attribute of the override objects: -# those get pickled out to the bridge's worker and would each carry a copy. +# those get sent to the bridge's worker and would each carry a copy. _tf_buffer = MultiTBuffer() @@ -233,9 +234,9 @@ def spot_camera_layout() -> rrb.Blueprint: # The solution is to anchor the image to the frame *at a specific time* # Probably hasn't been noticed earlier because its less obvious for color image and for 60/30zh # we should fix this in our rerun bridge but I'm not going to do that in this spot PR +@dataclass class _ImageBakedIntoAnchor: - def __init__(self, origin: str) -> None: - self.origin = origin + origin: str def __call__(self, image: Any) -> RerunData | None: anchor = _anchor_frame() diff --git a/dimos/hosted/fragment.py b/dimos/hosted/fragment.py index 3ba492fd45f..8bf3f66daae 100644 --- a/dimos/hosted/fragment.py +++ b/dimos/hosted/fragment.py @@ -18,10 +18,10 @@ from collections.abc import Mapping from dataclasses import dataclass -import pickle from typing import Any from dimos.core.coordination.blueprints import Blueprint +from dimos.protocol.encode.cbor import decode, encode @dataclass(frozen=True, slots=True) @@ -53,7 +53,7 @@ class PythonFragmentPayload: @dataclass(frozen=True, slots=True) class HostFragment: - """Deployment unit sent to one Host; the payload stays pickled until the fragment process.""" + """Deployment unit sent to one Host; the payload stays encoded until the fragment process.""" run_id: str host_id: str @@ -64,11 +64,11 @@ class HostFragment: def create( cls, *, run_id: str, host_id: str, revision: str, payload: PythonFragmentPayload ) -> HostFragment: - payload_bytes = pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL) + payload_bytes = encode(payload) return cls(run_id=run_id, host_id=host_id, revision=revision, payload=payload_bytes) def load_payload(self) -> PythonFragmentPayload: - payload = pickle.loads(self.payload) + payload = decode(self.payload) if not isinstance(payload, PythonFragmentPayload): raise TypeError("Fragment payload is not a PythonFragmentPayload") return payload diff --git a/dimos/hosted/test_fragment.py b/dimos/hosted/test_fragment.py index f68d88810b4..ed0007d78d0 100644 --- a/dimos/hosted/test_fragment.py +++ b/dimos/hosted/test_fragment.py @@ -12,12 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pickle import pytest from dimos.core.module import Module, ModuleConfig from dimos.hosted.fragment import HostFragment, PythonFragmentPayload +from dimos.protocol.encode.cbor import encode class PayloadConfig(ModuleConfig): @@ -46,6 +46,6 @@ def test_fragment_round_trips_the_blueprint_and_the_roots_values() -> None: def test_fragment_rejects_a_foreign_payload() -> None: - fragment = HostFragment("run-1", "host-1", "revision-1", pickle.dumps({"not": "a payload"})) + fragment = HostFragment("run-1", "host-1", "revision-1", encode({"not": "a payload"})) with pytest.raises(TypeError, match="PythonFragmentPayload"): fragment.load_payload() diff --git a/dimos/manipulation/planning/spec/models.py b/dimos/manipulation/planning/spec/models.py index 1f27ba801fb..1bd628f07d1 100644 --- a/dimos/manipulation/planning/spec/models.py +++ b/dimos/manipulation/planning/spec/models.py @@ -123,7 +123,7 @@ class Obstacle: color: tuple[float, float, float, float] = DEFAULT_OBSTACLE_RGBA mesh_path: str | None = None # Plain tuples rather than an array: obstacles are deepcopied, compared and - # pickled across worker RPC, and an ndarray field breaks all three. + # sent across worker RPC, and an ndarray field breaks comparison. points: tuple[tuple[float, float, float], ...] = () octree_resolution: float | None = None touch_links: tuple[str, ...] = () diff --git a/dimos/manipulation/planning/spec/test_config.py b/dimos/manipulation/planning/spec/test_config.py index 3077607925c..1829256b7b6 100644 --- a/dimos/manipulation/planning/spec/test_config.py +++ b/dimos/manipulation/planning/spec/test_config.py @@ -15,7 +15,6 @@ from dataclasses import replace import json from pathlib import Path -import pickle from typing import Any import pytest @@ -23,6 +22,7 @@ from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser from dimos.manipulation.manipulation_module import ManipulationModule from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.protocol.encode.cbor import decode, encode from dimos.robot.assets.model import RobotModel @@ -39,7 +39,7 @@ def test_robot_model_survives_blueprint_config_round_trip(tmp_path: Path) -> Non parsed = BlueprintConfigParser(blueprint).parse(environ={}) - kwargs = pickle.loads(pickle.dumps(parsed.module_kwargs(ManipulationModule.name))) + kwargs = decode(encode(parsed.module_kwargs(ManipulationModule.name))) assert "_loaded" not in kwargs["model"]["model"] restored = RobotModelConfig(**kwargs["model"]) assert restored.model.source_path == urdf diff --git a/dimos/manipulation/planning/spec/validation.py b/dimos/manipulation/planning/spec/validation.py index e56fd5796d6..766a144f71c 100644 --- a/dimos/manipulation/planning/spec/validation.py +++ b/dimos/manipulation/planning/spec/validation.py @@ -48,7 +48,7 @@ } # An octree obstacle carries one point per occupied cell across worker RPC, where -# every point is a pickled tuple. A wrist workspace at 5 cm sits in the low +# every point is an encoded tuple. A wrist workspace at 5 cm sits in the low # thousands; this is the ceiling that keeps a misaimed room-scale map from # quietly stalling the pipe. MAX_OCTREE_POINTS = 200_000 @@ -269,7 +269,7 @@ def _validate_octree(obstacle: Obstacle) -> None: if len(obstacle.points) > MAX_OCTREE_POINTS: raise ValueError( f"OCTREE obstacle has {len(obstacle.points)} points, over the " - f"{MAX_OCTREE_POINTS} limit. Each one crosses worker RPC as a pickled " + f"{MAX_OCTREE_POINTS} limit. Each one crosses worker RPC as an encoded " f"tuple, so an unbounded map stalls the pipe." ) points = np.asarray(obstacle.points, dtype=np.float64) diff --git a/dimos/manipulation/test_manipulation_spec.py b/dimos/manipulation/test_manipulation_spec.py index 915cb70a842..985b813be2a 100644 --- a/dimos/manipulation/test_manipulation_spec.py +++ b/dimos/manipulation/test_manipulation_spec.py @@ -16,7 +16,6 @@ from dataclasses import fields, is_dataclass import inspect -import pickle from types import UnionType from typing import get_args, get_origin, get_type_hints @@ -40,6 +39,7 @@ from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.protocol.encode.cbor import decode, encode from dimos.spec.utils import spec_annotation_compliance @@ -117,7 +117,7 @@ def test_state_repr_exposes_values_for_repl_use() -> None: assert "0.25" in text assert "0.75" in text assert "home" in text - assert repr(pickle.loads(pickle.dumps(snapshot))) == text + assert repr(decode(encode(snapshot))) == text def test_plan_result_repr_summarizes_trajectory_without_dumping_points() -> None: diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index e3c1838f6b6..d0dfdd8a9d7 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -17,7 +17,6 @@ from __future__ import annotations from pathlib import Path -import pickle from unittest.mock import ANY, DEFAULT, MagicMock, call import numpy as np @@ -76,6 +75,7 @@ from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState, TrajectoryStatus +from dimos.protocol.encode.cbor import decode, encode from dimos.robot.assets.model import LoadedRobotModel, RobotModel @@ -1159,7 +1159,7 @@ def test_replaced_plan_is_not_dispatched_or_consumed(self, module_factory): module = module_factory() config = _one_joint_config() _install_generated_plan(module, config, [0.0], [0.1]) - original = pickle.loads(pickle.dumps(module._last_plan)) + original = decode(encode(module._last_plan)) _install_generated_plan(module, config, [0.0], [0.2]) replacement = module._last_plan coordinator = _control_coordinator() @@ -1172,9 +1172,7 @@ def test_replaced_plan_is_not_dispatched_or_consumed(self, module_factory): assert _executed(coordinator) is None assert module._last_plan is replacement assert ( - module.execute( - plan_id=pickle.loads(pickle.dumps(replacement)).plan_id, blocking=False - ).status + module.execute(plan_id=decode(encode(replacement)).plan_id, blocking=False).status is ExecutionStatus.ACCEPTED ) assert _executed(coordinator) is not None diff --git a/dimos/mapping/experimental/mesh.py b/dimos/mapping/experimental/mesh.py index cefd1f823ad..98202125004 100644 --- a/dimos/mapping/experimental/mesh.py +++ b/dimos/mapping/experimental/mesh.py @@ -17,6 +17,7 @@ from __future__ import annotations from collections.abc import Sequence +from dataclasses import dataclass, field import threading import time from typing import TYPE_CHECKING, Any @@ -374,28 +375,23 @@ def _run(self) -> None: MESH_ENTITY = "world/mesh" +@dataclass class MeshColours: """Rerun entries for each mesh chunk, turbo by height over the mesh's own range. Once the range moves ``recolour_m``, every chunk is recoloured (colours only). ``z_range`` pins it. Works as a keyed rerun bridge renderer and in replays. + A dataclass so it crosses into the bridge's worker as blueprint config. """ - keyed_by_seq = True + alpha: float = 0.4 + z_range: tuple[float, float] | None = None + recolour_m: float = 0.5 + static: bool = True + _heights: dict[Key, np.ndarray] = field(default_factory=dict, init=False) + _range: tuple[float, float] | None = field(default=None, init=False) - def __init__( - self, - alpha: float = 0.4, - z_range: tuple[float, float] | None = None, - recolour_m: float = 0.5, - static: bool = True, - ) -> None: - self.alpha = alpha - self.z_range = z_range - self.recolour_m = recolour_m - self.static = static - self._heights: dict[Key, np.ndarray] = {} - self._range: tuple[float, float] | None = None + keyed_by_seq = True def __call__(self, msg: TriangleMesh) -> list[RerunEntry]: out = [RerunEntry(_chunk_path(msg.key), msg.to_rerun(), self.static)] diff --git a/dimos/models/base.py b/dimos/models/base.py index fac8adc72af..265e837f46b 100644 --- a/dimos/models/base.py +++ b/dimos/models/base.py @@ -18,11 +18,12 @@ from functools import cached_property import sys -from typing import Annotated, Any +from typing import Annotated, Any, cast import torch from dimos.core.resource import Resource +from dimos.protocol.encode.cbor import register_codec from dimos.protocol.service.spec import BaseConfig, Configurable # Device string type - 'cuda', 'cpu', 'cuda:0', 'cuda:1', 'mps', etc. @@ -38,6 +39,14 @@ def default_torch_device() -> str: return "cpu" +def _torch_dtype(name: str) -> torch.dtype: + return cast("torch.dtype", getattr(torch, name)) + + +# Model configs carry a torch.dtype default and ride into workers over the CBOR codec. +register_codec(torch.dtype, lambda dtype: str(dtype).removeprefix("torch."), _torch_dtype) + + class LocalModelConfig(BaseConfig): device: DeviceType = default_torch_device() dtype: torch.dtype = torch.float32 diff --git a/dimos/msgs/manipulation_msgs/GraspCandidate.py b/dimos/msgs/manipulation_msgs/GraspCandidate.py index f83db6016f4..a3c5469eea2 100644 --- a/dimos/msgs/manipulation_msgs/GraspCandidate.py +++ b/dimos/msgs/manipulation_msgs/GraspCandidate.py @@ -31,6 +31,13 @@ def __init__(self, pose: Pose | None = None, score: float = 0.0) -> None: if not math.isfinite(self.score): raise ValueError("GraspCandidate.score must be finite") + def cbor_encode(self) -> dict[str, Pose | float]: + return {"pose": self.pose, "score": self.score} + + @classmethod + def cbor_decode(cls, data: dict[str, Pose | float]) -> GraspCandidate: + return cls(data["pose"], data["score"]) # type: ignore[arg-type] + def encode(self) -> bytes: return pickle.dumps({"pose": self.pose, "score": self.score}) diff --git a/dimos/msgs/manipulation_msgs/GraspCandidateArray.py b/dimos/msgs/manipulation_msgs/GraspCandidateArray.py index 5144d0b5b91..3b9e7999b19 100644 --- a/dimos/msgs/manipulation_msgs/GraspCandidateArray.py +++ b/dimos/msgs/manipulation_msgs/GraspCandidateArray.py @@ -38,6 +38,13 @@ def __len__(self) -> int: def __iter__(self) -> Iterator[GraspCandidate]: return iter(self.candidates) + def cbor_encode(self) -> dict[str, Header | list[GraspCandidate]]: + return {"header": self.header, "candidates": self.candidates} + + @classmethod + def cbor_decode(cls, data: dict[str, Header | list[GraspCandidate]]) -> GraspCandidateArray: + return cls(data["header"], data["candidates"]) # type: ignore[arg-type] + def encode(self) -> bytes: """Encode using the repository's pickle transport convention.""" return pickle.dumps({"header": self.header, "candidates": self.candidates}) diff --git a/dimos/msgs/manipulation_msgs/test_grasp_candidate.py b/dimos/msgs/manipulation_msgs/test_grasp_candidate.py index 9538beda52e..77bdf0ad678 100644 --- a/dimos/msgs/manipulation_msgs/test_grasp_candidate.py +++ b/dimos/msgs/manipulation_msgs/test_grasp_candidate.py @@ -18,6 +18,7 @@ from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray from dimos.msgs.std_msgs.Header import Header +from dimos.protocol.encode.cbor import decode, encode def test_grasp_candidate_round_trip_preserves_pose_and_score() -> None: @@ -49,3 +50,17 @@ def test_grasp_candidate_array_round_trip_preserves_header_and_order() -> None: assert decoded.header.frame_id == "world" assert [candidate.score for candidate in decoded] == [0.9, 0.7] assert [candidate.pose.position.x for candidate in decoded] == [0.1, 0.2] + + +def test_grasp_candidate_array_crosses_rpc_as_cbor() -> None: + proposals = GraspCandidateArray( + Header(123.0, "world"), [GraspCandidate(Pose(0.1, 0.0, 0.2), 0.9)] + ) + + decoded = decode(encode(proposals)) + + assert type(decoded) is GraspCandidateArray + assert decoded.header.frame_id == "world" + assert type(decoded.candidates[0]) is GraspCandidate + assert decoded.candidates[0].score == 0.9 + assert decoded.candidates[0].pose.position.x == 0.1 diff --git a/dimos/msgs/std_msgs/Header.py b/dimos/msgs/std_msgs/Header.py index 7394f27f421..37ff8f54c2c 100644 --- a/dimos/msgs/std_msgs/Header.py +++ b/dimos/msgs/std_msgs/Header.py @@ -75,6 +75,12 @@ def __init__(self, header: object) -> None: else: raise ValueError(f"Cannot create Header from {type(header)}") + @classmethod + def lcm_decode(cls, data: bytes) -> Header: + # The generated decoder resolves nested field types from the decoding class's + # own annotations, which here hold only ``ts``: decode as the LCM type and wrap. + return cls(LCMHeader.lcm_decode(data)) + @classmethod def now(cls, frame_id: str = "", seq: int = 1) -> Header: """Create a Header with current timestamp.""" diff --git a/dimos/navigation/global_planner/test_viz.py b/dimos/navigation/global_planner/test_viz.py index 1deb9ba4dcb..2c1ed58f7a3 100644 --- a/dimos/navigation/global_planner/test_viz.py +++ b/dimos/navigation/global_planner/test_viz.py @@ -13,7 +13,6 @@ # limitations under the License. from functools import partial -import pickle import numpy as np import rerun as rr @@ -25,6 +24,7 @@ from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.navigation.global_planner import viz from dimos.navigation.global_planner.mls_planner.viz import render_surface_region +from dimos.protocol.encode.cbor import decode, encode from dimos.visualization.rerun.bridge import is_keyed_by_seq @@ -45,17 +45,17 @@ def test_empty_path_keeps_the_last_one_drawn() -> None: ) -def test_keyed_renderers_keep_their_flag_through_the_pickle_round_trip() -> None: - plain = pickle.loads(pickle.dumps(render_map_region)) - bound = pickle.loads(pickle.dumps(partial(render_map_region, voxel_size=0.1))) +def test_keyed_renderers_keep_their_flag_through_the_round_trip() -> None: + plain = decode(encode(render_map_region)) + bound = decode(encode(partial(render_map_region, voxel_size=0.1))) assert is_keyed_by_seq(plain) and is_keyed_by_seq(bound) assert not is_keyed_by_seq(viz.render_path) - assert not is_keyed_by_seq(pickle.loads(pickle.dumps(partial(viz.render_path)))) + assert not is_keyed_by_seq(decode(encode(partial(viz.render_path)))) -def test_bridge_config_pickles_for_the_workers() -> None: - static = pickle.loads(pickle.dumps(viz.nav_static(0.7, 0.3, 0.3, 0.1))) - overrides = pickle.loads(pickle.dumps(viz.nav_visual_override(2.0, 0.08, 0.1))) +def test_bridge_config_crosses_to_the_workers() -> None: + static = decode(encode(viz.nav_static(0.7, 0.3, 0.3, 0.1))) + overrides = decode(encode(viz.nav_visual_override(2.0, 0.08, 0.1))) assert len(static["world/robot_body"](rr)) == 2 surface = overrides["world/surface_map"] diff --git a/dimos/navigation/global_planner/viz.py b/dimos/navigation/global_planner/viz.py index 21f720bbd31..d44ed3cb2be 100644 --- a/dimos/navigation/global_planner/viz.py +++ b/dimos/navigation/global_planner/viz.py @@ -116,7 +116,7 @@ def nav_static( ) -> dict[str, Callable[[ModuleType], list[Archetype]]]: """Bridge static entities: the body box and clearance cylinder riding on base_link. - Module-level partials, since blueprint config is pickled out to the workers. + Module-level partials, since blueprint config is sent by reference to the workers. """ return { "world/robot_body": partial(body_on_base_link, length=length, width=width, height=height), diff --git a/dimos/porcelain/dimos.py b/dimos/porcelain/dimos.py index f2df937ec9b..fe5eaec3b42 100644 --- a/dimos/porcelain/dimos.py +++ b/dimos/porcelain/dimos.py @@ -82,7 +82,7 @@ def run(self, target: str | Blueprint | type[ModuleBase]) -> None: On a connected Dimos, `target` is sent to the daemon: strings and registered module classes take a name-based fast path, while Module - classes and Blueprint objects are pickled and unpickled on the daemon. + classes and Blueprint objects are encoded and rebuilt on the daemon. """ with self._lock: if self._stopped: diff --git a/dimos/porcelain/remote_module_source.py b/dimos/porcelain/remote_module_source.py index a01f12fcb15..44fcf471be9 100644 --- a/dimos/porcelain/remote_module_source.py +++ b/dimos/porcelain/remote_module_source.py @@ -16,7 +16,6 @@ import functools import importlib -import pickle import threading from typing import TYPE_CHECKING, Any @@ -144,7 +143,7 @@ def _get_stream(self, remote_name: str, name: str) -> RemoteOut[Any]: ) if descriptor is None: raise AttributeError(f"{remote_name!r} has no stream named {name!r}") - stream_type, transport = pickle.loads(descriptor.transport) + stream_type, transport = descriptor.type, descriptor.transport if isinstance( transport, (WebRTCTransport, WebRTCVideoTransport, WebRTCAudioTransport) ): diff --git a/dimos/porcelain/test_remote_module_source.py b/dimos/porcelain/test_remote_module_source.py index 0db92e544e0..12c99efdbfd 100644 --- a/dimos/porcelain/test_remote_module_source.py +++ b/dimos/porcelain/test_remote_module_source.py @@ -441,8 +441,11 @@ def test_connected_stream_uses_configured_lcm_bus( def test_get_stream_subscribes_when_module_class_is_unavailable( client, stream_module, collected, mocker ): + # Patch this module's importlib binding only: the shared importlib.import_module + # also serves the CBOR codec, which must keep resolving classes by reference. import_module = mocker.patch( - "dimos.porcelain.remote_module_source.importlib.import_module", side_effect=ImportError + "dimos.porcelain.remote_module_source.importlib", + **{"import_module.side_effect": ImportError}, ) module = client.get_module("remote/streams") mocker.stop(import_module) diff --git a/dimos/protocol/encode/cbor.py b/dimos/protocol/encode/cbor.py new file mode 100644 index 00000000000..22ebe0f353f --- /dev/null +++ b/dimos/protocol/encode/cbor.py @@ -0,0 +1,420 @@ +# 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. + +"""CBOR serialization for RPC payloads and the worker pipe. + +JSON-like values (None, bool, int, float, str, bytes, list, dict with primitive keys) +map to native CBOR. Every other supported type travels inside a CBOR tag so it comes +back as the same Python type in the other process: tuple, NamedTuple, set, frozenset, +numpy arrays and scalars, Enum, DimosMsg (as its class plus LCM bytes), dataclasses, pydantic models, +Path, classes (including ``list[float]``-style generics and ``X | Y`` unions), module-level +functions, functools.partial, and any class that implements +``cbor_encode``/``cbor_decode`` (see ``CborEncodable``). Anything else raises +``TypeError`` at encode time. There is no pickle-style fallback. + +Classes and functions are sent by reference as ``"module:qualname"`` and must be +importable under that name on the receiving side. +""" + +from __future__ import annotations + +from collections.abc import Callable +import dataclasses +import enum +import functools +import importlib +import operator +import pathlib +import sys +import types +from typing import Any, Protocol + +import cbor2 +from cbor2 import CBORTag +import numpy as np +from pydantic import BaseModel + +TAG_TUPLE = 4040 +TAG_NAMEDTUPLE = 4041 +TAG_SET = 4042 +TAG_FROZENSET = 4043 +TAG_NDARRAY = 4044 +TAG_NPSCALAR = 4045 +TAG_ENUM = 4046 +TAG_DIMOSMSG = 4047 +TAG_DATACLASS = 4048 +TAG_PYDANTIC = 4049 +TAG_PATH = 4050 +TAG_CLASS = 4051 +TAG_FUNCTION = 4052 +TAG_PARTIAL = 4053 +TAG_CUSTOM = 4054 +TAG_REGISTERED = 4055 +TAG_DICT = 4056 +TAG_GENERIC_ALIAS = 4057 +TAG_UNION = 4058 +TAG_ELLIPSIS = 4059 + +_PRIMITIVES = (type(None), bool, int, float, str, bytes) + + +class CborDecodeError(ValueError): + """The payload is not valid CBOR for this codec.""" + + +class CborEncodable(Protocol): + """Opt-in serialization for classes the generic rules do not cover. + + ``cbor_encode`` returns any encodable value; ``cbor_decode`` rebuilds an instance + from it. The returned object may be a stand-in of another class (a stream handle + decodes to its remote twin), which is why ``cbor_decode`` is typed as ``Any``. + """ + + def cbor_encode(self) -> Any: ... + + @classmethod + def cbor_decode(cls, data: Any) -> Any: ... + + +# type -> (encode, "module:qualname" of the decode function) +_REGISTRY: dict[type, tuple[Callable[[Any], Any], str]] = {} + + +def register_codec(type_: type, encode: Callable[[Any], Any], decode: Callable[[Any], Any]) -> None: + """Make instances of ``type_`` (and subclasses) encodable without touching the class. + + ``decode`` must be a module-level function: the payload names it by reference so the + decoding process imports the registering module on demand. + """ + _REGISTRY[type_] = (encode, _ref(decode, repr(decode))) + + +def encode(obj: Any) -> bytes: + """Encode ``obj`` to CBOR bytes. Raises ``TypeError`` for unsupported values.""" + return cbor2.dumps(_prepare(obj, "value", set())) + + +def decode(data: bytes | bytearray | memoryview) -> Any: + """Decode CBOR bytes produced by ``encode``. Raises ``CborDecodeError`` on bad input.""" + try: + return cbor2.loads(data, semantic_decoders=_DECODERS, tag_hook=_reject_unknown_tag) + except CborDecodeError: + raise + except cbor2.CBORDecodeError as e: + detail = f"{e}: {e.__cause__}" if e.__cause__ is not None else str(e) + raise CborDecodeError(detail) from e + + +def _prepare(obj: Any, path: str, active: set[int]) -> Any: + """Map an object graph to values cbor2 encodes natively (CBORTag for our tags).""" + t = type(obj) + if t in _PRIMITIVES: + return obj + if id(obj) in active: + raise TypeError(f"Cannot CBOR-encode a cyclic reference at {path}") + active.add(id(obj)) + try: + return _prepare_object(obj, t, path, active) + finally: + active.discard(id(obj)) + + +def _prepare_object(obj: Any, t: type, path: str, active: set[int]) -> Any: + if t is list: + return _prepare_items(obj, path, active) + if t is dict: + return _prepare_dict(obj, path, active) + if t is tuple: + return CBORTag(TAG_TUPLE, _prepare_items(obj, path, active)) + + # Always look hooks up on the class: proxies fabricate instance attributes in + # __getattr__ (MethodCallProxy would send a worker request) and lazy paths resolve a + # checkout in __getattribute__. For a class object, t is its metaclass, which lets a + # metaclass serialize classes that are not importable by name. + custom_encode = getattr(t, "cbor_encode", None) + if callable(custom_encode): + payload = _prepare(custom_encode(obj), f"{path}.cbor_encode()", active) + return CBORTag(TAG_CUSTOM, [_ref(t, path), payload]) + # Generic aliases and unions come before the class check: on Python 3.10 + # isinstance(list[int], type) is True. + if isinstance(obj, types.GenericAlias): + origin = _prepare(obj.__origin__, f"{path}.__origin__", active) + return CBORTag(TAG_GENERIC_ALIAS, [origin, _prepare_items(obj.__args__, path, active)]) + if isinstance(obj, types.UnionType): + return CBORTag(TAG_UNION, _prepare_items(obj.__args__, path, active)) + if isinstance(obj, type): + return CBORTag(TAG_CLASS, _ref(obj, path)) + if obj is Ellipsis: # the "..." in tuple[int, ...] + return CBORTag(TAG_ELLIPSIS, None) + for registered_type, (registered_encode, decode_ref) in _REGISTRY.items(): + if isinstance(obj, registered_type): + payload = _prepare(registered_encode(obj), f"{path}.encode()", active) + return CBORTag(TAG_REGISTERED, [decode_ref, payload]) + if isinstance(obj, enum.Enum): + return CBORTag(TAG_ENUM, [_ref(t, path), obj.name]) + if isinstance(obj, np.ndarray): + if obj.dtype.hasobject or obj.dtype.fields is not None: + raise TypeError(f"Cannot CBOR-encode ndarray of dtype {obj.dtype} at {path}") + return CBORTag(TAG_NDARRAY, [obj.dtype.str, list(obj.shape), obj.tobytes()]) + if isinstance(obj, np.generic): + if obj.dtype.hasobject: + raise TypeError(f"Cannot CBOR-encode numpy object scalar at {path}") + return CBORTag(TAG_NPSCALAR, [obj.dtype.str, obj.tobytes()]) + if isinstance(getattr(t, "msg_name", None), str) and callable(getattr(t, "lcm_encode", None)): + # The class, not msg_name: a subclass that inherits its parent's msg_name + # must still come back as the subclass. + return CBORTag(TAG_DIMOSMSG, [_ref(t, path), obj.lcm_encode()]) + if isinstance(obj, tuple) and hasattr(t, "_fields"): + return CBORTag(TAG_NAMEDTUPLE, [_ref(t, path), _prepare_items(obj, path, active)]) + if dataclasses.is_dataclass(obj): + fields = { + f.name: _prepare(getattr(obj, f.name), f"{path}.{f.name}", active) + for f in dataclasses.fields(obj) + } + return CBORTag(TAG_DATACLASS, [_ref(t, path), fields]) + if isinstance(obj, BaseModel): + fields = { + name: _prepare(getattr(obj, name), f"{path}.{name}", active) + for name in type(obj).model_fields + } + extra = _prepare_dict(obj.__pydantic_extra__ or {}, path, active) + return CBORTag(TAG_PYDANTIC, [_ref(t, path), fields, sorted(obj.model_fields_set), extra]) + if t is set: + return CBORTag(TAG_SET, _prepare_items(obj, path, active)) + if t is frozenset: + return CBORTag(TAG_FROZENSET, _prepare_items(obj, path, active)) + if isinstance(obj, pathlib.PurePath): + return CBORTag(TAG_PATH, str(obj)) + if isinstance(obj, functools.partial): + func = _prepare(obj.func, f"{path}.func", active) + args = _prepare_items(obj.args, f"{path}.args", active) + keywords = _prepare_dict(obj.keywords, f"{path}.keywords", active) + return CBORTag(TAG_PARTIAL, [func, args, keywords]) + if isinstance(obj, (types.FunctionType, types.BuiltinFunctionType)): + return CBORTag(TAG_FUNCTION, _ref(obj, path)) + raise TypeError( + f"Cannot CBOR-encode object of type {t.__module__}.{t.__qualname__} at {path}: " + "implement cbor_encode()/cbor_decode() on the class or call " + "dimos.protocol.encode.cbor.register_codec()" + ) + + +def _prepare_items(items: Any, path: str, active: set[int]) -> list[Any]: + return [_prepare(item, f"{path}[{i}]", active) for i, item in enumerate(items)] + + +def _prepare_dict(obj: dict[Any, Any], path: str, active: set[int]) -> Any: + if all(type(key) in _PRIMITIVES for key in obj): + return {key: _prepare(value, f"{path}[{key!r}]", active) for key, value in obj.items()} + # Keys of other types (tuples, enums, str subclasses) travel as a list of pairs so + # their type survives; cbor2 would coerce or refuse them as map keys. + pairs = [ + [_prepare(key, f"{path}.keys()", active), _prepare(value, f"{path}[{key!r}]", active)] + for key, value in obj.items() + ] + return CBORTag(TAG_DICT, pairs) + + +def _ref(obj: Any, path: str) -> str: + """``"module:qualname"`` of a class or function that the receiver can import again.""" + if obj is type(None): # claims builtins as its module, which does not export it + return "types:NoneType" + module = getattr(obj, "__module__", None) + qualname = getattr(obj, "__qualname__", None) + if not isinstance(module, str) or not isinstance(qualname, str): + raise TypeError(f"Cannot CBOR-encode {obj!r} at {path}: it has no module path") + found = _lookup(sys.modules.get(module), qualname) + # Identity for functions, as pickle does. A class only has to be found under its + # name: after importlib.reload the module holds a replacement class object, and + # the receiver resolves by name anyway. + if found is not obj and not (isinstance(obj, type) and isinstance(found, type)): + raise TypeError( + f"Cannot CBOR-encode {obj!r} at {path}: it is not importable as " + f"{module}.{qualname} (a lambda, a local definition or an unassigned runtime " + "object); define it at module level" + ) + return f"{module}:{qualname}" + + +def _lookup(target: Any, qualname: str) -> Any: + for part in qualname.split("."): + target = getattr(target, part, None) + if target is None: + return None + return target + + +def _resolve(ref: str) -> Any: + module_name, _, qualname = ref.partition(":") + try: + module = importlib.import_module(module_name) + except ImportError as e: + raise CborDecodeError(f"cannot import {module_name!r} for {ref!r}: {e}") from e + target = _lookup(module, qualname) + if target is None: + raise CborDecodeError(f"{ref!r} does not resolve to an object") + return target + + +def _resolve_class(ref: str, base: type) -> Any: + cls = _resolve(ref) + if not (isinstance(cls, type) and issubclass(cls, base)): + raise CborDecodeError(f"{ref!r} is not a {base.__name__} subclass") + return cls + + +def _decode_tuple(value: Any, _: bool) -> Any: + return tuple(value) + + +def _decode_namedtuple(value: Any, _: bool) -> Any: + ref, items = value + cls = _resolve_class(ref, tuple) + if not hasattr(cls, "_fields"): + raise CborDecodeError(f"{ref!r} is not a NamedTuple") + return cls._make(items) + + +def _decode_set(value: Any, _: bool) -> Any: + return set(value) + + +def _decode_frozenset(value: Any, _: bool) -> Any: + return frozenset(value) + + +def _decode_ndarray(value: Any, _: bool) -> Any: + dtype, shape, data = value + # frombuffer on the decoded bytes is read-only; the copy makes it writable. + return np.frombuffer(data, dtype=np.dtype(dtype)).reshape(shape).copy() + + +def _decode_npscalar(value: Any, _: bool) -> Any: + dtype, data = value + return np.frombuffer(data, dtype=np.dtype(dtype))[0] + + +def _decode_enum(value: Any, _: bool) -> Any: + ref, name = value + return _resolve_class(ref, enum.Enum)[name] + + +def _decode_dimosmsg(value: Any, _: bool) -> Any: + ref, data = value + lcm_decode = getattr(_resolve(ref), "lcm_decode", None) + if not callable(lcm_decode): + raise CborDecodeError(f"{ref!r} has no lcm_decode()") + return lcm_decode(data) + + +def _decode_dataclass(value: Any, _: bool) -> Any: + ref, fields = value + cls = _resolve(ref) + if not (isinstance(cls, type) and dataclasses.is_dataclass(cls)): + raise CborDecodeError(f"{ref!r} is not a dataclass") + # No __init__ or __post_init__: the sender already normalized the values, and this + # also handles frozen classes and init=False fields. + obj = cls.__new__(cls) + for f in dataclasses.fields(cls): + object.__setattr__(obj, f.name, fields[f.name]) + return obj + + +def _decode_pydantic(value: Any, _: bool) -> Any: + ref, fields, fields_set, extra = value + cls = _resolve_class(ref, BaseModel) + # model_construct skips validation (the sender validated) and keeps model_fields_set. + # It also keeps BaseSettings from re-reading the environment on the receiving side. + return cls.model_construct(_fields_set=set(fields_set), **fields, **extra) + + +def _decode_path(value: Any, _: bool) -> Any: + return pathlib.Path(value) + + +def _decode_class(value: Any, _: bool) -> Any: + cls = _resolve(value) + if not isinstance(cls, type): + raise CborDecodeError(f"{value!r} is not a class") + return cls + + +def _decode_function(value: Any, _: bool) -> Any: + func = _resolve(value) + if not callable(func): + raise CborDecodeError(f"{value!r} is not callable") + return func + + +def _decode_partial(value: Any, _: bool) -> Any: + func, args, keywords = value + return functools.partial(func, *args, **keywords) + + +def _decode_custom(value: Any, _: bool) -> Any: + ref, payload = value + custom_decode = getattr(_resolve(ref), "cbor_decode", None) + if not callable(custom_decode): + raise CborDecodeError(f"{ref!r} has no cbor_decode()") + return custom_decode(payload) + + +def _decode_registered(value: Any, _: bool) -> Any: + ref, payload = value + return _resolve(ref)(payload) + + +def _decode_dict(value: Any, _: bool) -> Any: + return dict(value) + + +def _decode_generic_alias(value: Any, _: bool) -> Any: + origin, args = value + return types.GenericAlias(origin, tuple(args)) + + +def _decode_union(value: Any, _: bool) -> Any: + return functools.reduce(operator.or_, value) + + +def _decode_ellipsis(_value: Any, _: bool) -> Any: + return Ellipsis + + +def _reject_unknown_tag(tag: CBORTag, _: bool) -> Any: + raise CborDecodeError(f"unknown CBOR tag {tag.tag}") + + +_DECODERS: dict[int, Callable[[Any, bool], Any]] = { + TAG_TUPLE: _decode_tuple, + TAG_NAMEDTUPLE: _decode_namedtuple, + TAG_SET: _decode_set, + TAG_FROZENSET: _decode_frozenset, + TAG_NDARRAY: _decode_ndarray, + TAG_NPSCALAR: _decode_npscalar, + TAG_ENUM: _decode_enum, + TAG_DIMOSMSG: _decode_dimosmsg, + TAG_DATACLASS: _decode_dataclass, + TAG_PYDANTIC: _decode_pydantic, + TAG_PATH: _decode_path, + TAG_CLASS: _decode_class, + TAG_FUNCTION: _decode_function, + TAG_PARTIAL: _decode_partial, + TAG_CUSTOM: _decode_custom, + TAG_REGISTERED: _decode_registered, + TAG_DICT: _decode_dict, + TAG_GENERIC_ALIAS: _decode_generic_alias, + TAG_UNION: _decode_union, + TAG_ELLIPSIS: _decode_ellipsis, +} diff --git a/dimos/protocol/encode/test_cbor.py b/dimos/protocol/encode/test_cbor.py new file mode 100644 index 00000000000..7b8a38374a0 --- /dev/null +++ b/dimos/protocol/encode/test_cbor.py @@ -0,0 +1,432 @@ +# 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 __future__ import annotations + +from dataclasses import dataclass, field +import enum +import functools +import math +from pathlib import Path +import threading +from types import MappingProxyType +from typing import Any, NamedTuple + +from cbor2 import CBORTag, dumps +import numpy as np +from pydantic import BaseModel, ConfigDict, Field +import pytest + +from dimos.core.global_config import GlobalConfig +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.protocol.encode.cbor import ( + TAG_CLASS, + TAG_DATACLASS, + TAG_DIMOSMSG, + TAG_ENUM, + CborDecodeError, + decode, + encode, + register_codec, +) +from dimos.types.robot_location import RobotLocation + + +class Point(NamedTuple): + x: int + y: str + + +class Color(enum.Enum): + RED = "r" + BLUE = "b" + + +class Level(enum.IntEnum): + LOW = 1 + HIGH = 2 + + +class Mode(str, enum.Enum): + FAST = "fast" + + +@dataclass(frozen=True) +class Frozen: + a: int + b: tuple[int, ...] = () + cache: int = field(default=7, init=False) + + +@dataclass(slots=True) +class Slotted: + x: float + tags: list[str] + + +class Settings(BaseModel): + model_config = ConfigDict(extra="allow", populate_by_name=True) + + name: str = Field(alias="Name") + count: int = 1 + nested: Slotted | None = None + + +class Temperature: + def __init__(self, kelvin: float) -> None: + self.kelvin = kelvin + + def cbor_encode(self) -> float: + return self.kelvin + + @classmethod + def cbor_decode(cls, data: float) -> Temperature: + return cls(data) + + +class RemoteHandle: + def __init__(self, name: str) -> None: + self.name = name + + +class LiveHandle: + """Decodes to a stand-in of another class, like a stream decoding to its remote twin.""" + + def __init__(self, name: str) -> None: + self.name = name + + def cbor_encode(self) -> str: + return self.name + + @classmethod + def cbor_decode(cls, data: str) -> RemoteHandle: + return RemoteHandle(data) + + +class TaggedMeta(type): + """Metaclass hooks serialize classes that are not importable by name.""" + + def cbor_encode(cls) -> str: + return cls.tag # type: ignore[attr-defined] + + @classmethod + def cbor_decode(cls, data: str) -> type: + return _CLASSES_BY_TAG[data] + + +class Tagged(metaclass=TaggedMeta): + tag = "tagged" + + +_CLASSES_BY_TAG = {"tagged": Tagged} + + +class Celsius: + def __init__(self, degrees: float) -> None: + self.degrees = degrees + + +def _decode_celsius(degrees: float) -> Celsius: + return Celsius(degrees) + + +register_codec(Celsius, lambda c: c.degrees, _decode_celsius) + + +class Proxy: + """Fabricates every attribute; the codec must never consult the instance.""" + + def __init__(self) -> None: + self.calls: list[str] = [] + + def __getattr__(self, name: str) -> Any: + self.calls.append(name) + return lambda *args: None + + +class Outer: + class Inner: + pass + + +def module_fn(a: int, k: int = 0) -> int: + return a + k + + +class Plain: + def method(self) -> None: + pass + + +def roundtrip(value: Any) -> Any: + return decode(encode(value)) + + +@pytest.mark.parametrize( + "value", + [ + None, + True, + 0, + -5, + 2**70, + 1.5, + math.inf, + "text", + b"bytes", + [1, "a", [2.5, None]], + {"a": 1, "b": {"c": [1, 2]}, 3: "int key", None: "none key"}, + (), + (1, "two", (3.0,)), + {1, (2, 3)}, + frozenset({"a", "b"}), + Path("/tmp/x/y.txt"), + Color.RED, + Level.HIGH, + Mode.FAST, + Point(1, "y"), + Frozen(1, (2, 3)), + Slotted(1.5, ["t"]), + Vector3(1, 2, 3), + PoseStamped(ts=1.5, frame_id="odom", position=[1, 2, 3]), + ], +) +def test_roundtrip_preserves_value_and_type(value: Any) -> None: + out = roundtrip(value) + assert out == value + assert type(out) is type(value) + + +def test_nan_roundtrip() -> None: + assert math.isnan(roundtrip(math.nan)) + + +def test_nested_container_types_survive() -> None: + value = {"t": (1, [2, (3,)]), "s": {frozenset({1})}, "l": [Point(1, "a")]} + out = roundtrip(value) + assert out == value + assert type(out["t"]) is tuple + assert type(out["t"][1]) is list + assert type(out["t"][1][1]) is tuple + assert type(next(iter(out["s"]))) is frozenset + assert type(out["l"][0]) is Point + + +def test_dict_with_non_primitive_keys() -> None: + value = {(1, 2): "tuple", Color.BLUE: "enum", np.float64(1.5): "np"} + out = roundtrip(value) + assert out == value + assert [type(k) for k in out] == [tuple, Color, np.float64] + + +@pytest.mark.parametrize("dtype", [np.bool_, np.uint8, np.uint16, np.int32, np.float32, np.float64]) +@pytest.mark.parametrize("shape", [(), (3,), (2, 3, 4)]) +def test_ndarray_roundtrip(dtype: Any, shape: tuple[int, ...]) -> None: + value = np.asarray(np.arange(math.prod(shape)).reshape(shape) % 2).astype(dtype) + out = roundtrip(value) + assert isinstance(out, np.ndarray) + assert out.dtype == value.dtype + assert out.shape == value.shape + assert np.array_equal(out, value) + assert out.flags.writeable + + +def test_fortran_ordered_array() -> None: + value = np.asfortranarray(np.arange(6).reshape(2, 3)) + assert np.array_equal(roundtrip(value), value) + + +@pytest.mark.parametrize("value", [np.float64(1.5), np.int64(-2), np.bool_(True), np.str_("hi")]) +def test_numpy_scalar_keeps_its_numpy_type(value: Any) -> None: + out = roundtrip(value) + assert type(out) is type(value) + assert out == value + + +def test_dataclass_post_init_is_not_rerun() -> None: + # __post_init__ pads the 2-tuple to 3 on construction; the receiver must get the + # normalized value as is, not run __post_init__ on it again. + loc = RobotLocation("kitchen", (1.0, 2.0), (0.5,), metadata={"k": [1]}) # type: ignore[arg-type] + assert loc.position == (1.0, 2.0, 0.0) + out = roundtrip(loc) + assert out == loc + assert type(out.position) is tuple + + +def test_frozen_dataclass_keeps_init_false_field() -> None: + value = Frozen(1) + object.__setattr__(value, "cache", 99) + assert roundtrip(value).cache == 99 + + +def test_pydantic_model_roundtrip() -> None: + model = Settings(Name="a", extra_field=3, nested=Slotted(2.0, ["t"])) # type: ignore[call-arg] + out = roundtrip(model) + assert type(out) is Settings + assert out.name == "a" + assert out.count == 1 + assert out.nested == Slotted(2.0, ["t"]) + assert out.model_fields_set == model.model_fields_set + assert out.model_extra == {"extra_field": 3} + + +def test_global_config_roundtrip() -> None: + config = GlobalConfig(robot_ip="10.0.0.1", n_workers=3) + out = roundtrip(config) + assert type(out) is GlobalConfig + assert out.model_dump() == config.model_dump() + + +@pytest.mark.parametrize( + "image", + [ + Image(data=np.arange(4 * 6 * 3, dtype=np.uint8).reshape(4, 6, 3), format=ImageFormat.RGB), + Image(data=np.full((2, 3), 700, dtype=np.uint16), format=ImageFormat.DEPTH16), + Image(data=np.full((2, 3), 0.25, dtype=np.float32), format=ImageFormat.DEPTH), + ], +) +def test_image_roundtrip(image: Image) -> None: + image.frame_id = "cam" + image.ts = 1.5 + out = roundtrip(image) + assert type(out) is Image + assert out.format == image.format + assert out.frame_id == "cam" + assert out.ts == 1.5 + assert out.data.dtype == image.data.dtype + assert np.array_equal(out.data, image.data) + + +def test_class_and_function_references() -> None: + assert roundtrip(Outer.Inner) is Outer.Inner + assert roundtrip(list[float]) == list[float] + assert roundtrip(dict[str, tuple[int, ...]]) == dict[str, tuple[int, ...]] + assert roundtrip(int | None) == int | None + assert roundtrip(module_fn) is module_fn + assert roundtrip(Color) is Color + partial = functools.partial(module_fn, 1, k=2) + out = roundtrip(partial) + assert isinstance(out, functools.partial) + assert out.func is module_fn + assert (out.args, out.keywords) == ((1,), {"k": 2}) + assert out() == 3 + + +def test_custom_protocol() -> None: + out = roundtrip({"t": Temperature(273.15)}) + assert type(out["t"]) is Temperature + assert out["t"].kelvin == 273.15 + + +def test_custom_protocol_may_return_a_stand_in() -> None: + out = roundtrip(LiveHandle("cam")) + assert type(out) is RemoteHandle + assert out.name == "cam" + + +def test_metaclass_hooks_serialize_class_objects() -> None: + assert roundtrip(Tagged) is Tagged + assert roundtrip([Tagged, Outer]) == [Tagged, Outer] + + +def test_registered_codec() -> None: + out = roundtrip(Celsius(21.0)) + assert type(out) is Celsius + assert out.degrees == 21.0 + + +def test_proxy_getattr_is_never_consulted() -> None: + proxy = Proxy() + with pytest.raises(TypeError, match="Cannot CBOR-encode"): + encode(proxy) + assert proxy.calls == [] + + +def _local_function() -> None: + pass + + +@pytest.mark.parametrize( + "value", + [ + lambda: None, + Plain(), + Plain().method, + threading.Lock(), + MappingProxyType({}), + np.array([object()]), + np.zeros(2, dtype=[("a", "i4")]), + bytearray(b"x"), + ], +) +def test_unsupported_values_raise_type_error(value: Any) -> None: + with pytest.raises(TypeError, match="Cannot CBOR-encode"): + encode(value) + + +def test_local_definitions_raise_type_error() -> None: + def inner() -> None: + pass + + class Local: + pass + + for value in (inner, Local, Local()): + with pytest.raises(TypeError, match="Cannot CBOR-encode"): + encode(value) + + +def test_error_names_the_path() -> None: + with pytest.raises(TypeError, match=r"value\[1\]\['k'\]"): + encode([1, {"k": threading.Lock()}]) + + +def test_cyclic_reference_raises() -> None: + cycle: list[Any] = [] + cycle.append(cycle) + with pytest.raises(TypeError, match="cyclic"): + encode(cycle) + + +def test_repeated_reference_is_not_a_cycle() -> None: + shared = [1] + assert roundtrip([shared, shared]) == [[1], [1]] + + +def test_register_codec_requires_module_level_decoder() -> None: + with pytest.raises(TypeError, match="Cannot CBOR-encode"): + register_codec(Plain, lambda p: None, lambda data: Plain()) + + +@pytest.mark.parametrize( + "payload", + [ + CBORTag(TAG_DATACLASS, ["builtins:int", {}]), + CBORTag(TAG_ENUM, ["builtins:int", "x"]), + CBORTag(TAG_CLASS, "math:pi"), + CBORTag(TAG_CLASS, "no_such_module_anywhere:X"), + CBORTag(TAG_DIMOSMSG, ["builtins:int", b""]), + CBORTag(999, 1), + ], +) +def test_bad_payloads_raise_decode_error(payload: Any) -> None: + with pytest.raises(CborDecodeError): + decode(dumps(payload)) + + +def test_truncated_input_raises_decode_error() -> None: + with pytest.raises(CborDecodeError): + decode(encode([1, 2, 3])[:-1]) diff --git a/dimos/protocol/pubsub/benchmark/test_replay_bench.py b/dimos/protocol/pubsub/benchmark/test_replay_bench.py index 4e1d772ed29..b16a63b8056 100644 --- a/dimos/protocol/pubsub/benchmark/test_replay_bench.py +++ b/dimos/protocol/pubsub/benchmark/test_replay_bench.py @@ -14,7 +14,6 @@ from __future__ import annotations -import pickle import statistics import threading import time @@ -26,6 +25,7 @@ from dimos.core.transport import LCMTransport, ZenohTransport from dimos.msgs.sensor_msgs.CompressedImage import CompressedImage from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.protocol.encode.cbor import decode, encode from dimos.protocol.pubsub.benchmark.tool_replay_bench import CompressedCodec from dimos.protocol.service.zenohservice import ZenohSessionPool @@ -86,10 +86,10 @@ def test_roundtrip_over_zenoh(retry_until, collector, session_pool) -> None: t.stop() -def test_pickle_roundtrip() -> None: - """The bench pins the codec into blueprints — it must pickle into forkserver workers.""" +def test_cbor_roundtrip() -> None: + """The bench pins the codec into blueprints — it must cross into forkserver workers.""" t = CompressedCodec(LCMTransport("dimos/test/codec_pickle", CompressedImage), quality=50) - t2 = pickle.loads(pickle.dumps(t)) + t2 = decode(encode(t)) assert isinstance(t2, CompressedCodec) assert t2.quality == 50 assert t2.topic.topic == t.topic.topic diff --git a/dimos/protocol/pubsub/benchmark/tool_replay_bench.py b/dimos/protocol/pubsub/benchmark/tool_replay_bench.py index ee80e5854e2..5742a14195b 100644 --- a/dimos/protocol/pubsub/benchmark/tool_replay_bench.py +++ b/dimos/protocol/pubsub/benchmark/tool_replay_bench.py @@ -63,8 +63,12 @@ def __init__(self, inner: PubSubTransport[Any], quality: int = 75) -> None: self.inner = inner self.quality = quality - def __reduce__(self): # type: ignore[no-untyped-def] - return (CompressedCodec, (self.inner, self.quality)) + def cbor_encode(self) -> dict[str, Any]: + return {"inner": self.inner, "quality": self.quality} + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> CompressedCodec[Any]: + return cls(data["inner"], data["quality"]) def broadcast(self, stream: Out[T] | None, msg: T) -> None: from dimos.msgs.sensor_msgs.CompressedImage import ( diff --git a/dimos/protocol/pubsub/encoders.py b/dimos/protocol/pubsub/encoders.py index ff5c78a2f52..042f75bdc50 100644 --- a/dimos/protocol/pubsub/encoders.py +++ b/dimos/protocol/pubsub/encoders.py @@ -21,10 +21,14 @@ from typing import TYPE_CHECKING, Generic, Protocol, TypeVar, cast from dimos.msgs.protocol import DimosMsg +from dimos.protocol.encode import cbor +from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from collections.abc import Callable +logger = setup_logger() + TopicT = TypeVar("TopicT") MsgT = TypeVar("MsgT") EncodingT = TypeVar("EncodingT") @@ -94,6 +98,24 @@ def decode(self, msg: bytes, _: TopicT) -> MsgT: return cast("MsgT", pickle.loads(msg)) +class CborEncoderMixin(PubSubEncoderMixin[TopicT, MsgT, bytes]): + """Encoder mixin using the CBOR codec (dimos.protocol.encode.cbor). + + Encoding an unsupported value raises TypeError out of publish(); a message + that does not decode is logged and dropped. + """ + + def encode(self, msg: MsgT, _: TopicT) -> bytes: + return cbor.encode(msg) + + def decode(self, msg: bytes, topic: TopicT) -> MsgT: + try: + return cast("MsgT", cbor.decode(msg)) + except cbor.CborDecodeError as e: + logger.error("Dropping undecodable CBOR message", topic=str(topic), error=str(e)) + raise DecodingError(str(e)) from e + + class LCMTopicProto(Protocol): """Protocol for topics usable with LCM encoders.""" diff --git a/dimos/protocol/pubsub/impl/lcmpubsub.py b/dimos/protocol/pubsub/impl/lcmpubsub.py index f47280f35cb..02aa8adf88a 100644 --- a/dimos/protocol/pubsub/impl/lcmpubsub.py +++ b/dimos/protocol/pubsub/impl/lcmpubsub.py @@ -22,6 +22,7 @@ from dimos.msgs.protocol import DimosMsg from dimos.protocol.pubsub.encoders import ( + CborEncoderMixin, LCMEncoderMixin, PickleEncoderMixin, ) @@ -158,3 +159,9 @@ class PickleLCM( PickleEncoderMixin, # type: ignore[type-arg] LCMPubSubBase, ): ... + + +class CborLCM( + CborEncoderMixin, # type: ignore[type-arg] + LCMPubSubBase, +): ... diff --git a/dimos/protocol/pubsub/impl/webrtc/providers/spec.py b/dimos/protocol/pubsub/impl/webrtc/providers/spec.py index 856f3645000..945a2d1c70e 100644 --- a/dimos/protocol/pubsub/impl/webrtc/providers/spec.py +++ b/dimos/protocol/pubsub/impl/webrtc/providers/spec.py @@ -14,9 +14,9 @@ """Provider contract for WebRTC DataChannel backends. -``Provider`` is the runtime interface; ``ProviderConfig`` is its picklable, -hashable description. Transports carry configs across process boundaries -(module workers receive their transports by pickle) and resolve them to a +``Provider`` is the runtime interface; ``ProviderConfig`` is its hashable +description that crosses process boundaries as data. Transports carry configs +(module workers receive their transports over the CBOR pipe) and resolve them to a per-process singleton provider — one PeerConnection per process, shared by every transport with an equal config. """ diff --git a/dimos/protocol/pubsub/impl/webrtc/test_transport.py b/dimos/protocol/pubsub/impl/webrtc/test_transport.py index d2a81dd5103..fa816169930 100644 --- a/dimos/protocol/pubsub/impl/webrtc/test_transport.py +++ b/dimos/protocol/pubsub/impl/webrtc/test_transport.py @@ -17,7 +17,6 @@ from __future__ import annotations from collections.abc import Callable -import pickle import struct from pydantic import ValidationError @@ -29,6 +28,7 @@ from dimos.core.coordination.module_coordinator import _materialize_transports from dimos.core.transport import WebRTCTransport from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped +from dimos.protocol.encode.cbor import decode, encode from dimos.protocol.pubsub.impl.webrtc.providers import spec from dimos.protocol.pubsub.impl.webrtc.providers.broker import BrokerConfig from dimos.protocol.pubsub.impl.webrtc.providers.spec import ( @@ -191,7 +191,7 @@ def test_pickle_roundtrip_preserves_everything() -> None: """Transports are pickled into module worker processes; topic, type, and provider config must all survive.""" t1 = MockTransport("cmd_unreliable", FakeLCMMsg, name="pickled") - t2 = pickle.loads(pickle.dumps(t1)) + t2 = decode(encode(t1)) assert type(t2) is MockTransport assert t2.topic == t1.topic @@ -290,7 +290,7 @@ def test_transport_overrides_apply_and_survive_pickle() -> None: transport = _materialize_transports(bp, {"mock": {"name": "overridden"}})[("topic", FakeLCMMsg)] assert transport._config.name == "overridden" - assert pickle.loads(pickle.dumps(transport))._config.name == "overridden" + assert decode(encode(transport))._config.name == "overridden" # No override → config built from the spec's defaults. default = _materialize_transports(bp, {})[("topic", FakeLCMMsg)] diff --git a/dimos/protocol/pubsub/impl/zenohpubsub.py b/dimos/protocol/pubsub/impl/zenohpubsub.py index ffdba172091..b80d8bb202f 100644 --- a/dimos/protocol/pubsub/impl/zenohpubsub.py +++ b/dimos/protocol/pubsub/impl/zenohpubsub.py @@ -148,27 +148,6 @@ def __init__(self, **kwargs: Any) -> None: self._subscriber_lock = threading.Lock() self._stopped = False - def __getstate__(self): # type: ignore[no-untyped-def] - """Drop the live publishers, subscribers and locks on top of the session.""" - state = super().__getstate__() # type: ignore[no-untyped-call] - for key in ( - "_publishers", - "_publisher_lock", - "_subscribers", - "_drain_stops", - "_subscriber_lock", - ): - state.pop(key, None) - return state - - def __setstate__(self, state) -> None: # type: ignore[no-untyped-def] - super().__setstate__(state) - self._publishers = {} - self._publisher_lock = threading.Lock() - self._subscribers = [] - self._drain_stops = [] - self._subscriber_lock = threading.Lock() - def _get_publisher(self, key_expr: str, qos: ZenohQoS | None) -> zenoh.Publisher: """Get or declare the cached publisher for a key expression. diff --git a/dimos/protocol/pubsub/test_encoder.py b/dimos/protocol/pubsub/test_encoder.py index dec1e429726..9998e003035 100644 --- a/dimos/protocol/pubsub/test_encoder.py +++ b/dimos/protocol/pubsub/test_encoder.py @@ -17,6 +17,10 @@ import json from typing import Any +import numpy as np +import pytest + +from dimos.protocol.pubsub.encoders import CborEncoderMixin from dimos.protocol.pubsub.impl.memory import Memory, MemoryWithJSONEncoder @@ -169,3 +173,31 @@ def callback(message: Any, topic: str) -> None: # Verify it's actually JSON decoded_raw = json.loads(raw_message.decode("utf-8")) assert decoded_raw == original_message + + +class MemoryWithCborEncoder(CborEncoderMixin, Memory): # type: ignore[type-arg, misc] + pass + + +def test_cbor_encoded_pubsub_keeps_python_types() -> None: + pubsub = MemoryWithCborEncoder() + received: list[Any] = [] + pubsub.subscribe("cbor_topic", lambda message, topic: received.append(message)) + + pubsub.publish("cbor_topic", {"t": (1, 2), "a": np.arange(3, dtype=np.uint8)}) + + assert type(received[0]["t"]) is tuple + assert received[0]["a"].dtype == np.uint8 + + with pytest.raises(TypeError, match="Cannot CBOR-encode"): + pubsub.publish("cbor_topic", object()) + + +def test_cbor_encoded_pubsub_drops_undecodable_messages() -> None: + pubsub = MemoryWithCborEncoder() + received: list[Any] = [] + pubsub.subscribe("cbor_topic", lambda message, topic: received.append(message)) + + Memory.publish(pubsub, "cbor_topic", b"\xff\xff not cbor") # bypasses the encoder + + assert received == [] diff --git a/dimos/protocol/rpc/pubsubrpc.py b/dimos/protocol/rpc/pubsubrpc.py index 02a36a9f044..270b84f7930 100644 --- a/dimos/protocol/rpc/pubsubrpc.py +++ b/dimos/protocol/rpc/pubsubrpc.py @@ -28,7 +28,7 @@ ) from dimos.constants import LCM_MAX_CHANNEL_NAME_LENGTH -from dimos.protocol.pubsub.impl.lcmpubsub import PickleLCM, Topic +from dimos.protocol.pubsub.impl.lcmpubsub import CborLCM, Topic from dimos.protocol.pubsub.spec import PubSub from dimos.protocol.rpc.rpc_utils import deserialize_exception, serialize_exception from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, DEFAULT_RPC_TIMEOUTS, Args, RPCSpec @@ -87,31 +87,6 @@ def __init__( self._msg_id_counter = 0 self._msg_id_lock = threading.Lock() - def __getstate__(self) -> dict[str, Any]: - state: dict[str, Any] - if hasattr(super(), "__getstate__"): - state = super().__getstate__() # type: ignore[misc, assignment] - else: - state = self.__dict__.copy() - - # Exclude unpicklable attributes when serializing. - state.pop("_response_subs", None) - state.pop("_response_subs_lock", None) - state.pop("_msg_id_lock", None) - - return state - - def __setstate__(self, state: dict[str, Any]) -> None: - if hasattr(super(), "__setstate__"): - super().__setstate__(state) # type: ignore[misc] - else: - self.__dict__.update(state) - - # Restore unserializable attributes. - self._response_subs = {} - self._response_subs_lock = threading.RLock() - self._msg_id_lock = threading.Lock() - @abstractmethod def topicgen(self, name: str, req_or_res: bool) -> TopicT: ... @@ -215,16 +190,21 @@ def shared_response_handler(msg: MsgT, _: TopicT) -> None: _, callbacks_dict = self._response_subs[topic_res_key] callbacks_dict[msg_id] = cb - # Publish request - self.publish(topic_req, self._encodeRPCReq(req)) # type: ignore[arg-type] - - # Return unsubscribe function that removes this callback from the dict + # Unsubscribe function that removes this callback from the dict def unsubscribe_callback() -> None: with self._response_subs_lock: if topic_res_key in self._response_subs: _, callbacks_dict = self._response_subs[topic_res_key] callbacks_dict.pop(msg_id, None) + # Publish request; an unencodable argument raises to the caller and leaves + # nothing pending. + try: + self.publish(topic_req, self._encodeRPCReq(req)) # type: ignore[arg-type] + except Exception: + unsubscribe_callback() + raise + return unsubscribe_callback def call_nowait(self, name: str, arguments: Args) -> None: @@ -276,7 +256,7 @@ def execute_and_respond() -> None: return self.subscribe(topic_req, receive_call) -class LCMRPC(PubSubRPCMixin[Topic, Any], PickleLCM): +class LCMRPC(PubSubRPCMixin[Topic, Any], CborLCM): def __init__( self, rpc_timeouts: dict[str, float] | None = None, @@ -285,7 +265,7 @@ def __init__( ) -> None: if rpc_timeouts is None: rpc_timeouts = dict(DEFAULT_RPC_TIMEOUTS) - PickleLCM.__init__(self, **kwargs) + CborLCM.__init__(self, **kwargs) PubSubRPCMixin.__init__( self, rpc_timeouts=rpc_timeouts, default_rpc_timeout=default_rpc_timeout, **kwargs ) diff --git a/dimos/protocol/rpc/rpc_utils.py b/dimos/protocol/rpc/rpc_utils.py index 26ab281e451..3a08bcb4813 100644 --- a/dimos/protocol/rpc/rpc_utils.py +++ b/dimos/protocol/rpc/rpc_utils.py @@ -19,6 +19,8 @@ import traceback from typing import Any, TypedDict +from dimos.protocol.encode.cbor import encode + class SerializedException(TypedDict): """Type for serialized exception data.""" @@ -64,11 +66,20 @@ def serialize_exception(exc: Exception) -> SerializedException: return SerializedException( type_name=type(exc).__name__, type_module=type(exc).__module__, - args=exc.args, + args=tuple(_encodable(arg) for arg in exc.args), traceback=tb_str, ) +def _encodable(value: Any) -> Any: + """``value`` if it can cross the wire, else its str(); the traceback keeps the detail.""" + try: + encode(value) + except TypeError: + return str(value) + return value + + def deserialize_exception(exc_data: SerializedException) -> Exception: """Reconstruct an exception from serialized data. @@ -97,7 +108,9 @@ def deserialize_exception(exc_data: SerializedException) -> Exception: # Add remote traceback as __cause__ for context exc.__cause__ = RemoteError(type_name, type_module, args, tb_str) return exc # type: ignore[no-any-return] - except (AttributeError, TypeError): + except Exception: + # Args that could not cross the wire (an ExceptionGroup's members, say) + # arrived as strings; RemoteError below still carries the traceback. pass # Use RemoteError for non-builtin or if reconstruction failed diff --git a/dimos/protocol/rpc/spec.py b/dimos/protocol/rpc/spec.py index c749c13f52a..c125d89768c 100644 --- a/dimos/protocol/rpc/spec.py +++ b/dimos/protocol/rpc/spec.py @@ -13,7 +13,7 @@ # limitations under the License. import asyncio -from collections.abc import Callable +from collections.abc import Callable, Sequence import threading from types import MappingProxyType from typing import Any, Protocol, overload @@ -22,7 +22,7 @@ class Empty: ... -Args = tuple[list, dict[str, Any]] # type: ignore[type-arg] +Args = tuple[Sequence[Any], dict[str, Any]] # module that we can inspect for RPCs diff --git a/dimos/protocol/rpc/test_spec.py b/dimos/protocol/rpc/test_spec.py index 7607142afe2..f6e53bb26f3 100644 --- a/dimos/protocol/rpc/test_spec.py +++ b/dimos/protocol/rpc/test_spec.py @@ -19,12 +19,16 @@ import asyncio from collections.abc import Callable from contextlib import contextmanager +from dataclasses import dataclass +import enum import threading import time from typing import Any +import numpy as np import pytest +from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.protocol.rpc.pubsubrpc import LCMRPC from dimos.protocol.rpc.rpc_utils import RemoteError from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, RPCSpec @@ -382,3 +386,108 @@ def make_call(a, b) -> None: finally: unsub() + + +@dataclass +class Payload: + name: str + values: tuple[int, ...] + + +class Level(enum.IntEnum): + LOW = 1 + HIGH = 2 + + +def typed_echo(payload, level, array, vector): # type: ignore[no-untyped-def] + return payload, level, array, vector + + +def returns_lock(): # type: ignore[no-untyped-def] + return threading.Lock() + + +def raises_with_lock(): # type: ignore[no-untyped-def] + raise ValueError("bad", threading.Lock()) + + +@pytest.mark.parametrize("rpc_context, impl_name", testdata) +def test_typed_arguments_and_results_keep_their_types(rpc_context, impl_name: str) -> None: + with rpc_context() as (server, client): + unsub = server.serve_rpc(typed_echo, "typed_echo") + try: + args = ( + Payload("p", (1, 2)), + Level.HIGH, + np.arange(3, dtype=np.float32), + Vector3(1, 2, 3), + ) + result, _ = client.call_sync("typed_echo", (args, {}), rpc_timeout=5.0) + payload, level, array, vector = result + assert type(result) is tuple + assert payload == args[0] and type(payload) is Payload + assert level is Level.HIGH + assert array.dtype == np.float32 and np.array_equal(array, args[2]) + assert vector == args[3] and type(vector) is Vector3 + finally: + unsub() + + +@pytest.mark.parametrize("rpc_context, impl_name", testdata) +def test_unencodable_result_reaches_the_caller_as_type_error(rpc_context, impl_name: str) -> None: + with rpc_context() as (server, client): + unsubs = [ + server.serve_rpc(returns_lock, "returns_lock"), + server.serve_rpc(add_function, "add"), + ] + try: + with pytest.raises(TypeError, match="Cannot CBOR-encode"): + client.call_sync("returns_lock", ([], {}), rpc_timeout=5.0) + # The server keeps serving. + result, _ = client.call_sync("add", ([1, 2], {}), rpc_timeout=5.0) + assert result == 3 + finally: + for unsub in unsubs: + unsub() + + +@pytest.mark.parametrize("rpc_context, impl_name", testdata) +def test_unencodable_argument_fails_at_the_caller(rpc_context, impl_name: str) -> None: + with rpc_context() as (server, client): + unsub = server.serve_rpc(add_function, "add") + try: + with pytest.raises(TypeError, match="Cannot CBOR-encode"): + client.call_sync("add", ([lambda: None, 2], {}), rpc_timeout=5.0) + if impl_name == "zenoh": + assert client._pending == {} + finally: + unsub() + + +@pytest.mark.parametrize("rpc_context, impl_name", testdata) +def test_unencodable_exception_argument_arrives_as_text(rpc_context, impl_name: str) -> None: + with rpc_context() as (server, client): + unsub = server.serve_rpc(raises_with_lock, "raises_with_lock") + try: + with pytest.raises(ValueError) as info: + client.call_sync("raises_with_lock", ([], {}), rpc_timeout=5.0) + assert info.value.args[0] == "bad" + assert "lock" in info.value.args[1] + finally: + unsub() + + +def test_zenoh_reply_decode_failure_raises_instead_of_timing_out(monkeypatch) -> None: + def corrupt(_payload: bytes) -> Any: + raise ValueError("corrupt reply") + + with zenoh_rpc_context() as (server, client): + unsub = server.serve_rpc(add_function, "add") + try: + monkeypatch.setattr("dimos.protocol.rpc.zenohrpc.decode", corrupt) + started = time.monotonic() + with pytest.raises(ValueError, match="corrupt reply"): + client.call_sync("add", ([1, 2], {}), rpc_timeout=10.0) + assert time.monotonic() - started < 5.0 + finally: + unsub() diff --git a/dimos/protocol/rpc/zenohrpc.py b/dimos/protocol/rpc/zenohrpc.py index 2d42b143b43..224667b0fe2 100644 --- a/dimos/protocol/rpc/zenohrpc.py +++ b/dimos/protocol/rpc/zenohrpc.py @@ -16,13 +16,13 @@ from concurrent.futures import ThreadPoolExecutor from itertools import count -import pickle import time from typing import TYPE_CHECKING, Any import zenoh import zenoh.handlers +from dimos.protocol.encode.cbor import decode, encode from dimos.protocol.rpc.rpc_utils import deserialize_exception, serialize_exception from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, DEFAULT_RPC_TIMEOUTS, Args, RPCSpec from dimos.protocol.service.call_pool import CallPool @@ -37,6 +37,14 @@ EXCEPTION_ENCODING = zenoh.Encoding("dimos/rpc-exception") +def _decode_reply(payload: bytes, convert: Callable[[Any], Any]) -> Any: + """The decoded reply, or the failure itself so the caller raises instead of timing out.""" + try: + return convert(decode(payload)) + except Exception as e: + return e + + class ZenohRPC(RPCSpec, ZenohService): def __init__( self, @@ -54,22 +62,6 @@ def __init__( self._pending: dict[int, Callable[..., Any]] = {} self._call_counter = count(1) - def __getstate__(self): # type: ignore[no-untyped-def] - """Drop the live queryables, call pool and in-flight calls on top of the session. - - A module carries its RPC service into the worker by pickle, and none of - these survive the trip; the far side rebuilds them on first use. - """ - state = super().__getstate__() # type: ignore[no-untyped-call] - for key in ("_queryables", "_pending"): - state.pop(key, None) - return state - - def __setstate__(self, state) -> None: # type: ignore[no-untyped-def] - super().__setstate__(state) - self._queryables = [] - self._pending = {} - def _get_call_thread_pool(self) -> ThreadPoolExecutor: return self._call_pool.get() @@ -88,11 +80,12 @@ def call_cb(self, name: str, arguments: Args, cb: Callable[..., Any]) -> Callabl timeout = self.rpc_timeouts.get(name) or self.rpc_timeouts.get( method, self.default_rpc_timeout ) + # Encode before registering the callback: an unencodable argument raises + # to the caller and leaves nothing pending. + payload = encode(arguments) call_id = next(self._call_counter) self._pending[call_id] = cb - self._issue_query( - call_id, f"dimos/rpc/{name}", pickle.dumps(arguments), time.monotonic() + timeout - ) + self._issue_query(call_id, f"dimos/rpc/{name}", payload, time.monotonic() + timeout) def unsubscribe_callback() -> None: self._pending.pop(call_id, None) @@ -105,12 +98,12 @@ def on_reply(reply: zenoh.Reply) -> None: if err is None: cb = self._pending.pop(call_id, None) if cb is not None: - cb(pickle.loads(reply.ok.payload.to_bytes())) # type: ignore[union-attr] + cb(_decode_reply(reply.ok.payload.to_bytes(), lambda value: value)) # type: ignore[union-attr] return if err.encoding == EXCEPTION_ENCODING: cb = self._pending.pop(call_id, None) if cb is not None: - cb(deserialize_exception(pickle.loads(err.payload.to_bytes()))) + cb(_decode_reply(err.payload.to_bytes(), deserialize_exception)) return self._pending.pop(call_id, None) @@ -147,7 +140,7 @@ def call_nowait(self, name: str, arguments: Args) -> None: consolidation=zenoh.ConsolidationMode.NONE, congestion_control=zenoh.CongestionControl.BLOCK, timeout=timeout, - payload=pickle.dumps(arguments), + payload=encode(arguments), attachment=b"nowait", ) @@ -167,17 +160,18 @@ def unsubscribe() -> None: return unsubscribe def _execute_rpc(self, f: Callable[..., Any], name: str, query: zenoh.Query) -> None: - args = pickle.loads(query.payload.to_bytes()) # type: ignore[union-attr] nowait = query.attachment is not None try: + args = decode(query.payload.to_bytes()) # type: ignore[union-attr] response = f(*args[0], **args[1]) if not nowait: - query.reply(query.key_expr, pickle.dumps(response)) + query.reply(query.key_expr, encode(response)) except Exception as e: logger.exception(f"Exception in RPC handler for {name}: {e}", exc_info=e) if not nowait: - query.reply_err(pickle.dumps(serialize_exception(e)), encoding=EXCEPTION_ENCODING) - query.drop() + query.reply_err(encode(serialize_exception(e)), encoding=EXCEPTION_ENCODING) + finally: + query.drop() def stop(self) -> None: with self._call_pool.stop_scope(): diff --git a/dimos/protocol/service/call_pool.py b/dimos/protocol/service/call_pool.py index ae8183619e3..6240eaec44e 100644 --- a/dimos/protocol/service/call_pool.py +++ b/dimos/protocol/service/call_pool.py @@ -33,13 +33,6 @@ def __init__(self, max_workers: int) -> None: self._stop_calls = 0 self._stopped = False - def __getstate__(self) -> dict[str, int]: - # Only configuration crosses into a worker process. - return {"max_workers": self._max_workers} - - def __setstate__(self, state: dict[str, int]) -> None: - CallPool.__init__(self, state["max_workers"]) - def start(self) -> None: with self.lock: if self._stop_calls: diff --git a/dimos/protocol/service/lcmservice.py b/dimos/protocol/service/lcmservice.py index 3875cc64c03..d915fd90871 100644 --- a/dimos/protocol/service/lcmservice.py +++ b/dimos/protocol/service/lcmservice.py @@ -15,6 +15,7 @@ from __future__ import annotations from concurrent.futures import ThreadPoolExecutor +import copy import os import platform import threading @@ -88,28 +89,17 @@ def __init__(self, **kwargs: Any) -> None: self._thread = None self._call_pool = CallPool(max_workers=4) - def __getstate__(self): # type: ignore[no-untyped-def] - """Exclude unpicklable runtime attributes when serializing.""" - state = self.__dict__.copy() - # Remove unpicklable attributes - state.pop("l", None) - state.pop("_stop_event", None) - state.pop("_startup_complete", None) - state.pop("_thread", None) - state.pop("_l_lock", None) - state.pop("_pending_unsubs", None) - return state - - def __setstate__(self, state) -> None: # type: ignore[no-untyped-def] - """Restore object from pickled state.""" - self.__dict__.update(state) - # Reinitialize runtime attributes - self.l = None - self._stop_event = threading.Event() - self._startup_complete = threading.Event() - self._thread = None - self._l_lock = threading.Lock() - self._pending_unsubs = [] + # Services ride into workers as blueprint config (``pubsubs=[LCM(...)]``): only + # their settings cross, and a copy is a fresh, unstarted service. + def cbor_encode(self) -> dict[str, Any]: + return self.config.model_dump(exclude={"lcm"}) + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> LCMService: + return cls(**data) + + def __deepcopy__(self, memo: dict[int, Any]) -> LCMService: + return type(self).cbor_decode(copy.deepcopy(self.cbor_encode(), memo)) def start(self) -> None: """Start dispatch, or raise if a previous shutdown is still in progress. @@ -122,7 +112,7 @@ def start(self) -> None: raise RuntimeError("LCM service is still stopping; cannot restart yet") self._call_pool.start() if self._thread is None or not self._thread.is_alive(): - # Reinitialize LCM if it's None (e.g., after unpickling) + # Reinitialize LCM if it's None (e.g. after stop()) if self.l is None: # Any pending unsubscribes reference the destroyed instance. self._pending_unsubs.clear() diff --git a/dimos/protocol/service/system_configurator/base.py b/dimos/protocol/service/system_configurator/base.py index 8c9e2b47e7d..6f07e2653bb 100644 --- a/dimos/protocol/service/system_configurator/base.py +++ b/dimos/protocol/service/system_configurator/base.py @@ -17,6 +17,7 @@ from abc import ABC, abstractmethod import os import subprocess +from typing import Any from dimos.utils import prompt from dimos.utils.logging_config import setup_logger @@ -54,6 +55,16 @@ def _write_sysctl_int(name: str, value: int) -> None: class SystemConfigurator(ABC): critical: bool = False + # Configurators ride inside a Blueprint over RPC; plain attributes are all they hold. + def cbor_encode(self) -> dict[str, Any]: + return dict(vars(self)) + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> SystemConfigurator: + configurator = cls.__new__(cls) + configurator.__dict__.update(data) + return configurator + @abstractmethod def check(self) -> bool: """Return True if configured. Log errors and return False on uncertainty.""" diff --git a/dimos/protocol/service/test_call_pool.py b/dimos/protocol/service/test_call_pool.py deleted file mode 100644 index 9967e5ffe5b..00000000000 --- a/dimos/protocol/service/test_call_pool.py +++ /dev/null @@ -1,38 +0,0 @@ -# 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 contextlib import ExitStack -import pickle - -import pytest - -from dimos.protocol.service.call_pool import CallPool - - -@pytest.mark.parametrize("stopped", [False, True]) -def test_pickle_rebuilds_pool_without_live_workers_or_shutdown_state(stopped): - with ExitStack() as cleanup: - pool = CallPool(max_workers=2) - cleanup.callback(pool.shutdown) - assert pool.get().submit(lambda: 42).result(timeout=2.0) == 42 - if stopped: - pool.shutdown() - - clone = pickle.loads(pickle.dumps(pool)) - cleanup.callback(clone.shutdown) - - assert clone.get().submit(lambda: 43).result(timeout=2.0) == 43 - clone.shutdown() - with pytest.raises(RuntimeError, match="cannot accept calls"): - clone.get() diff --git a/dimos/protocol/service/test_lcmservice.py b/dimos/protocol/service/test_lcmservice.py index d04ad0ec679..854f740a66f 100644 --- a/dimos/protocol/service/test_lcmservice.py +++ b/dimos/protocol/service/test_lcmservice.py @@ -16,7 +16,6 @@ from contextlib import contextmanager import errno import os -import pickle import threading import time from unittest.mock import MagicMock, create_autospec, patch @@ -396,51 +395,6 @@ def test_start_and_stop(self) -> None: time.sleep(0.1) assert service._thread is None - def test_getstate_excludes_unpicklable_attrs(self) -> None: - with patch("dimos.protocol.service.lcmservice.lcm_mod.LCM") as mock_lcm_class: - mock_lcm_instance = create_autospec(LCM, spec_set=True, instance=True) - mock_lcm_class.return_value = mock_lcm_instance - - service = LCMService() - state = service.__getstate__() - - assert "l" not in state - assert "_stop_event" not in state - assert "_thread" not in state - assert "_l_lock" not in state - - def test_unpickling_reinitializes_runtime_attrs(self) -> None: - with patch("dimos.protocol.service.lcmservice.lcm_mod.LCM") as mock_lcm_class: - mock_lcm_instance = create_autospec(LCM, spec_set=True, instance=True) - mock_lcm_class.return_value = mock_lcm_instance - - service = LCMService() - new_service = pickle.loads(pickle.dumps(service)) - - assert new_service.l is None - assert isinstance(new_service._stop_event, threading.Event) - assert new_service._thread is None - # threading.Lock is a factory function, not a type - # Just check that the lock exists and has acquire/release methods - assert hasattr(new_service._l_lock, "acquire") - assert hasattr(new_service._l_lock, "release") - - def test_start_reinitializes_lcm_after_unpickling(self) -> None: - with patch("dimos.protocol.service.lcmservice.lcm_mod.LCM") as mock_lcm_class: - mock_lcm_instance = create_autospec(LCM, spec_set=True, instance=True) - mock_lcm_class.return_value = mock_lcm_instance - - service = LCMService() - new_service = pickle.loads(pickle.dumps(service)) - - # Start should reinitialize LCM - new_service.start() - - # LCM should be created again - assert mock_lcm_class.call_count == 2 - - new_service.stop() - def test_stop_cleans_up_lcm_instance(self) -> None: with patch("dimos.protocol.service.lcmservice.lcm_mod.LCM") as mock_lcm_class: mock_lcm_instance = create_autospec(LCM, spec_set=True, instance=True) diff --git a/dimos/protocol/service/test_zenohservice.py b/dimos/protocol/service/test_zenohservice.py index 103238aabb2..03784ec2dd6 100644 --- a/dimos/protocol/service/test_zenohservice.py +++ b/dimos/protocol/service/test_zenohservice.py @@ -17,7 +17,6 @@ from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError import json import os -import pickle import threading from types import SimpleNamespace from typing import Any @@ -25,8 +24,6 @@ import pytest import zenoh -from dimos.protocol.pubsub.impl.zenohpubsub import ZenohPubSubBase -from dimos.protocol.rpc.zenohrpc import ZenohRPC from dimos.protocol.service import zenohservice from dimos.protocol.service.zenohservice import ZenohConfig, ZenohService, ZenohSessionPool @@ -161,37 +158,6 @@ def test_different_modes_produce_different_keys() -> None: assert peer.session_key != client.session_key -def test_pickle_round_trip_sheds_the_session(session_pool) -> None: - """A module travels to its worker by pickle and re-acquires on the far side.""" - - svc = ZenohService(session_pool=session_pool) - svc.start() - clone = pickle.loads(pickle.dumps(svc)) - - assert clone._session is None - assert clone._session_pool is zenohservice.default_session_pool - assert clone.config.session_key == svc.config.session_key - - clone.start() - assert clone.session is not None - clone.stop() - - -def test_pickled_pubsub_and_rpc_rebuild_their_runtime(session_pool) -> None: - pubsub = pickle.loads(pickle.dumps(ZenohPubSubBase(session_pool=session_pool))) - assert pubsub._publishers == {} - assert pubsub._subscribers == [] - assert pubsub._drain_stops == [] - - rpc = pickle.loads(pickle.dumps(ZenohRPC(session_pool=session_pool))) - assert rpc._pending == {} - assert rpc._queryables == [] - try: - assert rpc._get_call_thread_pool().submit(lambda: 42).result(timeout=2.0) == 42 - finally: - rpc.stop() - - def test_start_creates_session(session_pool) -> None: svc = ZenohService(session_pool=session_pool) svc.start() diff --git a/dimos/protocol/service/zenohservice.py b/dimos/protocol/service/zenohservice.py index cc0f1b18ab9..20742c20144 100644 --- a/dimos/protocol/service/zenohservice.py +++ b/dimos/protocol/service/zenohservice.py @@ -15,6 +15,7 @@ from __future__ import annotations from concurrent.futures import Future +import copy from dataclasses import dataclass, field import json import os @@ -402,22 +403,18 @@ def __init__(self, *, session_pool: ZenohSessionPool | None = None, **kwargs: An self._session_pool = session_pool or default_session_pool self._session: zenoh.Session | None = None - def __getstate__(self): # type: ignore[no-untyped-def] - """Drop the live session, which pyo3 cannot pickle. + # Services ride into workers as blueprint config (``pubsubs=[Zenoh(...)]``): only + # their settings cross, the far side uses the default session pool, and a copy + # is a fresh, unstarted service. + def cbor_encode(self) -> dict[str, Any]: + return self.config.model_dump() - A module travels to its worker by pickle, so anything holding a session - has to shed it and re-acquire from the pool on the far side -- the same - move LCMService makes with its own runtime handles. - """ - state = self.__dict__.copy() - state.pop("_session", None) - state.pop("_session_pool", None) - return state - - def __setstate__(self, state) -> None: # type: ignore[no-untyped-def] - self.__dict__.update(state) - self._session_pool = default_session_pool - self._session = None + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> ZenohService: + return cls(**data) + + def __deepcopy__(self, memo: dict[int, Any]) -> ZenohService: + return type(self).cbor_decode(copy.deepcopy(self.cbor_encode(), memo)) def start(self) -> None: try: diff --git a/dimos/robot/assets/source.py b/dimos/robot/assets/source.py index c0bff41af60..09bf37201cd 100644 --- a/dimos/robot/assets/source.py +++ b/dimos/robot/assets/source.py @@ -17,7 +17,7 @@ from __future__ import annotations from pathlib import Path -from typing import SupportsIndex +from typing import Any, SupportsIndex from dimos.robot.assets.git_cache import GitAssetCache @@ -69,6 +69,13 @@ def __eq__(self, other: object) -> bool: def __hash__(self) -> int: return hash((self.url, self.ref)) + def cbor_encode(self) -> dict[str, str]: + return {"url": self.url, "ref": self.ref} + + @classmethod + def cbor_decode(cls, data: dict[str, str]) -> RobotDescriptionSource: + return cls(data["url"], data["ref"]) + def __reduce__(self) -> tuple[type[RobotDescriptionSource], tuple[str, str]]: """Recreate the source without serializing checkout state.""" return (RobotDescriptionSource, (self.url, self.ref)) @@ -116,6 +123,7 @@ def __getattribute__(self, name: str) -> object: "_resolve", "__reduce__", "__reduce_ex__", + "cbor_encode", }: return object.__getattribute__(self, name) @@ -164,6 +172,16 @@ def __hash__(self) -> int: ) ) + def cbor_encode(self) -> dict[str, Any]: + return { + "source": object.__getattribute__(self, "_robot_description_source"), + "relative_path": str(object.__getattribute__(self, "_robot_description_relative_path")), + } + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> RobotDescriptionPath: + return cls(data["source"], Path(data["relative_path"])) + def __reduce__(self) -> tuple[type[RobotDescriptionPath], tuple[RobotDescriptionSource, Path]]: """Preserve the lazy source handle across worker serialization.""" source: RobotDescriptionSource = object.__getattribute__(self, "_robot_description_source") diff --git a/dimos/robot/assets/test_model.py b/dimos/robot/assets/test_model.py index 3039ef468d8..97a3aab0958 100644 --- a/dimos/robot/assets/test_model.py +++ b/dimos/robot/assets/test_model.py @@ -14,12 +14,12 @@ from collections.abc import Callable from pathlib import Path -import pickle import xml.etree.ElementTree as ET import pytest from pytest_mock import MockerFixture +from dimos.protocol.encode.cbor import decode, encode import dimos.robot.assets.model as robot_model from dimos.utils.data import LfsPath @@ -703,7 +703,7 @@ def test_with_fixed_frame_rejects_invalid_model( model.load() -def test_pickle_keeps_model_edits_but_not_materialized_xml(tmp_path: Path) -> None: +def test_roundtrip_keeps_model_edits_but_not_materialized_xml(tmp_path: Path) -> None: urdf = tmp_path / "robot.urdf" urdf.write_text( "" @@ -721,7 +721,7 @@ def test_pickle_keeps_model_edits_but_not_materialized_xml(tmp_path: Path) -> No ) assert ET.fromstring(model.load().xml).get("name") == "before" - restored = pickle.loads(pickle.dumps(model)) + restored = decode(encode(model)) urdf.write_text(urdf.read_text().replace("name='before'", "name='after'")) root = ET.fromstring(restored.load().xml) diff --git a/dimos/robot/assets/test_source.py b/dimos/robot/assets/test_source.py index aa9246c0789..cb257e87bcd 100644 --- a/dimos/robot/assets/test_source.py +++ b/dimos/robot/assets/test_source.py @@ -14,10 +14,10 @@ import os from pathlib import Path -import pickle import pytest +from dimos.protocol.encode.cbor import decode, encode from dimos.robot.assets.git_cache import GitAssetCache from dimos.robot.assets.source import RobotDescriptionPath, RobotDescriptionSource @@ -117,13 +117,13 @@ def test_custom_source_handle_needs_no_registration(tmp_path: Path) -> None: ] -def test_lazy_path_survives_pickle_without_checkout( +def test_lazy_path_survives_the_round_trip_without_checkout( robot_source: tuple[RobotDescriptionSource, RecordingGitAssetCache], ) -> None: source, git_cache = robot_source model_path = source / "robots" / "testbot" / "model.urdf" - restored = pickle.loads(pickle.dumps(model_path)) + restored = decode(encode(model_path)) assert isinstance(restored, RobotDescriptionPath) assert restored.suffix == ".urdf" diff --git a/dimos/robot/test_blueprints_cross_processes.py b/dimos/robot/test_blueprints_cross_processes.py new file mode 100644 index 00000000000..3826a9acdc7 --- /dev/null +++ b/dimos/robot/test_blueprints_cross_processes.py @@ -0,0 +1,43 @@ +# 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. + +"""Every built-in blueprint must survive the RPC codec. + +`dimos shell` sends whole Blueprints to the daemon and the coordinator sends each +module's kwargs into its worker, so every config value a blueprint holds has to be +encodable (see dimos/protocol/encode/cbor.py). +""" + +import importlib + +import pytest + +from dimos.protocol.encode.cbor import decode, encode +from dimos.robot.all_blueprints import all_blueprints + + +@pytest.mark.parametrize("name", sorted(all_blueprints)) +def test_blueprint_roundtrips_through_the_codec(name: str) -> None: + module_path, attr = all_blueprints[name].split(":") + try: + blueprint = getattr(importlib.import_module(module_path), attr) + except Exception as error: # an optional dependency or robot setting this machine lacks + pytest.skip(f"{name} is not importable here: {error}") + + restored = decode(encode(blueprint)) + + assert [atom.module for atom in restored.blueprints] == [ + atom.module for atom in blueprint.blueprints + ] + assert restored.transport_map.keys() == blueprint.transport_map.keys() diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index ddf2bd8dd2e..f92c84629af 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -166,7 +166,7 @@ def _go2_rerun_blueprint() -> Blueprint: "world/navigation_costmap": _convert_navigation_costmap, "world/pgo_keyframes": _convert_pgo_keyframes, "world/pgo_loops": _convert_pgo_loops, - # partial, not a closure: this config is pickled to the viewer's worker + # partial, not a closure: this config is sent by reference to the viewer's worker "world/mapper_frame_ms": partial(_plot_telemetry, "mapper_frame_ms"), "world/pgo_loop_ms": partial(_plot_telemetry, "pgo_loop_ms"), "world/pgo_rebuild_ms": partial(_plot_telemetry, "pgo_rebuild_ms"), diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_joystick_record.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_joystick_record.py index 41f2baf74d4..4244a869cd2 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_joystick_record.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_joystick_record.py @@ -25,6 +25,7 @@ dimos --replay-db recordings//memory.db run replay """ +from dataclasses import dataclass, field from functools import partial from typing import Any @@ -47,11 +48,11 @@ def _plot_twist(name: str, twist: Any) -> Any: return [(f"plots/{name}", rr.Scalars([twist.linear.x, twist.linear.y, twist.angular.z]))] +@dataclass class _KeysText: """Joystick -> one text line per key change, e.g. ``W + SHIFT``.""" - def __init__(self) -> None: - self._last: list[int] | None = None + _last: list[int] | None = field(default=None, init=False) def __call__(self, joy: Any) -> Any: import rerun as rr diff --git a/dimos/utils/data.py b/dimos/utils/data.py index 257d3e881f3..5b5a6946b71 100644 --- a/dimos/utils/data.py +++ b/dimos/utils/data.py @@ -385,6 +385,7 @@ def __getattribute__(self, name: str) -> object: "_lfs_filename", "_lfs_resolved_cache", "_ensure_downloaded", + "cbor_encode", "__class__", "__copy__", "__deepcopy__", @@ -397,6 +398,13 @@ def __getattribute__(self, name: str) -> object: resolved = object.__getattribute__(self, "_ensure_downloaded")() return getattr(resolved, name) + def cbor_encode(self) -> str: + return str(object.__getattribute__(self, "_lfs_filename")) + + @classmethod + def cbor_decode(cls, filename: str) -> "LfsPath": + return cls(filename) + def __copy__(self) -> "LfsPath": copied = LfsPath(self._lfs_filename) object.__setattr__(copied, "_lfs_resolved_cache", self._lfs_resolved_cache) diff --git a/dimos/visualization/rerun/test_bridge_config.py b/dimos/visualization/rerun/test_bridge_config.py index e2218d9f243..3b69a66dbce 100644 --- a/dimos/visualization/rerun/test_bridge_config.py +++ b/dimos/visualization/rerun/test_bridge_config.py @@ -13,13 +13,13 @@ # limitations under the License. from pathlib import Path -import pickle from typing import Any import pytest from dimos.core.coordination.blueprint_config.errors import BlueprintConfigError from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser +from dimos.protocol.encode.cbor import decode, encode from dimos.simulation.scene_assets.spec import SceneMeshAlignment from dimos.visualization.rerun.bridge import Config, RerunBridgeModule from dimos.visualization.rerun.scene_package import SceneVisualFactory @@ -46,7 +46,7 @@ def test_parsed_bridge_config_preserves_callable_factories(tmp_path: Path) -> No parsed = BlueprintConfigParser(blueprint).parse(environ={}) # Worker deployment transports constructor kwargs over a multiprocessing pipe. - kwargs = pickle.loads(pickle.dumps(parsed.module_kwargs(RerunBridgeModule.name))) + kwargs = decode(encode(parsed.module_kwargs(RerunBridgeModule.name))) config = Config(**kwargs) assert config.visual_override == blueprint.blueprints[0].kwargs["visual_override"] diff --git a/dimos/visualization/rerun/test_scene_package.py b/dimos/visualization/rerun/test_scene_package.py index b5a5f185f27..ed68f070a49 100644 --- a/dimos/visualization/rerun/test_scene_package.py +++ b/dimos/visualization/rerun/test_scene_package.py @@ -13,11 +13,11 @@ # limitations under the License. from pathlib import Path -import pickle from typing import Any import pytest +from dimos.protocol.encode.cbor import encode from dimos.simulation.scene_assets.spec import SceneMeshAlignment, ScenePackage from dimos.visualization.rerun.scene_package import ( resolve_scene_package_for_rerun, @@ -64,7 +64,7 @@ def test_resolves_composed_mjb_to_scene_package_visual(tmp_path: Path) -> None: static_entities = scene_package_static_entities(composed_path) assert set(static_entities) == {"world/scene"} - pickle.dumps(static_entities) + encode(static_entities) transform, asset = static_entities["world/scene"](FakeRerun) assert transform.kwargs == { diff --git a/dimos/web/cockpit.py b/dimos/web/cockpit.py index 7c82897db49..ecefb7c9ebd 100644 --- a/dimos/web/cockpit.py +++ b/dimos/web/cockpit.py @@ -88,6 +88,13 @@ class _FrozenDict(dict[str, Any]): def __reduce__(self) -> tuple[Any, ...]: return (_FrozenDict, (dict(self),)) + def cbor_encode(self) -> dict[str, Any]: + return dict(self) + + @classmethod + def cbor_decode(cls, data: dict[str, Any]) -> _FrozenDict: + return cls(data) + def __setitem__(self, key: str, value: Any) -> None: raise TypeError("Channel params are immutable") diff --git a/dimos/web/codecs.py b/dimos/web/codecs.py index e0129908846..4547e8d1ce0 100644 --- a/dimos/web/codecs.py +++ b/dimos/web/codecs.py @@ -29,7 +29,7 @@ def decode_goal(value: dict[str, Any], context: PublishContext) -> PoseStamped: compiler resolves encoding ids to callables at blueprint definition time and ships the resolved callable inside an immutable runtime channel spec, so a worker never depends on the user's codec module having been imported there. -Codec functions must therefore live at module level (standard pickle ships +Codec functions must therefore live at module level (the worker pipe sends them by reference); built-ins register in dimos/web/relay_bridge/builtin_codecs.py. @@ -167,18 +167,18 @@ def _describe(func: Callable[..., Any]) -> str: def _check_registrable(func: Callable[..., Any], label: str) -> None: """Decoration-time importability gate. - The full pickle-by-reference walk (`_resolves_by_reference`) cannot run + The full by-reference walk (`_resolves_by_reference`) cannot run here: a decorator fires before the def statement binds the module attribute. resolve_encoder() completes the check once the module exists. """ if not inspect.isfunction(func): raise TypeError(f"{label} must be a module-level function, got {func!r}") if func.__name__ == "": - raise ValueError(f"{label} cannot be a lambda; standard pickle ships codecs by reference") + raise ValueError(f"{label} cannot be a lambda; the worker pipe sends codecs by reference") if "" in func.__qualname__: raise ValueError( f"{label} {_describe(func)} is defined inside a function; move it to " - "module level so it can be pickled by reference" + "module level so it can be sent by reference" ) if func.__module__ == "__main__": raise ValueError( @@ -188,7 +188,7 @@ def _check_registrable(func: Callable[..., Any], label: str) -> None: def _resolves_by_reference(func: Callable[..., Any]) -> bool: - """Mirror pickle's save-by-reference lookup for a function.""" + """Mirror the CBOR codec's by-reference lookup for a function.""" obj: Any = sys.modules.get(func.__module__) for part in func.__qualname__.split("."): obj = getattr(obj, part, None) @@ -389,7 +389,7 @@ def decode_json_v1(value: Any) -> Any: def resolve_decoder(encoding: str, message_type: type[Any]) -> DecoderDef: """The decoder a publish channel (encoding, message type) compiles to; ValueError when the pair is unsupported. Runs in the parent at blueprint - authoring time, so this also finishes the pickle-by-reference check + authoring time, so this also finishes the by-reference check deferred from decoration (mirrors resolve_encoder).""" definition = _decoders.get(encoding) if definition is not None: @@ -401,7 +401,7 @@ def resolve_decoder(encoding: str, message_type: type[Any]) -> DecoderDef: if not _resolves_by_reference(definition.decode): raise ValueError( f"decoder {_describe(definition.decode)} for {encoding!r} cannot be " - "pickled by reference; it must be importable under its own module " + "sent by reference; it must be importable under its own module " "and qualified name" ) return definition @@ -430,7 +430,7 @@ def is_generic_lcm_encoding(encoding: str) -> bool: def resolve_encoder(encoding: str, message_type: type[Any]) -> EncoderDef: """The encoder a channel (encoding, message type) compiles to; ValueError when the pair is unsupported. Runs in the parent at blueprint authoring - time, so this also finishes the pickle-by-reference check deferred from + time, so this also finishes the by-reference check deferred from decoration.""" definition = _encoders.get(encoding) if definition is not None: @@ -442,7 +442,7 @@ def resolve_encoder(encoding: str, message_type: type[Any]) -> EncoderDef: if not _resolves_by_reference(definition.encode): raise ValueError( f"encoder {_describe(definition.encode)} for {encoding!r} cannot be " - "pickled by reference; it must be importable under its own module " + "sent by reference; it must be importable under its own module " "and qualified name" ) return definition diff --git a/dimos/web/relay_bridge/channels.py b/dimos/web/relay_bridge/channels.py index 07a45faf54f..42105c07b92 100644 --- a/dimos/web/relay_bridge/channels.py +++ b/dimos/web/relay_bridge/channels.py @@ -34,7 +34,7 @@ class RuntimeChannelSpec: """One channel's immutable runtime contract. Compiled by cockpit() in the parent (codecs resolved from the registry, - ready to pickle by reference into the worker) or resolved from the + ready to send by reference into the worker) or resolved from the manifest against BUILTIN_CHANNELS at module start. The bridge's rx encode behavior and its tx publish decode behavior are driven entirely by these specs. diff --git a/dimos/web/relay_bridge/dynamic.py b/dimos/web/relay_bridge/dynamic.py index e6f84ce752d..8dcf15af668 100644 --- a/dimos/web/relay_bridge/dynamic.py +++ b/dimos/web/relay_bridge/dynamic.py @@ -12,22 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Pickle-safe dynamically generated RelayBridgeModule subclasses. +"""Dynamically generated RelayBridgeModule subclasses that survive the worker pipe. DimOS discovers a module's streams from class annotations and ships the class -object itself into forkserver workers (DeployModuleRequest) and back out again -(Actor.__reduce__), all over standard pickle. A class built at runtime has no -importable module attribute for pickle's save-by-reference path, so -`make_relay_bridge_class` builds classes through a dedicated metaclass -registered with copyreg: pickle consults copyreg's dispatch table before its -class special-case, serializes the canonical port specs, and reconstructs -through the memoized factory on load. A metaclass `__reduce__` would not work; -pickle never calls it for class objects. +object itself into forkserver workers (DeployModuleRequest) over the CBOR +codec, which sends classes by importable name. A class built at runtime has no +such name, so `make_relay_bridge_class` builds classes through a dedicated +metaclass whose `cbor_encode`/`cbor_decode` serialize the canonical port specs +and reconstruct through the memoized factory on the far side. """ from abc import ABCMeta -from collections.abc import Callable, Sequence -import copyreg +from collections.abc import Sequence from dataclasses import dataclass import hashlib import keyword @@ -60,12 +56,22 @@ class _DynamicRelayBridgeMeta(ABCMeta): """Marker metaclass for generated bridge classes. Derives from ABCMeta because that is Module's metaclass (a bare `type` - metaclass would conflict). copyreg dispatches on the exact metaclass, so - ordinary ABCMeta classes are unaffected by the reducer below. + metaclass would conflict). The CBOR codec looks for the hooks on a class + object's metaclass, so ordinary ABCMeta classes are unaffected. """ dynamic_port_specs: tuple[DynamicPortSpec, ...] + def cbor_encode(cls) -> tuple[tuple[str, type[Any], str], ...]: + # Plain field tuples, not DynamicPortSpec instances: instances would pin + # the payload to one generation of the dataclass, so re-encoding a class + # that survived an importlib.reload of this module would fail. + return _spec_fields(cls.dynamic_port_specs) + + @classmethod + def cbor_decode(cls, fields: Sequence[Sequence[Any]]) -> type[RelayBridgeModule]: + return _rebuild_relay_bridge_class(tuple(tuple(f) for f in fields)) + _CLASS_NAME_PREFIX = "DynamicRelayBridge_" _DIGEST_LEN = 10 @@ -79,7 +85,7 @@ class _DynamicRelayBridgeMeta(ABCMeta): | {"ref", "rpc", "encoded"} ) -# Process-global on purpose: repeated unpickles of the same class must return +# Process-global on purpose: repeated decodes of the same class must return # the identical class object, because blueprints and the coordinator compare # module classes with `is`. _memo: dict[tuple[DynamicPortSpec, ...], type[RelayBridgeModule]] = {} @@ -87,7 +93,7 @@ class _DynamicRelayBridgeMeta(ABCMeta): def _resolves_by_reference(message_type: type[Any]) -> bool: - """Mirror pickle's save-by-reference lookup for a class.""" + """Mirror the CBOR codec's by-reference lookup for a class.""" obj: Any = sys.modules.get(message_type.__module__) for part in message_type.__qualname__.split("."): obj = getattr(obj, part, None) @@ -142,12 +148,12 @@ def _validate_specs(specs: tuple[DynamicPortSpec, ...]) -> None: if "" in message_type.__qualname__: raise ValueError( f"stream {stream!r}: message type {message_type.__qualname__} is defined " - "inside a function; move it to module level so it can be pickled by reference" + "inside a function; move it to module level so it can be sent by reference" ) if not _resolves_by_reference(message_type): raise ValueError( f"stream {stream!r}: message type {message_type.__module__}." - f"{message_type.__qualname__} cannot be pickled by reference; it must be " + f"{message_type.__qualname__} cannot be sent by reference; it must be " "importable under its own module and qualified name" ) seen: set[str] = set() @@ -238,22 +244,3 @@ def _rebuild_relay_bridge_class( for stream, message_type, direction in fields ] ) - - -def _reduce_dynamic_class( - cls: _DynamicRelayBridgeMeta, -) -> tuple[Callable[..., Any], tuple[tuple[tuple[str, type[Any], str], ...]]]: - # Plain field tuples, not DynamicPortSpec instances: instances would pin - # the pickle to one generation of the dataclass, so re-pickling a class - # that survived an importlib.reload of this module would fail with - # "it's not the same object as ...DynamicPortSpec". - return (_rebuild_relay_bridge_class, (_spec_fields(cls.dynamic_port_specs),)) - - -# Registered at import time on purpose: the reducer must be active in every -# process that pickles a generated class (workers re-pickle it through -# Actor.__reduce__), and unpickling resolves _rebuild_relay_bridge_class by -# importing this module, so the import itself activates the reducer wherever -# the class travels. copyreg.pickle overwrites the same dispatch-table slot, -# making repeated registration (including importlib.reload) idempotent. -copyreg.pickle(_DynamicRelayBridgeMeta, _reduce_dynamic_class) diff --git a/dimos/web/relay_bridge/test_dynamic.py b/dimos/web/relay_bridge/test_dynamic.py index a28566b5cbc..0c5cc2ac38d 100644 --- a/dimos/web/relay_bridge/test_dynamic.py +++ b/dimos/web/relay_bridge/test_dynamic.py @@ -12,14 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for pickle-safe dynamically generated relay bridge classes.""" +"""Tests for dynamically generated relay bridge classes that cross the worker pipe.""" # No `from __future__ import annotations` here on purpose: under PEP 563 a # Module class whose annotations become strings silently loses its streams # (see the module-level note in test_relay_bridge_module.py). from dataclasses import dataclass -import pickle import re import subprocess import sys @@ -38,6 +37,7 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid from dimos.msgs.sensor_msgs.Image import Image +from dimos.protocol.encode.cbor import decode, encode from dimos.web.cockpit import Channel, cockpit from dimos.web.relay_bridge import dynamic from dimos.web.relay_bridge.config import RelayBridgeConfig @@ -129,9 +129,9 @@ def test_memoization_and_canonical_order(): assert generated.dynamic_port_specs == tuple(sorted(_MAIN_SPECS, key=lambda s: s.stream)) -def test_pickle_roundtrip_returns_identical_class_object(): +def test_cbor_roundtrip_returns_identical_class_object(): generated = make_relay_bridge_class(_MAIN_SPECS) - assert pickle.loads(pickle.dumps(generated)) is generated + assert decode(encode(generated)) is generated def test_two_shapes_coexist_in_one_process(): @@ -139,8 +139,8 @@ def test_two_shapes_coexist_in_one_process(): shape_b = make_relay_bridge_class([DynamicPortSpec("marker_feed", _MarkerB, "rx")]) assert shape_a is not shape_b assert shape_a.__name__ != shape_b.__name__ - assert pickle.loads(pickle.dumps(shape_a)) is shape_a - assert pickle.loads(pickle.dumps(shape_b)) is shape_b + assert decode(encode(shape_a)) is shape_a + assert decode(encode(shape_b)) is shape_b def test_generated_class_is_published_for_reload_lookup(): @@ -169,24 +169,25 @@ def test_digest_name_collision_with_different_specs_raises(monkeypatch): _RELOAD_SCRIPT = """ -import importlib, pickle, sys +import importlib, sys from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.protocol.encode.cbor import decode, encode import dimos.web.relay_bridge.dynamic as dynamic C = dynamic.make_relay_bridge_class([dynamic.DynamicPortSpec("reload_feed", Vector3, "rx")]) -data = pickle.dumps(C) +data = encode(C) mod = importlib.reload(sys.modules["dimos.web.relay_bridge.dynamic"]) assert getattr(mod, C.__name__) is C, "restart getattr lost the class" -assert pickle.loads(data) is C, "pre-reload pickle lost identity" -assert pickle.loads(pickle.dumps(C)) is C, "re-pickle after reload failed" +assert decode(data) is C, "pre-reload payload lost identity" +assert decode(encode(C)) is C, "re-encode after reload failed" again = mod.make_relay_bridge_class([mod.DynamicPortSpec("reload_feed", Vector3, "rx")]) assert again is C, "factory lost identity after reload" """ -def test_reload_then_repickle_preserves_identity(): +def test_reload_then_reencode_preserves_identity(): """The coordinator's restart path reloads dynamic.py and then re-deploys - the surviving class object, which re-pickles it. Runs in a subprocess so + the surviving class object, which re-encodes it. Runs in a subprocess so the reload cannot poison this process's module state.""" result = subprocess.run( [sys.executable, "-c", _RELOAD_SCRIPT], capture_output=True, text=True, timeout=120 @@ -236,7 +237,7 @@ class LocalType: def test_unimportable_message_type_rejected(): ghost = type("GhostType", (), {"__module__": __name__}) - with pytest.raises(ValueError, match="pickled by reference"): + with pytest.raises(ValueError, match="sent by reference"): make_relay_bridge_class([DynamicPortSpec("ghost_feed", ghost, "rx")]) @@ -275,12 +276,12 @@ def test_autoconnect_matches_generated_ports_with_static_modules(): ] -def test_blueprint_pickle_preserves_generated_class_identity(): +def test_blueprint_roundtrip_preserves_generated_class_identity(): generated = make_relay_bridge_class(_AUTOCONNECT_SPECS) blueprint = autoconnect( _PingProducer.blueprint(), generated.blueprint(), _NoteConsumer.blueprint() ) - restored = pickle.loads(pickle.dumps(blueprint)) + restored = decode(encode(blueprint)) assert restored == blueprint restored_atom = next(a for a in restored.blueprints if a.name == generated.name) assert restored_atom.module is generated @@ -304,7 +305,7 @@ def test_forkserver_worker_predates_class_deploy_rpc_and_stream( worker_manager, deployed_proxies, ping_transport, retry_until ): """The W5 flagship proof: the worker process exists before the class does, - so the class can only arrive in the worker through pickle, not fork + so the class can only arrive in the worker through the pipe, not fork inheritance. Deploy, call an RPC, wire a generated stream, see data.""" assert worker_manager.workers[0].pid is not None diff --git a/dimos/web/relay_bridge/test_relay_bridge_authoring.py b/dimos/web/relay_bridge/test_relay_bridge_authoring.py index 8b59103dc61..3e9e64e79cf 100644 --- a/dimos/web/relay_bridge/test_relay_bridge_authoring.py +++ b/dimos/web/relay_bridge/test_relay_bridge_authoring.py @@ -23,7 +23,6 @@ import base64 from dataclasses import asdict, replace import json -import pickle import struct from typing import Any @@ -43,6 +42,7 @@ from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid from dimos.msgs.nav_msgs.Path import Path as NavPath from dimos.msgs.sensor_msgs.Image import Image +from dimos.protocol.encode.cbor import decode, encode from dimos.web.cockpit import Channel, Chat, Map2D, Stats, Video, cockpit from dimos.web.codecs import EncodedPayload, PublishContext, web_decoder, web_encoder from dimos.web.relay_bridge import builtin_codecs, pacing, relay_bridge_module @@ -164,7 +164,7 @@ def _encode_boom(msg: NavPath) -> bytes: _NAV_PATH_PAYLOAD = struct.pack(" None: +def test_config_accepts_runtime_specs_by_identity_and_crosses_the_pipe() -> None: spec = RuntimeChannelSpec( ch="odom", message_type=PoseStamped, @@ -179,7 +179,7 @@ def test_config_accepts_runtime_specs_by_identity_and_pickles() -> None: # Pydantic must pass the frozen dataclass through untouched; a rebuilt # copy would break the is-checks blueprints rely on. assert config.channels is not None and config.channels[0] is spec - restored = pickle.loads(pickle.dumps(config)) + restored = decode(encode(config)) assert restored.channels[0].encoder is builtin_codecs.encode_pose assert restored.channels[0].params == {"a": 1} diff --git a/dimos/web/test_cockpit.py b/dimos/web/test_cockpit.py index 5c45a4184a4..d8e6db34fbb 100644 --- a/dimos/web/test_cockpit.py +++ b/dimos/web/test_cockpit.py @@ -15,7 +15,6 @@ """Authoring-API tests: cockpit()/panels/layout compile to pinned manifests.""" from dataclasses import dataclass -import pickle import struct import subprocess import sys @@ -44,6 +43,7 @@ from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.protocol.encode.cbor import decode, encode from dimos.web.cockpit import ( Channel, ChannelRequest, @@ -169,11 +169,11 @@ def test_default_preset() -> None: assert manifest == GO2_MANIFEST -def test_blueprint_pickles() -> None: +def test_blueprint_crosses_the_worker_pipe() -> None: # Blueprint kwargs cross the forkserver Pipe; the manifest dict must - # survive a pickle round-trip unchanged. + # survive the round trip unchanged. blueprint = cockpit(layout=GO2_LAYOUT) - restored = pickle.loads(pickle.dumps(blueprint)) + restored = decode(encode(blueprint)) assert manifest_of(restored) == GO2_MANIFEST @@ -415,7 +415,7 @@ def test_channel_params_are_immutable() -> None: channel.params["nested"]["a"] = [] # Nested sequences freeze to tuples: no append surface at all. assert channel.params["nested"]["a"] == (1, 2) - restored = pickle.loads(pickle.dumps(channel)) + restored = decode(encode(channel)) assert restored.params == channel.params with pytest.raises(TypeError, match="immutable"): restored.params["scale"] = 10 @@ -474,8 +474,8 @@ def test_channels_only_blueprint() -> None: assert specs[0].encoder_takes_params is False assert specs[1].encoder is encode_json_v1 and specs[1].message_type is _OpsNote # Blueprint kwargs cross the forkserver Pipe: class identity and the - # by-reference encoder must survive pickling. - restored = pickle.loads(pickle.dumps(blueprint)) + # by-reference encoder must survive the trip. + restored = decode(encode(blueprint)) (ratom,) = restored.blueprints assert ratom.module is atom.module assert ratom.kwargs["channels"][0].encoder is _encode_path_xy @@ -512,8 +512,8 @@ def test_publish_tx_channel_blueprint() -> None: assert spec.decoder is decode_text and spec.decoder_takes_context is False assert spec.encoder is None # Blueprint kwargs cross the forkserver Pipe: the by-reference decoder - # must survive pickling. - restored = pickle.loads(pickle.dumps(blueprint)) + # must survive the trip. + restored = decode(encode(blueprint)) (ratom,) = restored.blueprints assert ratom.module is atom.module assert ratom.kwargs["channels"][0].decoder is decode_text @@ -549,7 +549,7 @@ def test_chat_panel_blueprint() -> None: assert specs["human_input"].decoder is decode_text assert specs["audio_in"].decoder is decode_audio_chunk assert not specs["audio_in"].decoder_takes_context and specs["audio_in"].encoder is None - restored = pickle.loads(pickle.dumps(blueprint)) + restored = decode(encode(blueprint)) (ratom,) = restored.blueprints assert {s.ch: s.decoder for s in ratom.kwargs["channels"]}["audio_in"] is decode_audio_chunk # Pacing is the chat panel's, not the stream's: the same streams declared @@ -603,7 +603,7 @@ def test_map2d_nav_channels_blueprint() -> None: assert specs["path"].paced and specs["path"].resend_on_subscribe assert specs["clicked_point"].decoder is decode_point assert specs["stop_movement"].decoder is decode_bool - restored = pickle.loads(pickle.dumps(blueprint)) + restored = decode(encode(blueprint)) (ratom,) = restored.blueprints assert {s.ch: s.decoder for s in ratom.kwargs["channels"]}["clicked_point"] is decode_point @@ -630,7 +630,7 @@ def test_map3d_blueprint() -> None: spec = next(s for s in atom.kwargs["channels"] if s.ch == "global_map") assert spec.encoder is encode_voxels assert spec.resend_on_subscribe and not spec.paced - restored = pickle.loads(pickle.dumps(blueprint)) + restored = decode(encode(blueprint)) (ratom,) = restored.blueprints assert {s.ch: s.encoder for s in ratom.kwargs["channels"]}["global_map"] is encode_voxels @@ -775,7 +775,7 @@ def test_stats_panel_blueprint() -> None: assert ports[("resource_stats", "in")] is dict spec = next(s for s in atom.kwargs["channels"] if s.ch == "resource_stats") assert spec.encoder is encode_stats and not spec.paced - restored = pickle.loads(pickle.dumps(blueprint)) + restored = decode(encode(blueprint)) (ratom,) = restored.blueprints assert {s.ch: s.encoder for s in ratom.kwargs["channels"]}["resource_stats"] is encode_stats @@ -941,7 +941,7 @@ def test_lcm_channel_blueprint() -> None: assert any( s.name == "pose" and s.type is PoseStamped and s.direction == "in" for s in atom.streams ) - restored = pickle.loads(pickle.dumps(blueprint)) + restored = decode(encode(blueprint)) assert restored.blueprints[0].kwargs["channels"][0].encoder is encode_lcm_v1 diff --git a/dimos/web/test_codecs.py b/dimos/web/test_codecs.py index d97d6f03e0b..6f5225a0c20 100644 --- a/dimos/web/test_codecs.py +++ b/dimos/web/test_codecs.py @@ -345,13 +345,13 @@ def test_resolve_unknown_encoding() -> None: resolve_encoder("nope.v1", dict) -def test_resolve_rechecks_pickle_by_reference(monkeypatch: pytest.MonkeyPatch) -> None: +def test_resolve_rechecks_by_reference(monkeypatch: pytest.MonkeyPatch) -> None: # Decoration cannot verify the module binding (the def statement has not # bound the name yet); resolve does. Simulate a codec whose module-level # name was deleted after registration. web_encoder("t.enc.unref.v1")(_enc_ok) monkeypatch.delattr(sys.modules[__name__], "_enc_ok") - with pytest.raises(ValueError, match="cannot be pickled by reference"): + with pytest.raises(ValueError, match="cannot be sent by reference"): resolve_encoder("t.enc.unref.v1", _Point) @@ -377,10 +377,10 @@ def test_resolve_decoder_rejections() -> None: resolve_decoder("json.v1", _Point) -def test_resolve_decoder_rechecks_pickle_by_reference(monkeypatch: pytest.MonkeyPatch) -> None: +def test_resolve_decoder_rechecks_by_reference(monkeypatch: pytest.MonkeyPatch) -> None: web_decoder("t.dec.unref.v1")(_dec_ok) monkeypatch.delattr(sys.modules[__name__], "_dec_ok") - with pytest.raises(ValueError, match="cannot be pickled by reference"): + with pytest.raises(ValueError, match="cannot be sent by reference"): resolve_decoder("t.dec.unref.v1", _Point) diff --git a/docs/coding-agents/code-quality-rules.md b/docs/coding-agents/code-quality-rules.md index 1cad8b4e24e..90c279f8990 100644 --- a/docs/coding-agents/code-quality-rules.md +++ b/docs/coding-agents/code-quality-rules.md @@ -15,7 +15,7 @@ Rules dimos code is expected to follow. They address recurring issues found in c * Specify only what differs from defaults. Don't restate defaults like `tick_rate=100.0`, `publish_joint_state=True`, or default topics (`/cmd_vel`, `/odom`). * `.transports({...})` applies to all matching modules, so define a remap once, not twice across sub-blueprints. -* No lambdas -- they can't be pickled to worker processes. Use named functions. +* No lambdas -- they can't be sent to worker processes (functions travel by import path). Use module-level named functions. * Do no work at import time: no subprocesses, viewers, model parsing, or network. In particular don't call `get_data(...)` (it blocks import until the download finishes) -- use `LfsPath` (resolved at access time) or build the config in `start`/`build`. Any process you start must be managed (shut down when not needed). * Blueprint files define blueprints, not modules/classes. * Helper blueprints not meant to run alone must start with `_` (the in-repo `all_blueprints.py` generator skips them); demo/non-shared built-in ones get a `demo_` prefix (hidden from `dimos list`). Externally packaged blueprints are not added to `all_blueprints.py`; expose them with `dimos.blueprints` Python package entry points. diff --git a/docs/usage/distributed/index.md b/docs/usage/distributed/index.md index 305cd189c00..439449749d7 100644 --- a/docs/usage/distributed/index.md +++ b/docs/usage/distributed/index.md @@ -4,9 +4,9 @@ Run one blueprint across several machines: the robot's drivers on its Jetson, a ## WARNING: dimos host is essentially remote code execution -Dimos RPC is pickle for now, and a fragment is a pickled blueprint. Anyone who can reach a -Host's router can make it run any code they like, as the user the Host runs as. There is no -authentication or encryption yet. +Dimos RPC sends classes and functions by import path, and a fragment is an encoded blueprint. +Anyone who can reach a Host's router can make it import and run any installed code, as the user +the Host runs as. There is no authentication or encryption yet. - A Host listens only on its own machine until you open it with `listen=`. - Open it only on networks you trust: ideally a VPN such as Tailscale diff --git a/docs/usage/python-api.md b/docs/usage/python-api.md index 26fafb46f0c..27e3d68509a 100644 --- a/docs/usage/python-api.md +++ b/docs/usage/python-api.md @@ -203,9 +203,9 @@ app.unload(SomeModule) # stop and remove it; the rest keeps running ``` Strings and registered Module classes take a name-based fast path. Other -Module classes and `Blueprint` objects are pickled and unpickled on the +Module classes and `Blueprint` objects are encoded and rebuilt on the daemon, so their module classes must be importable there and all kwargs must -be picklable. +be encodable by the RPC codec ([cbor.py](/dimos/protocol/encode/cbor.py)). ## Limitations diff --git a/docs/usage/transports/index.md b/docs/usage/transports/index.md index a07032b7d95..8a594049900 100644 --- a/docs/usage/transports/index.md +++ b/docs/usage/transports/index.md @@ -446,6 +446,7 @@ Transports often need to serialize messages before sending and deserialize after | Mixin | Encoding | Use case | |----------------------|-----------------|------------------------------------| | `PickleEncoderMixin` | Python pickle | Any Python object, Python-only | +| `CborEncoderMixin` | CBOR | Python values with their types kept (tuples, numpy, enums, dataclasses, messages); what RPC uses | | `LCMEncoderMixin` | LCM binary | Cross-language (C/C++/Python/Go/…) | | `JpegEncoderMixin` | JPEG compressed | Image data, reduces bandwidth | diff --git a/pyproject.toml b/pyproject.toml index fecabfbf166..be358c7a9fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,6 +122,7 @@ dependencies = [ "sortedcontainers==2.4.0", "pydantic", "pydantic-core", # blueprint_config, sensor_msgs/CameraInfo + "cbor2>=6.1.5,<7", # RPC and worker-pipe serialization (protocol/encode/cbor.py) "python-dotenv", "packaging>=24.0", "requests>=2.32", diff --git a/uv.lock b/uv.lock index 23fff2ea94a..c5450c58cf3 100644 --- a/uv.lock +++ b/uv.lock @@ -741,6 +741,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d8/2b/a40e1488fdfa02d3f9a653a61a5935ea08b3c2225ee818db6a76c7ba9695/cattrs-25.3.0-py3-none-any.whl", hash = "sha256:9896e84e0a5bf723bc7b4b68f4481785367ce07a8a02e7e9ee6eb2819bc306ff", size = 70738, upload-time = "2025-10-07T12:26:06.603Z" }, ] +[[package]] +name = "cbor2" +version = "6.1.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/39/34/d443914ea562a985ccb357682e17b7190d5d58eff797c741379be47a8f31/cbor2-6.1.5.tar.gz", hash = "sha256:6eb06160c42315ac0c4ded461c7d84d92fa18c69d13d17fc1dfc1fae96580c95", size = 94232, upload-time = "2026-10-01T18:09:33.621Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/08/8bb3abca3820c20cd5efa51f0f37033f8bc514b4d6f38afb559257a4b17d/cbor2-6.1.5-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:519f3f0d0d9467091c678f4a19a31e1b8756c10bbd6294cb3f906092f3da1597", size = 416900, upload-time = "2026-10-01T18:07:47.646Z" }, + { url = "https://files.pythonhosted.org/packages/89/7a/39d6a60076cd9ffda49cb6cfa87cb57fc8bb9fdc2bec1b7eb4e934bb2ab2/cbor2-6.1.5-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:fe81e4ff1b6bab72856d020dab89d86d4dcfbe18af4ff3fe2f391e1b03d0793c", size = 460595, upload-time = "2026-10-01T18:07:50.116Z" }, + { url = "https://files.pythonhosted.org/packages/a0/b5/40618405d7925149c59e4e2874c7247670ccb562ede141b4c3b46f826d02/cbor2-6.1.5-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:1ebbc6e2d5ea8acf44cc2247d48ca4ccae724fcdb97eaa673903e2d87f0ffc5d", size = 468860, upload-time = "2026-10-01T18:07:51.915Z" }, + { url = "https://files.pythonhosted.org/packages/3f/3d/e9dfa478e4964e741cf6a9c5a098644264d4e0f5bef18a51d3ec4e2610d3/cbor2-6.1.5-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:4db32eefe9fc173939d114fb78e09f967e69627714ad2e3bca807d0ea9d386ad", size = 527958, upload-time = "2026-10-01T18:07:53.916Z" }, + { url = "https://files.pythonhosted.org/packages/e0/39/13fa54e47a466414ea4a7b9d384b188539e869f6c2b57771b7e2f7429413/cbor2-6.1.5-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:0fa113902a302c22429b32e2454251a8fd14b18204fdff647c869a54114c3ed1", size = 536415, upload-time = "2026-10-01T18:07:55.668Z" }, + { url = "https://files.pythonhosted.org/packages/51/b7/f12c7b555ab56633c0285e10294d5ea8a1d6c3aba3699ca9a44b0d2d267b/cbor2-6.1.5-cp310-cp310-win32.whl", hash = "sha256:c87272763122be24213c7bb3d47750a3af034da8755fbd3fcb0694c1efb6c3e8", size = 285715, upload-time = "2026-10-01T18:07:57.406Z" }, + { url = "https://files.pythonhosted.org/packages/72/2a/fcf9348216a376bd3607fdd15f46aec50e665deff677b936fccc77d931b7/cbor2-6.1.5-cp310-cp310-win_amd64.whl", hash = "sha256:994b09c578e9dd7c5687a9f151f545bde705d12e47427b5a78c9d6cc970187f5", size = 308357, upload-time = "2026-10-01T18:07:58.892Z" }, + { url = "https://files.pythonhosted.org/packages/ed/15/4f3f573eb75cd7f2b709983bf567021d3d1018f101b6fb62f2e3d4d917c0/cbor2-6.1.5-cp310-cp310-win_arm64.whl", hash = "sha256:eba54489d82683e8cdb9af80a2e55c2089e439e76b60cdb9fd4dfdc62ecfee3c", size = 300702, upload-time = "2026-10-01T18:08:00.439Z" }, + { url = "https://files.pythonhosted.org/packages/84/62/6bd7ab55dda27ce4c0eefdf31a05b647c74a46e794bbf8ad5c3c26928e5b/cbor2-6.1.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:5a5859d1f82dce094a1bdd6a5b318411b750262070bf5d37fbc9607d185f0b1b", size = 416295, upload-time = "2026-10-01T18:08:01.813Z" }, + { url = "https://files.pythonhosted.org/packages/b2/22/9151b86062cc63d7155c86968971013dd6b01aeabd252a6dea015b16cfd9/cbor2-6.1.5-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:7de5383eb059498291415f5b07f99e54dac4603dc99960eb0e2307c9cb2dc352", size = 458485, upload-time = "2026-10-01T18:08:03.502Z" }, + { url = "https://files.pythonhosted.org/packages/44/d3/9aecf0948c50e54302ae8859c85358a82310331ca00e210f8984760a2e3c/cbor2-6.1.5-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:dd3e4f08aaf25bca5db6274ac40e4d138b0e09890510c1fda20d5b7840e505fa", size = 467049, upload-time = "2026-10-01T18:08:05.254Z" }, + { url = "https://files.pythonhosted.org/packages/b0/13/bf133682c99f162662395dafe3b2525ed0bdafa558e52ac840e7a134d5bc/cbor2-6.1.5-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:bb58549a45e3f6355338345a2df449f42f45d55e4a20af24d4302d76a1578650", size = 526819, upload-time = "2026-10-01T18:08:06.758Z" }, + { url = "https://files.pythonhosted.org/packages/a6/9b/7dda5b13258f740d529c9b3f5ed418d2c1aa4dbcbf886a35fb2f3f41970b/cbor2-6.1.5-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:a4956f498cbf5eab192e0f838cc787e09bef4caab57f05ccbf00451935cacb8b", size = 535239, upload-time = "2026-10-01T18:08:08.829Z" }, + { url = "https://files.pythonhosted.org/packages/0e/43/b72cb7b71c25b506a181ea9ec5bf634783c38e284873847ae6cb610c0f45/cbor2-6.1.5-cp311-cp311-win32.whl", hash = "sha256:f02c339ab9942578b63a5d54c8956191f6e88f3d8b2c918024ff565f7faa1bde", size = 285268, upload-time = "2026-10-01T18:08:10.591Z" }, + { url = "https://files.pythonhosted.org/packages/73/e5/9e51e3e43d6d42e71e93781d50b2f28cdcacc7f647681e07cbdaaf670e03/cbor2-6.1.5-cp311-cp311-win_amd64.whl", hash = "sha256:015ed73f10e1f7b67306d41e36e0d7dc40e4a2100bc5c29b7a7f039ad3dc9061", size = 307786, upload-time = "2026-10-01T18:08:12.034Z" }, + { url = "https://files.pythonhosted.org/packages/b7/7c/8514bf3a7a8af8347b8ba33cb9b3a9943200b37d81103b783543ab831ecb/cbor2-6.1.5-cp311-cp311-win_arm64.whl", hash = "sha256:f0bd6334302a5016a2b0f5530b7aea3ff588b6894523fd8491b49f7ce9e67f11", size = 300234, upload-time = "2026-10-01T18:08:13.579Z" }, + { url = "https://files.pythonhosted.org/packages/a0/d6/8278f1abd5b6b5bcfc94158226a737b62fa0e50ba1d8d0b77f42edbf74f8/cbor2-6.1.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0c1565bcd74a389b581e292592ccab0ed9c46286c6e986256820bc68c9ad7e8c", size = 407737, upload-time = "2026-10-01T18:08:14.982Z" }, + { url = "https://files.pythonhosted.org/packages/fa/1b/a58d72ecbe15273e4e4842ac2149361e2bc0ad75fcab117c06da3c31782f/cbor2-6.1.5-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:f8f85a49db66df77546d278de4d249772a4557d715df07ba8ae155cfa6a7fb31", size = 451924, upload-time = "2026-10-01T18:08:16.618Z" }, + { url = "https://files.pythonhosted.org/packages/72/28/72c76aee7aa74e5dc53b79505dc6c168805d20c8e75166143076c5b61906/cbor2-6.1.5-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b70d7c47ea84d456034d2be02e89d92eef7044cfcedf6f05058e21d4452f0fef", size = 463316, upload-time = "2026-10-01T18:08:18.293Z" }, + { url = "https://files.pythonhosted.org/packages/0b/a4/d81e9351c9ad37da4d999edcd05c6542a24e8899bb0ee8f91990e9e52981/cbor2-6.1.5-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:694f75fdcdb8c6b9a71ab77f789f56be1deab20bbdbf948d5ff53cd7c2543dfc", size = 519564, upload-time = "2026-10-01T18:08:20.123Z" }, + { url = "https://files.pythonhosted.org/packages/af/c7/f7da3d0d46022a1c802074e13966863972d68f29cf07301cce2c8e98febc/cbor2-6.1.5-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:09eeb76177758a0fdf1627a9428b384756872b048c6c0d7d158106b29b207d2c", size = 530974, upload-time = "2026-10-01T18:08:21.83Z" }, + { url = "https://files.pythonhosted.org/packages/5f/e3/74fddce015b171ee087a6e0185a233f3d29c7fda80cfa3041c796a67d100/cbor2-6.1.5-cp312-cp312-win32.whl", hash = "sha256:789ef813f416d353aecd5c8824860ee4be94e0f1179a385eb2beccfbeb615e4f", size = 281010, upload-time = "2026-10-01T18:08:23.614Z" }, + { url = "https://files.pythonhosted.org/packages/5e/f5/ecc8d6a9ff9322405b23a4d3226504e7d7a44424e0d831a02b49bac8e605/cbor2-6.1.5-cp312-cp312-win_amd64.whl", hash = "sha256:9677ce1c3c0cb1fa5a4f721a127fc2cc06e8efc43ee8e5f94e292186d6b51953", size = 304308, upload-time = "2026-10-01T18:08:25.077Z" }, + { url = "https://files.pythonhosted.org/packages/a8/90/23b702147b0858dbbc8a3136f288248118bb32f2785cc35c470a3b3f5571/cbor2-6.1.5-cp312-cp312-win_arm64.whl", hash = "sha256:b73d982e35a60e602a200feb2a9d272e850efdc9ff767b0f4887bdbc16d23e52", size = 293958, upload-time = "2026-10-01T18:08:26.493Z" }, +] + [[package]] name = "certifi" version = "2026.1.4" @@ -1843,6 +1875,7 @@ dependencies = [ { name = "annotation-protocol" }, { name = "attrs" }, { name = "bleak" }, + { name = "cbor2" }, { name = "cmeel-tinyxml2" }, { name = "cryptography" }, { name = "dimos-lcm" }, @@ -2539,6 +2572,7 @@ requires-dist = [ { name = "bosdyn-client", marker = "extra == 'spot'", specifier = ">=4.0.0" }, { name = "bosdyn-core", marker = "extra == 'spot'", specifier = ">=4.0.0" }, { name = "can-motor-control", marker = "extra == 'control'", specifier = ">=0.0.8" }, + { name = "cbor2", specifier = ">=6.1.5,<7" }, { name = "chromadb", marker = "extra == 'perception'", specifier = ">=1.0.0" }, { name = "cmeel-tinyxml2", specifier = ">=11,<12" }, { name = "coacd", marker = "extra == 'scene'", specifier = ">=1.0.0" }, From b4e64f5520dbfc1c793a831414c1cb4223c50e76 Mon Sep 17 00:00:00 2001 From: Paul Nechifor Date: Sun, 25 Oct 2026 07:22:32 +0200 Subject: [PATCH 2/3] fix(rpc): close the CBOR codec and send stream descriptors as data Review fixes for the pickle-to-CBOR commit. Decode no longer calls anything the payload names. A registered codec travels as its type plus the registering module; the receiver imports that module and takes the decoder from its own registry, so a payload naming builtins:eval is refused. StreamDescriptor carries a TransportSpec instead of a live transport, so list_streams and dimos shell no longer open an LCM socket per stream on every call. TransportSpec moves to dimos/core/stream.py and describes a live transport with TransportSpec.of(). Every transport's cbor_encode() returns constructor kwargs and PubSubTransport.cbor_decode is cls(**data); ZenohTransport and pZenohTransport take a qos kwarg, DDSTransport accepts the qos dict it encodes, and ROSTransport drops its unused kwargs. LCM RPC wraps bodies as bytes in a primitives-only envelope with a wire version: an undecodable request gets an error reply, an undecodable reply raises at the caller instead of timing out, and traffic from another dimos version is logged and ignored. Codec: encode() only raises TypeError and hook failures keep their path, enums travel by value so Flag combinations survive, typing.List, Optional, Union, Literal and Any are accepted, union and generic-alias payloads must hold types, a missing dataclass field names the class and field, and the fields-only dataclass rule and the pydantic private-state rule are documented and tested. The worker decodes replies outside its lock. Obstacle.points is an (N, 3) ndarray with an array-aware __eq__, so the voxel map crosses RPC as one tag instead of one tuple per cell. A cbor2 native-encoder hybrid was measured and rejected; _prepare says why. Also: isolated_python sends its config over stdin as CBOR, stale pickle comments and deepcopy-hook docstrings are updated, and the LCM and Zenoh service, Zenoh pubsub and DDS transport round trips are tested again. --- AGENTS.md | 2 +- .../coordination/blueprint_config/parser.py | 2 +- .../coordination/blueprint_config/schema.py | 3 +- .../blueprint_config/test_parsed.py | 3 +- .../blueprint_config/test_parser.py | 4 +- dimos/core/coordination/blueprints.py | 27 +-- dimos/core/coordination/module_coordinator.py | 4 +- dimos/core/coordination/python_worker.py | 13 +- .../coordination/test_module_coordinator.py | 3 + dimos/core/introspection/module/info.py | 11 +- dimos/core/module.py | 4 +- dimos/core/stream.py | 44 ++++- dimos/core/test_stream.py | 3 +- dimos/core/test_transport.py | 50 +++++ dimos/core/transport.py | 105 ++++------- .../experimental/isolated_python/bootstrap.py | 4 +- dimos/experimental/isolated_python/module.py | 5 +- dimos/experimental/robot/bosdyn/spot/rerun.py | 4 +- dimos/hosted/test_fragment_compiler.py | 4 +- dimos/manipulation/manipulation_module.py | 2 +- .../manipulation/planning/gripper_contact.py | 2 +- dimos/manipulation/planning/spec/models.py | 27 ++- .../manipulation/planning/spec/validation.py | 6 +- dimos/manipulation/test_roboplan.py | 18 +- dimos/memory/replay_module.py | 4 +- dimos/porcelain/remote_module_source.py | 11 +- dimos/porcelain/test_remote_module_source.py | 10 +- dimos/protocol/encode/cbor.py | 174 ++++++++++++++---- dimos/protocol/encode/test_cbor.py | 130 ++++++++++++- dimos/protocol/pubsub/encoders.py | 10 +- .../pubsub/impl/webrtc/providers/spec.py | 3 +- dimos/protocol/rpc/pubsubrpc.py | 120 +++++++----- dimos/protocol/rpc/test_spec.py | 72 +++++++- dimos/protocol/service/lcmservice.py | 3 +- dimos/protocol/service/test_lcmservice.py | 19 ++ dimos/protocol/service/test_zenohservice.py | 34 ++++ dimos/protocol/service/zenohservice.py | 3 +- dimos/robot/assets/model.py | 6 +- dimos/robot/assets/source.py | 8 +- .../blueprints/basic/r1pro_coordinator.py | 3 +- .../basic/unitree_go2_holonomic_benchmark.py | 4 +- .../basic/unitree_go2_rpp_benchmark.py | 4 +- dimos/utils/data.py | 4 +- dimos/web/cockpit.py | 5 +- 44 files changed, 711 insertions(+), 266 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index dc31ad28dc4..5c65c7743b3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -321,7 +321,7 @@ class MySkillContainer(Module): If multiple modules match the spec, use `.remappings()` to resolve. Source: `dimos/spec/utils.py`, `dimos/core/coordination/blueprints.py`. -RPC arguments, results and worker-pipe messages are CBOR (`dimos/protocol/encode/cbor.py`), not pickle. JSON-like values, tuples, sets, numpy arrays, enums, dataclasses, pydantic models, `Path`, message types with `lcm_encode`, classes and module-level functions cross as themselves. Anything else raises `TypeError` at the call site; give such a class `cbor_encode()` and a `cbor_decode()` classmethod. +RPC arguments, results and worker-pipe messages are CBOR (`dimos/protocol/encode/cbor.py`), not pickle. JSON-like values, tuples, sets, numpy arrays, enums, dataclasses (their fields only: state that `__post_init__` derives stays behind), pydantic models (private attributes stay behind), `Path`, message types with `lcm_encode`, classes and module-level functions cross as themselves. Anything else raises `TypeError` at the call site; give such a class `cbor_encode()` and a `cbor_decode()` classmethod. ### Adding a New Skill diff --git a/dimos/core/coordination/blueprint_config/parser.py b/dimos/core/coordination/blueprint_config/parser.py index a96ed466e38..573d0edbf23 100644 --- a/dimos/core/coordination/blueprint_config/parser.py +++ b/dimos/core/coordination/blueprint_config/parser.py @@ -80,11 +80,11 @@ ) from dimos.core.coordination.blueprints import ( Blueprint, - TransportSpec, config_key, transport_config_name, ) from dimos.core.global_config import GlobalConfig +from dimos.core.stream import TransportSpec def split_run_arguments(tokens: Sequence[str]) -> tuple[tuple[str, ...], tuple[str, ...]]: diff --git a/dimos/core/coordination/blueprint_config/schema.py b/dimos/core/coordination/blueprint_config/schema.py index 40720db8efe..0f7c684ba60 100644 --- a/dimos/core/coordination/blueprint_config/schema.py +++ b/dimos/core/coordination/blueprint_config/schema.py @@ -28,7 +28,8 @@ all_annotation_types, scalar_annotation_types, ) -from dimos.core.coordination.blueprints import BlueprintAtom, TransportSpec +from dimos.core.coordination.blueprints import BlueprintAtom +from dimos.core.stream import TransportSpec @dataclass(frozen=True, slots=True) diff --git a/dimos/core/coordination/blueprint_config/test_parsed.py b/dimos/core/coordination/blueprint_config/test_parsed.py index 8593779d077..fe1bdd070b3 100644 --- a/dimos/core/coordination/blueprint_config/test_parsed.py +++ b/dimos/core/coordination/blueprint_config/test_parsed.py @@ -21,9 +21,8 @@ from dimos.core.coordination.blueprint_config.errors import BlueprintConfigError from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser -from dimos.core.coordination.blueprints import TransportSpec from dimos.core.module import Module, ModuleConfig -from dimos.core.stream import Stream, Transport +from dimos.core.stream import Stream, Transport, TransportSpec class NestedConfig(BaseModel): diff --git a/dimos/core/coordination/blueprint_config/test_parser.py b/dimos/core/coordination/blueprint_config/test_parser.py index 3c207cca239..716ead29a82 100644 --- a/dimos/core/coordination/blueprint_config/test_parser.py +++ b/dimos/core/coordination/blueprint_config/test_parser.py @@ -25,9 +25,9 @@ BlueprintConfigParser, split_run_arguments, ) -from dimos.core.coordination.blueprints import TransportSpec, autoconnect +from dimos.core.coordination.blueprints import autoconnect from dimos.core.module import Module, ModuleConfig -from dimos.core.stream import Stream, Transport +from dimos.core.stream import Stream, Transport, TransportSpec class NestedConfig(BaseModel): diff --git a/dimos/core/coordination/blueprints.py b/dimos/core/coordination/blueprints.py index 58de4d95a10..92640cfa559 100644 --- a/dimos/core/coordination/blueprints.py +++ b/dimos/core/coordination/blueprints.py @@ -31,13 +31,11 @@ get_type_hints, ) -from pydantic import BaseModel - if TYPE_CHECKING: from dimos.protocol.service.system_configurator.base import SystemConfigurator from dimos.core.module import ModuleBase, is_module_type -from dimos.core.stream import IO, In, Out, Transport +from dimos.core.stream import IO, In, Out, Transport, TransportSpec from dimos.core.transport import PubSubTransport from dimos.spec.utils import Spec, is_spec from dimos.utils.logging_config import setup_logger @@ -190,29 +188,6 @@ def create(cls, module: type[ModuleBase], kwargs: dict[str, Any]) -> Self: ) -@dataclass(frozen=True) -class TransportSpec: - """Deferred transport construction: a transport class plus its ctor args. - - Blueprint authors declare transports via ``Cls.spec(...)`` so nothing is - constructed at blueprint-definition time. The coordinator materializes - specs at build time, once CLI/env/config overrides have resolved. - """ - - cls: type[Transport[Any]] - args: tuple[Any, ...] = () - kwargs: Mapping[str, Any] = field(default_factory=dict) - - @property - def config_cls(self) -> type[BaseModel] | None: - # Set by transports that expose a pydantic config override surface - return self.cls._config_cls - - def build(self, config: BaseModel | None = None) -> Transport[Any]: - extra = {"config": config} if config is not None else {} - return self.cls(*self.args, **self.kwargs, **extra) - - # MappingProxyType fields: sent as plain dicts, rebuilt on decode. _PROXY_FIELDS = ("transport_map", "global_config_overrides", "remapping_map") diff --git a/dimos/core/coordination/module_coordinator.py b/dimos/core/coordination/module_coordinator.py index 13ba57037e7..cd0390a1bba 100644 --- a/dimos/core/coordination/module_coordinator.py +++ b/dimos/core/coordination/module_coordinator.py @@ -28,7 +28,7 @@ from typing import TYPE_CHECKING, Any, NamedTuple, cast from dimos.core.coordination.blueprint_config.values import deep_merge, plain -from dimos.core.coordination.blueprints import TransportSpec, transport_config_name +from dimos.core.coordination.blueprints import transport_config_name from dimos.core.coordination.coordinator_rpc import CoordinatorRPC from dimos.core.coordination.worker_manager import WorkerManager from dimos.core.coordination.worker_manager_python import WorkerManagerPython @@ -36,7 +36,7 @@ from dimos.core.introspection.module.info import StreamDescriptor from dimos.core.module import ModuleBase, ModuleSpec, is_module_type from dimos.core.resource import Resource -from dimos.core.stream import Transport +from dimos.core.stream import Transport, TransportSpec from dimos.core.transport import ( LCMTransport, PubSubTransport, diff --git a/dimos/core/coordination/python_worker.py b/dimos/core/coordination/python_worker.py index 016a969f98f..1bdd0bf15ee 100644 --- a/dimos/core/coordination/python_worker.py +++ b/dimos/core/coordination/python_worker.py @@ -134,16 +134,18 @@ def cbor_decode(cls, data: dict[str, int]) -> Actor: def _send_request_to_worker(self, request: WorkerRequest) -> Any: if self._conn is None: raise RuntimeError("Actor connection not available - cannot send requests") - # Encode before taking the lock: an unencodable argument raises to the - # caller without holding up the other modules on this worker. + # Encode and decode outside the lock: an unencodable argument raises to the + # caller, and a large result is decoded, without holding up the other modules + # on this worker. payload = encode(request) if self._lock is not None: with self._lock: self._conn.send_bytes(payload) - response: WorkerResponse = _recv(self._conn) + reply = self._conn.recv_bytes() else: self._conn.send_bytes(payload) - response = _recv(self._conn) + reply = self._conn.recv_bytes() + response: WorkerResponse = decode(reply) if response.error: if "AttributeError" in response.error: # TODO: better error handling raise AttributeError(response.error) @@ -263,7 +265,8 @@ def deploy_module( payload = encode(request) with self._lock: self._conn.send_bytes(payload) - response: WorkerResponse = _recv(self._conn) + reply = self._conn.recv_bytes() + response: WorkerResponse = decode(reply) if response.error: raise RuntimeError(f"Failed to deploy module: {response.error}") diff --git a/dimos/core/coordination/test_module_coordinator.py b/dimos/core/coordination/test_module_coordinator.py index 8ab3edd3ccc..ebbabf707d9 100644 --- a/dimos/core/coordination/test_module_coordinator.py +++ b/dimos/core/coordination/test_module_coordinator.py @@ -131,6 +131,9 @@ def __init__( super().__init__(topic) self._config = config or self._config_cls(**config_kwargs) + def cbor_encode(self) -> dict[str, Any]: + return {"topic": self.topic, "config": self._config} + def start(self) -> None: pass diff --git a/dimos/core/introspection/module/info.py b/dimos/core/introspection/module/info.py index 1254038320a..97567d7a1cf 100644 --- a/dimos/core/introspection/module/info.py +++ b/dimos/core/introspection/module/info.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any, NamedTuple if TYPE_CHECKING: - from dimos.core.stream import StreamDirection, Transport + from dimos.core.stream import StreamDirection, TransportSpec # Internal RPCs to hide from io() output # Every module has these; io() shows what a module adds. @@ -42,7 +42,12 @@ class StreamDescriptor(NamedTuple): - """Wire metadata plus the message type and transport of one stream.""" + """Wire metadata plus the message type and transport of one stream. + + ``transport`` is a description, not a live transport: descriptors cross process + boundaries for inspection, and decoding one must not open sockets. Whoever needs + the transport calls ``transport.build()``. + """ module: str name: str @@ -50,7 +55,7 @@ class StreamDescriptor(NamedTuple): type_name: str channel: str type: type - transport: Transport[Any] + transport: TransportSpec @dataclass diff --git a/dimos/core/module.py b/dimos/core/module.py index 6b4b29d4bfe..629f2786a0c 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -42,7 +42,7 @@ from dimos.core.introspection.module.info import StreamDescriptor, extract_module_info from dimos.core.introspection.module.render import render_module_io from dimos.core.resource import CompositeResource -from dimos.core.stream import IO, In, Out, RemoteOut, Transport +from dimos.core.stream import IO, In, Out, RemoteOut, Transport, TransportSpec from dimos.core.transport import PubSubTransport from dimos.core.transport_factory import rpc_backend from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, DEFAULT_RPC_TIMEOUTS, RPCSpec @@ -362,7 +362,7 @@ def get_stream_descriptors(self) -> list[StreamDescriptor]: if isinstance(transport, PubSubTransport) else type(transport).__name__, type=stream.type, - transport=transport, + transport=TransportSpec.of(transport), ) ) return descriptors diff --git a/dimos/core/stream.py b/dimos/core/stream.py index cfe629c9a1e..8d883cf67aa 100644 --- a/dimos/core/stream.py +++ b/dimos/core/stream.py @@ -14,6 +14,7 @@ from __future__ import annotations +from dataclasses import dataclass, field import enum from typing import ( TYPE_CHECKING, @@ -37,7 +38,7 @@ from dimos.utils.reactive import backpressure if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Mapping from reactivex.observable import Observable @@ -108,6 +109,47 @@ def subscribe( def publish(self, msg: T) -> None: self.broadcast(None, msg) + @classmethod + def spec(cls, *args: Any, **kwargs: Any) -> TransportSpec: + """Defer construction: capture ctor args for the coordinator to build later.""" + return TransportSpec(cls, args, kwargs) + + +@dataclass(frozen=True) +class TransportSpec: + """Deferred transport construction: a transport class plus its ctor args. + + Blueprint authors declare transports via ``Cls.spec(...)`` so nothing is + constructed at blueprint-definition time. The coordinator materializes + specs at build time, once CLI/env/config overrides have resolved. A live + transport is described the same way (``TransportSpec.of``) when stream + metadata crosses a process boundary, so listing streams opens no sockets. + """ + + cls: type[Transport[Any]] + args: tuple[Any, ...] = () + kwargs: Mapping[str, Any] = field(default_factory=dict) + + @classmethod + def of(cls, transport: Transport[Any]) -> TransportSpec: + """The spec that rebuilds ``transport``: its class plus the kwargs ``cbor_encode()`` reports.""" + encode = getattr(type(transport), "cbor_encode", None) + if encode is None: + raise TypeError( + f"{type(transport).__name__} has no cbor_encode(), so it cannot be " + "described to another process" + ) + return cls(type(transport), kwargs=encode(transport)) + + @property + def config_cls(self) -> type[BaseModel] | None: + # Set by transports that expose a pydantic config override surface + return self.cls._config_cls + + def build(self, config: BaseModel | None = None) -> Transport[Any]: + extra = {"config": config} if config is not None else {} + return self.cls(*self.args, **self.kwargs, **extra) + class Stream(Generic[T]): _transport: Transport[T] | None diff --git a/dimos/core/test_stream.py b/dimos/core/test_stream.py index bc49f0ec35a..cc67f270658 100644 --- a/dimos/core/test_stream.py +++ b/dimos/core/test_stream.py @@ -306,5 +306,6 @@ def test_module_stream_crosses_processes_as_a_remote_handle() -> None: assert (remote.owner._worker_id, remote.owner._module_id) == (1, 2) assert remote.transport is None - with pytest.raises(ValueError, match="owner ref"): + # encode() turns whatever a hook raises into a TypeError that names the path. + with pytest.raises(TypeError, match=r"In at value: .*owner ref"): encode(In(Odometry, "odom")) diff --git a/dimos/core/test_transport.py b/dimos/core/test_transport.py index 43fb0b5ed21..b87071face0 100644 --- a/dimos/core/test_transport.py +++ b/dimos/core/test_transport.py @@ -13,9 +13,14 @@ # limitations under the License. +from dataclasses import dataclass + import pytest +from dimos.core.stream import TransportSpec from dimos.core.transport import ( + DDS_AVAILABLE, + DDSTransport, JpegLcmTransport, LCMTransport, pLCMTransport, @@ -23,6 +28,15 @@ from dimos.msgs.sensor_msgs.Image import Image from dimos.protocol.encode.cbor import decode, encode +if DDS_AVAILABLE: + # cyclonedds is the optional dds extra; see DDS_AVAILABLE. + from cyclonedds.idl import IdlStruct + from cyclonedds.idl.types import uint8 + + @dataclass + class DDSSample(IdlStruct): # type: ignore[misc] + value: uint8 # type: ignore[valid-type] + @pytest.mark.parametrize("transport_class", [pLCMTransport, LCMTransport, JpegLcmTransport]) def test_lcm_transport_roundtrip_preserves_connection_settings(transport_class): @@ -53,3 +67,39 @@ def test_with_topic_keeps_every_other_setting() -> None: moved.stop() finally: transport.stop() + + +def test_transport_spec_of_a_live_transport_decodes_without_building_it(mocker) -> None: + transport = LCMTransport("/described", Image, url="udpm://239.255.76.68:7798?ttl=2", ttl=2) + try: + spec = TransportSpec.of(transport) + assert spec.cls is LCMTransport + assert spec.kwargs["topic"] == "/described" + + lcm_ctor = mocker.patch("dimos.protocol.service.lcmservice.lcm_mod.LCM") + restored = decode(encode([spec] * 30)) # stream descriptors carry these + lcm_ctor.assert_not_called() # no sockets: descriptions are data + mocker.stop(lcm_ctor) + + built = restored[0].build() + try: + assert type(built) is LCMTransport + assert built.topic == transport.topic + assert built.lcm.config == transport.lcm.config + finally: + built.stop() + finally: + transport.stop() + + +@pytest.mark.skipif(not DDS_AVAILABLE, reason="cyclonedds is not installed") +def test_dds_transport_roundtrip_preserves_qos_and_domain() -> None: + from cyclonedds.qos import Policy, Qos # cyclonedds: see DDS_AVAILABLE + + qos = Qos(Policy.Reliability.BestEffort, Policy.History.KeepLast(depth=1)) + transport = DDSTransport("/dds_roundtrip", DDSSample, qos=qos, domain_id=3) + restored = decode(encode(transport)) + assert type(restored) is DDSTransport + assert restored.topic == transport.topic + assert restored.dds.qos.asdict() == qos.asdict() + assert restored.dds.config.domain_id == 3 diff --git a/dimos/core/transport.py b/dimos/core/transport.py index 3e3a242d507..33edbcd8f10 100644 --- a/dimos/core/transport.py +++ b/dimos/core/transport.py @@ -14,6 +14,7 @@ from __future__ import annotations +from abc import abstractmethod import copy import importlib.util import threading @@ -37,6 +38,7 @@ PickleZenoh, Topic as ZenohTopic, Zenoh, + ZenohQoS, ) from dimos.stream.audio.base import AudioEvent from dimos.utils import colors @@ -47,7 +49,6 @@ if TYPE_CHECKING: from collections.abc import Callable - from dimos.core.coordination.blueprints import TransportSpec from dimos.protocol.pubsub.impl.rospubsub import DimosROS T = TypeVar("T") @@ -91,22 +92,16 @@ def channel(self) -> str: """The channel string this transport publishes and subscribes on.""" return str(self.topic) - @classmethod - def spec(cls, *args: Any, **kwargs: Any) -> TransportSpec: - """Defer construction: capture ctor args for the coordinator to build later.""" - from dimos.core.coordination.blueprints import TransportSpec - - return TransportSpec(cls, args, kwargs) - - # Transports cross process boundaries as their constructor settings (never the - # live pubsub): a dict with the plain "topic" string plus whatever the subclass - # needs to rebuild itself. - def cbor_encode(self) -> dict[str, Any]: - raise NotImplementedError + # Transports cross process boundaries as their constructor kwargs, never as the + # live pubsub: cbor_encode() returns a dict with the plain "topic" string plus + # whatever else cls(**kwargs) needs, and StreamDescriptor carries the same dict as + # a TransportSpec so that listing streams builds nothing. + @abstractmethod + def cbor_encode(self) -> dict[str, Any]: ... @classmethod def cbor_decode(cls, data: dict[str, Any]) -> PubSubTransport[Any]: - raise NotImplementedError + return cls(**data) def __deepcopy__(self, memo: dict[int, Any]) -> PubSubTransport[T]: # A fresh instance from the same settings, never a copy of the live pubsub: @@ -132,11 +127,6 @@ def __init__(self, topic: str, **kwargs) -> None: # type: ignore[no-untyped-def def cbor_encode(self) -> dict[str, Any]: return {"topic": self.topic, **self.lcm.config.model_dump(exclude={"lcm"})} - @classmethod - def cbor_decode(cls, data: dict[str, Any]) -> pLCMTransport[Any]: - data = dict(data) - return cls(data.pop("topic"), **data) - def broadcast(self, _: Stream[T] | None, msg: T) -> None: if not self._started: self.start() @@ -182,11 +172,6 @@ def cbor_encode(self) -> dict[str, Any]: **self.lcm.config.model_dump(exclude={"lcm"}), } - @classmethod - def cbor_decode(cls, data: dict[str, Any]) -> LCMTransport[Any]: - data = dict(data) - return cls(data.pop("topic"), data.pop("type"), **data) - def broadcast(self, _, msg) -> None: # type: ignore[no-untyped-def] if not self._started: self.start() @@ -222,20 +207,15 @@ def stop(self) -> None: class ROSTransport(PubSubTransport[DimosMsg]): _ros: DimosROS | None = None - def __init__(self, topic: str, msg_type: type[DimosMsg], **kwargs: Any) -> None: + def __init__(self, topic: str, msg_type: type[DimosMsg]) -> None: # rclpy (about 0.1 s with rospubsub): only loaded when a ROSTransport exists. from dimos.protocol.pubsub.impl.rospubsub import ROSTopic super().__init__(ROSTopic(topic, msg_type)) - self._kwargs = kwargs def cbor_encode(self) -> dict[str, Any]: return {"topic": self.topic.topic, "msg_type": self.topic.msg_type} - @classmethod - def cbor_decode(cls, data: dict[str, Any]) -> ROSTransport: - return cls(data["topic"], data["msg_type"]) - def broadcast(self, _: Stream[DimosMsg] | None, msg: DimosMsg) -> None: if self._ros is None: self.start() @@ -254,7 +234,7 @@ def start(self) -> None: if self._ros is None: from dimos.protocol.pubsub.impl.rospubsub import DimosROS # rclpy: see __init__ - self._ros = DimosROS(**self._kwargs) + self._ros = DimosROS() self._ros.start() def stop(self) -> None: @@ -269,12 +249,16 @@ def stop(self) -> None: class DDSTransport(PubSubTransport[T]): - def __init__(self, topic: str, type: type, **kwargs) -> None: # type: ignore[no-untyped-def] + def __init__(self, topic: str, type: type, qos: Any = None, **kwargs) -> None: # type: ignore[no-untyped-def] # cyclonedds: see DDS_AVAILABLE. + from cyclonedds.qos import Qos + from dimos.protocol.pubsub.impl.ddspubsub import DDS, Topic as DDSTopic super().__init__(DDSTopic(topic, type)) - self.dds = DDS(**kwargs) + if isinstance(qos, dict): # the wire form cbor_encode() produces + qos = Qos.fromdict(qos) + self.dds = DDS(qos=qos, **kwargs) self._started: bool = False self._start_lock = threading.RLock() @@ -287,19 +271,6 @@ def cbor_encode(self) -> dict[str, Any]: **self.dds.config.model_dump(), } - @classmethod - def cbor_decode(cls, data: dict[str, Any]) -> DDSTransport[Any]: - from cyclonedds.qos import Qos # cyclonedds: see DDS_AVAILABLE. - - data = dict(data) - qos = data.pop("qos") - return cls( - data.pop("topic"), - data.pop("type"), - qos=Qos.fromdict(qos) if qos is not None else None, - **data, - ) - def start(self) -> None: with self._start_lock: if not self._started: @@ -373,10 +344,6 @@ def __init__( def cbor_encode(self) -> dict[str, Any]: return {"topic": self.topic, "msg_type": self._msg_type, "config": self._config} - @classmethod - def cbor_decode(cls, data: dict[str, Any]) -> WebRTCTransport[Any]: - return cls(data["topic"], data["msg_type"], config=data["config"]) - def broadcast(self, _: Stream[M] | None, msg: M) -> None: if not self._started: self.start() @@ -449,19 +416,13 @@ class _ProviderTransport(Transport[T]): def __init__(self, *, config: ProviderConfig | None = None, **config_kwargs: Any) -> None: self._config = config or self._config_cls(**config_kwargs) - def cbor_encode(self) -> ProviderConfig: - return self._config - - @classmethod - def cbor_decode(cls, config: ProviderConfig) -> _ProviderTransport[Any]: - return cls(config=config) + # Constructor kwargs, like PubSubTransport.cbor_encode. + def cbor_encode(self) -> dict[str, Any]: + return {"config": self._config} @classmethod - def spec(cls, *args: Any, **kwargs: Any) -> TransportSpec: - """Defer construction: capture ctor args for the coordinator to build later.""" - from dimos.core.coordination.blueprints import TransportSpec - - return TransportSpec(cls, args, kwargs) + def cbor_decode(cls, data: dict[str, Any]) -> _ProviderTransport[Any]: + return cls(**data) def stop(self) -> None: pass # shared provider is process-scoped (see WebRTCTransport.stop) @@ -563,9 +524,15 @@ class ZenohTransport(PubSubTransport[T]): _started: bool = False - def __init__(self, topic: str | ZenohTopic, type: type | None = None, **kwargs: Any) -> None: + def __init__( + self, + topic: str | ZenohTopic, + type: type | None = None, + qos: ZenohQoS | None = None, + **kwargs: Any, + ) -> None: if isinstance(topic, str): - topic = ZenohTopic(topic, type) + topic = ZenohTopic(topic, type, qos=qos) super().__init__(topic) self.zenoh = Zenoh(**kwargs) self._start_lock = threading.RLock() @@ -582,10 +549,6 @@ def publish_qos(self) -> dict[str, str] | None: def cbor_encode(self) -> dict[str, Any]: return {"topic": self.topic.topic, "type": self.topic.lcm_type, "qos": self.topic.qos} - @classmethod - def cbor_decode(cls, data: dict[str, Any]) -> ZenohTransport[Any]: - return cls(ZenohTopic(data["topic"], data["type"], qos=data["qos"])) - def start(self) -> None: with self._start_lock: if not self._started: @@ -620,8 +583,8 @@ class pZenohTransport(PubSubTransport[T]): _started: bool = False - def __init__(self, topic: str | ZenohTopic, **kwargs: Any) -> None: - self._zenoh_topic = ZenohTopic(topic) if isinstance(topic, str) else topic + def __init__(self, topic: str | ZenohTopic, qos: ZenohQoS | None = None, **kwargs: Any) -> None: + self._zenoh_topic = ZenohTopic(topic, qos=qos) if isinstance(topic, str) else topic super().__init__(self._zenoh_topic.pattern) self.zenoh = PickleZenoh(**kwargs) self._start_lock = threading.RLock() @@ -638,10 +601,6 @@ def publish_qos(self) -> dict[str, str] | None: def cbor_encode(self) -> dict[str, Any]: return {"topic": self._zenoh_topic.topic, "qos": self._zenoh_topic.qos} - @classmethod - def cbor_decode(cls, data: dict[str, Any]) -> pZenohTransport[Any]: - return cls(ZenohTopic(data["topic"], qos=data["qos"])) - def start(self) -> None: with self._start_lock: if not self._started: diff --git a/dimos/experimental/isolated_python/bootstrap.py b/dimos/experimental/isolated_python/bootstrap.py index 1a0d24b89a6..0ca072237c0 100644 --- a/dimos/experimental/isolated_python/bootstrap.py +++ b/dimos/experimental/isolated_python/bootstrap.py @@ -19,7 +19,6 @@ from importlib.metadata import EntryPoint import inspect import os -import pickle import signal import sys import threading @@ -28,6 +27,7 @@ import typer from dimos.experimental.isolated_python.module import IsolatedPythonModule, contract_rpc_names +from dimos.protocol.encode.cbor import decode from dimos.spec.utils import _signatures_compatible @@ -91,7 +91,7 @@ def main( if not issubclass(runtime_class, IsolatedPythonModule): raise TypeError(f"{implementation!r} is not an IsolatedPythonModule subclass") validate_runtime(declaration_class, runtime_class) - kwargs = pickle.load(sys.stdin.buffer) + kwargs = decode(sys.stdin.buffer.read()) kwargs["instance_name"] = instance_name module = runtime_class(_isolated_python_runtime=True, **kwargs) os.write(handshake_fd, b"READY\n") diff --git a/dimos/experimental/isolated_python/module.py b/dimos/experimental/isolated_python/module.py index 1587ba9cefe..a763e3e6c19 100644 --- a/dimos/experimental/isolated_python/module.py +++ b/dimos/experimental/isolated_python/module.py @@ -19,7 +19,6 @@ from hashlib import sha256 import os from pathlib import Path -import pickle import select import subprocess import threading @@ -31,6 +30,7 @@ from dimos.core.module import Module from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.rpc_client import RPCClient +from dimos.protocol.encode.cbor import encode from dimos.utils.data import get_project_root from dimos.utils.generic import short_id from dimos.utils.logging_config import setup_logger @@ -200,7 +200,8 @@ def _spawn_runtime(self) -> None: assert self._process.stdin is not None kwargs = self.config.model_dump() kwargs["instance_name"] = self._runtime_name - pickle.dump(kwargs, self._process.stdin) + # The same codec as the worker pipe; the runtime reads stdin to EOF. + self._process.stdin.write(encode(kwargs)) self._process.stdin.close() watchdog = threading.Thread( diff --git a/dimos/experimental/robot/bosdyn/spot/rerun.py b/dimos/experimental/robot/bosdyn/spot/rerun.py index fc128193cc4..3f56d9ec934 100644 --- a/dimos/experimental/robot/bosdyn/spot/rerun.py +++ b/dimos/experimental/robot/bosdyn/spot/rerun.py @@ -112,8 +112,8 @@ def _camera_info_pinhole(camera_info: CameraInfo, origin: Callable[[str], str]) ] -# Module-level (not closures) so the RerunBridgeModule config stays picklable -# when it is shipped to its worker process. +# Module-level (not closures) so the RerunBridgeModule config stays encodable +# when it is shipped to its worker process: functions cross by import path. def _grayscale_info_to_pinhole(camera_info: CameraInfo) -> RerunData | None: return _camera_info_pinhole(camera_info, _grayscale_origin) diff --git a/dimos/hosted/test_fragment_compiler.py b/dimos/hosted/test_fragment_compiler.py index 62ec4350c31..9c43d41ca03 100644 --- a/dimos/hosted/test_fragment_compiler.py +++ b/dimos/hosted/test_fragment_compiler.py @@ -15,9 +15,9 @@ import pytest from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser -from dimos.core.coordination.blueprints import TransportSpec, autoconnect +from dimos.core.coordination.blueprints import autoconnect from dimos.core.module import Module, ModuleConfig -from dimos.core.stream import In, Out +from dimos.core.stream import In, Out, TransportSpec from dimos.core.transport import ZenohTransport from dimos.hosted.fragment import RemoteModuleReference from dimos.hosted.fragment_compiler import compile_fragments, resolve_hosted_assignments diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 48b73ef6bc7..46cf4e62eee 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -1321,7 +1321,7 @@ def _apply_voxel_map(self, cloud: PointCloud2) -> None: name=VOXEL_MAP_OBSTACLE_ID, obstacle_type=ObstacleType.OCTREE, pose=PoseStamped(frame_id=frame), - points=tuple(map(tuple, points.tolist())), + points=points, octree_resolution=self.config.voxel_map_resolution, ) try: diff --git a/dimos/manipulation/planning/gripper_contact.py b/dimos/manipulation/planning/gripper_contact.py index 0e26a194244..eb9ee05453b 100644 --- a/dimos/manipulation/planning/gripper_contact.py +++ b/dimos/manipulation/planning/gripper_contact.py @@ -325,7 +325,7 @@ def _octree( name=name, obstacle_type=ObstacleType.OCTREE, pose=PoseStamped(frame_id=frame_id), - points=tuple((x, y, z) for x, y, z in finite.tolist()), + points=finite, octree_resolution=cell, touch_links=touch_links, ) diff --git a/dimos/manipulation/planning/spec/models.py b/dimos/manipulation/planning/spec/models.py index 1bd628f07d1..f144f3550d1 100644 --- a/dimos/manipulation/planning/spec/models.py +++ b/dimos/manipulation/planning/spec/models.py @@ -17,10 +17,12 @@ from __future__ import annotations from collections.abc import Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields from typing import TYPE_CHECKING, TypeAlias from uuid import uuid4 +import numpy as np + from dimos.manipulation.planning.spec.enums import ( IKStatus, ObstacleType, @@ -28,7 +30,6 @@ ) if TYPE_CHECKING: - import numpy as np from numpy.typing import NDArray from dimos.manipulation.planning.groups.models import PlanningGroup @@ -110,7 +111,8 @@ class Obstacle: - MESH: Not used color: RGBA color tuple (0-1 range) mesh_path: Path to mesh file (for MESH type) - points: Occupied cell centers in the obstacle's frame (for OCTREE type) + points: Occupied cell centers in the obstacle's frame as an (N, 3) float + array (for OCTREE type). Sequences of triples are accepted and converted. octree_resolution: Edge length of an OCTREE cell (meters) touch_links: Robot links allowed to touch the obstacle, by exact link name; their collisions with it are ignored @@ -122,12 +124,25 @@ class Obstacle: dimensions: tuple[float, ...] = () color: tuple[float, float, float, float] = DEFAULT_OBSTACLE_RGBA mesh_path: str | None = None - # Plain tuples rather than an array: obstacles are deepcopied, compared and - # sent across worker RPC, and an ndarray field breaks comparison. - points: tuple[tuple[float, float, float], ...] = () + # An array, not one tuple per cell: it crosses worker RPC as a single ndarray tag. + # __eq__ below keeps obstacles comparable, which the world's snapshots rely on. + points: NDArray[np.float64] = field(default_factory=lambda: np.empty((0, 3))) octree_resolution: float | None = None touch_links: tuple[str, ...] = () + def __post_init__(self) -> None: + self.points = np.asarray(self.points, dtype=np.float64) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Obstacle): + return NotImplemented + for f in fields(self): + mine, theirs = getattr(self, f.name), getattr(other, f.name) + same = np.array_equal(mine, theirs) if f.name == "points" else mine == theirs + if not same: + return False + return True + @dataclass class IKResult: diff --git a/dimos/manipulation/planning/spec/validation.py b/dimos/manipulation/planning/spec/validation.py index 766a144f71c..aa8c9549251 100644 --- a/dimos/manipulation/planning/spec/validation.py +++ b/dimos/manipulation/planning/spec/validation.py @@ -264,13 +264,13 @@ def validate_obstacle( def _validate_octree(obstacle: Obstacle) -> None: if obstacle.octree_resolution is None or not obstacle.octree_resolution > 0.0: raise ValueError("OCTREE obstacle requires a positive octree_resolution") - if not obstacle.points: + if len(obstacle.points) == 0: raise ValueError("OCTREE obstacle requires at least one point") if len(obstacle.points) > MAX_OCTREE_POINTS: raise ValueError( f"OCTREE obstacle has {len(obstacle.points)} points, over the " - f"{MAX_OCTREE_POINTS} limit. Each one crosses worker RPC as an encoded " - f"tuple, so an unbounded map stalls the pipe." + f"{MAX_OCTREE_POINTS} limit. The whole map crosses worker RPC on every " + f"update, so an unbounded one stalls the pipe." ) points = np.asarray(obstacle.points, dtype=np.float64) if points.ndim != 2 or points.shape[1] != 3: diff --git a/dimos/manipulation/test_roboplan.py b/dimos/manipulation/test_roboplan.py index b19dd48081c..e96a3d454f6 100644 --- a/dimos/manipulation/test_roboplan.py +++ b/dimos/manipulation/test_roboplan.py @@ -50,6 +50,7 @@ from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.protocol.encode.cbor import decode, encode from dimos.robot.assets.model import RobotModel from dimos.utils.transform_utils import pose_to_matrix @@ -775,9 +776,8 @@ def test_octree_obstacle_survives_the_deepcopy_and_equality_the_world_does( fake_roboplan: None, robot_config: RobotModelConfig, ) -> None: - # The world snapshots obstacles with deepcopy and hands copies back out. A - # raw ndarray field would make that snapshot uncomparable, so the points - # stay plain tuples. + # The world snapshots obstacles with deepcopy and hands copies back out. The + # points are an ndarray, so Obstacle.__eq__ has to keep the snapshot comparable. world = _make_world(fake_roboplan, robot_config) obstacle = Obstacle( name="voxel-map", @@ -790,8 +790,18 @@ def test_octree_obstacle_survives_the_deepcopy_and_equality_the_world_does( stored = world.get_obstacles()[0] assert stored == world.get_obstacles()[0] - assert stored.points == ((0.0, 0.0, 0.0),) + assert stored == obstacle + assert stored != Obstacle( + name="voxel-map", + obstacle_type=ObstacleType.OCTREE, + pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + points=((0.0, 0.0, 1.0),), + octree_resolution=0.05, + ) + np.testing.assert_array_equal(stored.points, [(0.0, 0.0, 0.0)]) assert stored.octree_resolution == 0.05 + # The reason for the array: the map crosses worker RPC as one ndarray tag. + assert decode(encode(stored)) == stored def test_octree_obstacle_rejects_geometry_it_cannot_build( diff --git a/dimos/memory/replay_module.py b/dimos/memory/replay_module.py index b3259dc29e4..57567fa7409 100644 --- a/dimos/memory/replay_module.py +++ b/dimos/memory/replay_module.py @@ -187,8 +187,8 @@ def _port_names(topics: str, types: dict[str, type]) -> list[str]: def replay_module(dataset: str, topics: str = "*", name: str = "Replay") -> type[ReplayModule]: """Build a :class:`ReplayModule` subclass with an ``Out`` per stream in *dataset*. - Assign the result to *name* at module level: deploying to a worker pickles the - class by that path. An empty *dataset* yields a port-less class. + Assign the result to *name* at module level: deploying to a worker sends the + class by that import path. An empty *dataset* yields a port-less class. """ ports: dict[str, Any] = {} if dataset: diff --git a/dimos/porcelain/remote_module_source.py b/dimos/porcelain/remote_module_source.py index 44fcf471be9..68dcc27ccbb 100644 --- a/dimos/porcelain/remote_module_source.py +++ b/dimos/porcelain/remote_module_source.py @@ -143,16 +143,17 @@ def _get_stream(self, remote_name: str, name: str) -> RemoteOut[Any]: ) if descriptor is None: raise AttributeError(f"{remote_name!r} has no stream named {name!r}") - stream_type, transport = descriptor.type, descriptor.transport - if isinstance( - transport, (WebRTCTransport, WebRTCVideoTransport, WebRTCAudioTransport) + spec = descriptor.transport + if issubclass( + spec.cls, (WebRTCTransport, WebRTCVideoTransport, WebRTCAudioTransport) ): raise NotImplementedError( "WebRTC stream subscriptions require an operator client; " "use the hosted teleop client." ) - # Expose a subscription view for inputs, outputs and IO ports. - self._streams[key] = RemoteOut(stream_type, name, remote_name, transport) + # Expose a subscription view for inputs, outputs and IO ports. Only the + # requested stream's transport is built; descriptors are data. + self._streams[key] = RemoteOut(descriptor.type, name, remote_name, spec.build()) return self._streams[key] def invalidate(self, name: str) -> None: diff --git a/dimos/porcelain/test_remote_module_source.py b/dimos/porcelain/test_remote_module_source.py index 12c99efdbfd..1cfd8c72e19 100644 --- a/dimos/porcelain/test_remote_module_source.py +++ b/dimos/porcelain/test_remote_module_source.py @@ -30,7 +30,7 @@ from dimos.core.demos.stress_test_module import StressTestModule from dimos.core.global_config import GlobalConfig from dimos.core.module import Module -from dimos.core.stream import IO, In, Out +from dimos.core.stream import IO, In, Out, TransportSpec from dimos.core.transport import LCMTransport, pLCMTransport from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.porcelain.dimos import Dimos @@ -369,6 +369,14 @@ def test_connect_discovers_only_wired_streams(running_app, client, stream_module assert all(s.channel.endswith(f"/{s.name}") for s in streams) assert client.list_streams(client.get_module("remote/streams")) == streams assert {s.module_name for s in client.list_streams()} >= {"remote/streams"} + # Descriptors are data on both sides: nothing is built until a stream is asked for. + for app in (running_app, client): + descriptors = [ + d for d in app._source.list_stream_descriptors() if d.module == "remote/streams" + ] + assert len(descriptors) == 3 + assert all(type(d.transport) is TransportSpec for d in descriptors) + assert all(d.transport.cls is pLCMTransport for d in descriptors) @pytest.mark.parametrize("name", ["messages", "commands", "state"]) diff --git a/dimos/protocol/encode/cbor.py b/dimos/protocol/encode/cbor.py index 22ebe0f353f..92cabaf3bb2 100644 --- a/dimos/protocol/encode/cbor.py +++ b/dimos/protocol/encode/cbor.py @@ -17,14 +17,30 @@ JSON-like values (None, bool, int, float, str, bytes, list, dict with primitive keys) map to native CBOR. Every other supported type travels inside a CBOR tag so it comes back as the same Python type in the other process: tuple, NamedTuple, set, frozenset, -numpy arrays and scalars, Enum, DimosMsg (as its class plus LCM bytes), dataclasses, pydantic models, -Path, classes (including ``list[float]``-style generics and ``X | Y`` unions), module-level -functions, functools.partial, and any class that implements -``cbor_encode``/``cbor_decode`` (see ``CborEncodable``). Anything else raises -``TypeError`` at encode time. There is no pickle-style fallback. +numpy arrays and scalars, Enum (by value, so Flag combinations survive), DimosMsg (as +its class plus LCM bytes), dataclasses, pydantic models, Path, classes, type forms +(``list[float]``, ``typing.List[float]``, ``X | Y``, ``Optional[X]``, ``Literal[...]``, +``Any``), module-level functions, functools.partial, and any class that implements +``cbor_encode``/``cbor_decode`` (see ``CborEncodable``) or is registered with +``register_codec``. Anything else raises ``TypeError`` at encode time, naming the +value's path in the payload. There is no pickle-style fallback. Classes and functions are sent by reference as ``"module:qualname"`` and must be importable under that name on the receiving side. + +Dataclasses travel as their ``fields()`` and are rebuilt without ``__init__`` or +``__post_init__``: the sender already normalized the values, and anything +``__post_init__`` or a cached property derives stays behind. A dataclass that needs +such state on the receiving side recomputes it lazily or implements the hooks. + +Pydantic models travel as their fields and extras and are rebuilt with +``model_construct``: validation does not run again, ``model_post_init`` does, and +private attributes start from their defaults, so private state is process-local. + +Decoding never calls a callable the payload chose. It imports modules by name and +calls only ``cbor_decode``, ``lcm_decode``, ``model_construct``, ``_make``, +``__new__`` and ``__class_getitem__`` of the classes it resolves, plus the decoders +this process registered with ``register_codec``. """ from __future__ import annotations @@ -38,6 +54,7 @@ import pathlib import sys import types +import typing from typing import Any, Protocol import cbor2 @@ -65,6 +82,7 @@ TAG_GENERIC_ALIAS = 4057 TAG_UNION = 4058 TAG_ELLIPSIS = 4059 +TAG_LITERAL = 4060 _PRIMITIVES = (type(None), bool, int, float, str, bytes) @@ -79,6 +97,7 @@ class CborEncodable(Protocol): ``cbor_encode`` returns any encodable value; ``cbor_decode`` rebuilds an instance from it. The returned object may be a stand-in of another class (a stream handle decodes to its remote twin), which is why ``cbor_decode`` is typed as ``Any``. + The codec looks both hooks up on the class, never on the instance. """ def cbor_encode(self) -> Any: ... @@ -87,22 +106,30 @@ def cbor_encode(self) -> Any: ... def cbor_decode(cls, data: Any) -> Any: ... -# type -> (encode, "module:qualname" of the decode function) -_REGISTRY: dict[type, tuple[Callable[[Any], Any], str]] = {} +# type -> (encode, decode, name of the module that registered them) +_REGISTRY: dict[type, tuple[Callable[[Any], Any], Callable[[Any], Any], str]] = {} def register_codec(type_: type, encode: Callable[[Any], Any], decode: Callable[[Any], Any]) -> None: """Make instances of ``type_`` (and subclasses) encodable without touching the class. - ``decode`` must be a module-level function: the payload names it by reference so the - decoding process imports the registering module on demand. + ``decode`` must be a module-level function, and this call must run when its module + is imported: a payload names the type and that module, the receiving process imports + the module and then uses the decoder from its own registry. A decoder is never taken + from the payload. """ - _REGISTRY[type_] = (encode, _ref(decode, repr(decode))) + _ref(decode, repr(decode)) + _REGISTRY[type_] = (encode, decode, decode.__module__) def encode(obj: Any) -> bytes: - """Encode ``obj`` to CBOR bytes. Raises ``TypeError`` for unsupported values.""" - return cbor2.dumps(_prepare(obj, "value", set())) + """Encode ``obj`` to CBOR bytes. Raises ``TypeError`` for any value that cannot cross.""" + try: + return cbor2.dumps(_prepare(obj, "value", set())) + except TypeError: + raise + except Exception as e: # cbor2 itself, e.g. a str holding a lone surrogate + raise TypeError(f"Cannot CBOR-encode value: {e}") from e def decode(data: bytes | bytearray | memoryview) -> Any: @@ -117,7 +144,14 @@ def decode(data: bytes | bytearray | memoryview) -> Any: def _prepare(obj: Any, path: str, active: set[int]) -> Any: - """Map an object graph to values cbor2 encodes natively (CBORTag for our tags).""" + """Map an object graph to values cbor2 encodes natively (CBORTag for our tags). + + This walk is the cost of the codec (about 10x pickle on many small objects). A + cbor2 ``encoders=`` hybrid was measured in 2026-10 and did not help: the C encoder + resolves types along the MRO, so ``str``/``int``/``float`` need exact-type Python + wrappers to keep ``Key(str)``, ``IntEnum`` and ``np.float64`` from being flattened, + and those wrappers cost as much as this walk. Send arrays, not tuples of scalars. + """ t = type(obj) if t in _PRIMITIVES: return obj @@ -144,25 +178,29 @@ def _prepare_object(obj: Any, t: type, path: str, active: set[int]) -> Any: # metaclass serialize classes that are not importable by name. custom_encode = getattr(t, "cbor_encode", None) if callable(custom_encode): - payload = _prepare(custom_encode(obj), f"{path}.cbor_encode()", active) - return CBORTag(TAG_CUSTOM, [_ref(t, path), payload]) - # Generic aliases and unions come before the class check: on Python 3.10 - # isinstance(list[int], type) is True. - if isinstance(obj, types.GenericAlias): - origin = _prepare(obj.__origin__, f"{path}.__origin__", active) - return CBORTag(TAG_GENERIC_ALIAS, [origin, _prepare_items(obj.__args__, path, active)]) - if isinstance(obj, types.UnionType): - return CBORTag(TAG_UNION, _prepare_items(obj.__args__, path, active)) + payload = _hook(functools.partial(custom_encode, obj), t, path) + return CBORTag( + TAG_CUSTOM, [_ref(t, path), _prepare(payload, f"{path}.cbor_encode()", active)] + ) + # Type forms come before the class check: on Python 3.10 isinstance(list[int], type) + # is True. typing's aliases and special forms all live in the typing module, but so + # does the metaclass of Protocol subclasses, which are ordinary classes. + if isinstance(obj, (types.GenericAlias, types.UnionType)) or ( + t.__module__ == "typing" and not isinstance(obj, type) + ): + return _prepare_type_form(obj, path, active) if isinstance(obj, type): return CBORTag(TAG_CLASS, _ref(obj, path)) if obj is Ellipsis: # the "..." in tuple[int, ...] return CBORTag(TAG_ELLIPSIS, None) - for registered_type, (registered_encode, decode_ref) in _REGISTRY.items(): + for registered_type, (registered_encode, _, module) in _REGISTRY.items(): if isinstance(obj, registered_type): - payload = _prepare(registered_encode(obj), f"{path}.encode()", active) - return CBORTag(TAG_REGISTERED, [decode_ref, payload]) + payload = _hook(functools.partial(registered_encode, obj), t, path) + prepared = _prepare(payload, f"{path}.encode()", active) + return CBORTag(TAG_REGISTERED, [_ref(registered_type, path), module, prepared]) if isinstance(obj, enum.Enum): - return CBORTag(TAG_ENUM, [_ref(t, path), obj.name]) + # The value, not the name: a Flag combination has no member name. + return CBORTag(TAG_ENUM, [_ref(t, path), _prepare(obj.value, f"{path}.value", active)]) if isinstance(obj, np.ndarray): if obj.dtype.hasobject or obj.dtype.fields is not None: raise TypeError(f"Cannot CBOR-encode ndarray of dtype {obj.dtype} at {path}") @@ -174,7 +212,7 @@ def _prepare_object(obj: Any, t: type, path: str, active: set[int]) -> Any: if isinstance(getattr(t, "msg_name", None), str) and callable(getattr(t, "lcm_encode", None)): # The class, not msg_name: a subclass that inherits its parent's msg_name # must still come back as the subclass. - return CBORTag(TAG_DIMOSMSG, [_ref(t, path), obj.lcm_encode()]) + return CBORTag(TAG_DIMOSMSG, [_ref(t, path), _hook(obj.lcm_encode, t, path)]) if isinstance(obj, tuple) and hasattr(t, "_fields"): return CBORTag(TAG_NAMEDTUPLE, [_ref(t, path), _prepare_items(obj, path, active)]) if dataclasses.is_dataclass(obj): @@ -210,6 +248,35 @@ def _prepare_object(obj: Any, t: type, path: str, active: set[int]) -> Any: ) +def _hook(call: Callable[[], Any], t: type, path: str) -> Any: + """Run an encode hook; whatever it raises becomes a TypeError that names the path.""" + try: + return call() + except Exception as e: + raise TypeError( + f"Cannot CBOR-encode object of type {t.__module__}.{t.__qualname__} at {path}: {e}" + ) from e + + +def _prepare_type_form(obj: Any, path: str, active: set[int]) -> Any: + """Generic aliases, unions and typing special forms. + + typing's spellings are normalized to the builtin ones: ``typing.List[int]`` arrives + as ``list[int]`` and ``Optional[X]`` as ``X | None`` (they compare equal). + """ + if obj is typing.Any: # not a class on Python 3.10 + return CBORTag(TAG_CLASS, "typing:Any") + origin, args = typing.get_origin(obj), typing.get_args(obj) + if origin is typing.Union or origin is types.UnionType: + return CBORTag(TAG_UNION, _prepare_items(args, path, active)) + if origin is typing.Literal: + return CBORTag(TAG_LITERAL, _prepare_items(args, path, active)) + if isinstance(origin, type): + prepared_origin = _prepare(origin, f"{path}.__origin__", active) + return CBORTag(TAG_GENERIC_ALIAS, [prepared_origin, _prepare_items(args, path, active)]) + raise TypeError(f"Cannot CBOR-encode {obj!r} at {path}: unsupported typing form") + + def _prepare_items(items: Any, path: str, active: set[int]) -> list[Any]: return [_prepare(item, f"{path}[{i}]", active) for i, item in enumerate(items)] @@ -255,13 +322,16 @@ def _lookup(target: Any, qualname: str) -> Any: return target -def _resolve(ref: str) -> Any: - module_name, _, qualname = ref.partition(":") +def _import(module_name: str, what: str) -> types.ModuleType: try: - module = importlib.import_module(module_name) + return importlib.import_module(module_name) except ImportError as e: - raise CborDecodeError(f"cannot import {module_name!r} for {ref!r}: {e}") from e - target = _lookup(module, qualname) + raise CborDecodeError(f"cannot import {module_name!r} for {what}: {e}") from e + + +def _resolve(ref: str) -> Any: + module_name, _, qualname = ref.partition(":") + target = _lookup(_import(module_name, repr(ref)), qualname) if target is None: raise CborDecodeError(f"{ref!r} does not resolve to an object") return target @@ -274,6 +344,11 @@ def _resolve_class(ref: str, base: type) -> Any: return cls +def _is_type_form(value: Any) -> bool: + """A class, ``Any`` or a parameterized form: what ``|`` and ``[]`` accept.""" + return isinstance(value, type) or value is typing.Any or typing.get_origin(value) is not None + + def _decode_tuple(value: Any, _: bool) -> Any: return tuple(value) @@ -306,8 +381,8 @@ def _decode_npscalar(value: Any, _: bool) -> Any: def _decode_enum(value: Any, _: bool) -> Any: - ref, name = value - return _resolve_class(ref, enum.Enum)[name] + ref, member_value = value + return _resolve_class(ref, enum.Enum)(member_value) def _decode_dimosmsg(value: Any, _: bool) -> Any: @@ -324,9 +399,12 @@ def _decode_dataclass(value: Any, _: bool) -> Any: if not (isinstance(cls, type) and dataclasses.is_dataclass(cls)): raise CborDecodeError(f"{ref!r} is not a dataclass") # No __init__ or __post_init__: the sender already normalized the values, and this - # also handles frozen classes and init=False fields. + # also handles frozen classes and init=False fields. Only fields come back; see the + # module docstring for derived state. obj = cls.__new__(cls) for f in dataclasses.fields(cls): + if f.name not in fields: + raise CborDecodeError(f"{ref!r} payload has no field {f.name!r}") object.__setattr__(obj, f.name, fields[f.name]) return obj @@ -345,7 +423,7 @@ def _decode_path(value: Any, _: bool) -> Any: def _decode_class(value: Any, _: bool) -> Any: cls = _resolve(value) - if not isinstance(cls, type): + if not (isinstance(cls, type) or cls is typing.Any): # Any is not a class on 3.10 raise CborDecodeError(f"{value!r} is not a class") return cls @@ -371,8 +449,14 @@ def _decode_custom(value: Any, _: bool) -> Any: def _decode_registered(value: Any, _: bool) -> Any: - ref, payload = value - return _resolve(ref)(payload) + type_ref, module_name, payload = value + # Importing the registering module runs its register_codec() call; the decoder + # itself comes from this process's registry, never from the payload. + _import(module_name, f"the codec of {type_ref!r}") + entry = _REGISTRY.get(_decode_class(type_ref, False)) + if entry is None: + raise CborDecodeError(f"{type_ref!r} has no codec registered by {module_name!r}") + return entry[1](payload) def _decode_dict(value: Any, _: bool) -> Any: @@ -381,13 +465,24 @@ def _decode_dict(value: Any, _: bool) -> Any: def _decode_generic_alias(value: Any, _: bool) -> Any: origin, args = value - return types.GenericAlias(origin, tuple(args)) + if not isinstance(origin, type): + raise CborDecodeError(f"{origin!r} is not a class and cannot be parameterized") + # Subscripting rather than types.GenericAlias(): a Generic subclass then gives back + # its own typing alias, and list[(int,)] is list[int]. + generic: Any = origin # a plain `type` is not indexable for mypy + return generic[tuple(args)] def _decode_union(value: Any, _: bool) -> Any: + if not value or not all(_is_type_form(member) for member in value): + raise CborDecodeError(f"{value!r} is not a union of types") return functools.reduce(operator.or_, value) +def _decode_literal(value: Any, _: bool) -> Any: + return typing.Literal[tuple(value)] + + def _decode_ellipsis(_value: Any, _: bool) -> Any: return Ellipsis @@ -416,5 +511,6 @@ def _reject_unknown_tag(tag: CBORTag, _: bool) -> Any: TAG_DICT: _decode_dict, TAG_GENERIC_ALIAS: _decode_generic_alias, TAG_UNION: _decode_union, + TAG_LITERAL: _decode_literal, TAG_ELLIPSIS: _decode_ellipsis, } diff --git a/dimos/protocol/encode/test_cbor.py b/dimos/protocol/encode/test_cbor.py index 7b8a38374a0..2295339fb2d 100644 --- a/dimos/protocol/encode/test_cbor.py +++ b/dimos/protocol/encode/test_cbor.py @@ -14,6 +14,7 @@ from __future__ import annotations +import collections.abc from dataclasses import dataclass, field import enum import functools @@ -21,11 +22,12 @@ from pathlib import Path import threading from types import MappingProxyType -from typing import Any, NamedTuple +import typing +from typing import Any, Generic, NamedTuple, TypeVar from cbor2 import CBORTag, dumps import numpy as np -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr import pytest from dimos.core.global_config import GlobalConfig @@ -37,6 +39,9 @@ TAG_DATACLASS, TAG_DIMOSMSG, TAG_ENUM, + TAG_GENERIC_ALIAS, + TAG_REGISTERED, + TAG_UNION, CborDecodeError, decode, encode, @@ -44,6 +49,8 @@ ) from dimos.types.robot_location import RobotLocation +T = TypeVar("T") + class Point(NamedTuple): x: int @@ -64,6 +71,51 @@ class Mode(str, enum.Enum): FAST = "fast" +class Permission(enum.Flag): + R = 1 + W = 2 + + +class Bits(enum.IntFlag): + A = 1 + B = 4 + + +class Box(Generic[T]): + pass + + +class Greeter(typing.Protocol): + """Its metaclass lives in typing, but it is an ordinary class to the codec.""" + + def hi(self) -> str: ... + + +@dataclass +class WithDerived: + seed: int = 1 + + def __post_init__(self) -> None: + self._double = self.seed * 2 + + +class WithPrivate(BaseModel): + value: int = 1 + _derived: int = PrivateAttr(default=0) + + def model_post_init(self, _context: Any) -> None: + self._derived = self.value * 2 + + +class RaisingHook: + def cbor_encode(self) -> None: + raise ValueError("no owner") + + @classmethod + def cbor_decode(cls, data: Any) -> RaisingHook: + return cls() + + @dataclass(frozen=True) class Frozen: a: int @@ -197,6 +249,8 @@ def roundtrip(value: Any) -> Any: Color.RED, Level.HIGH, Mode.FAST, + Permission.R | Permission.W, + Bits.A | Bits.B, Point(1, "y"), Frozen(1, (2, 3)), Slotted(1.5, ["t"]), @@ -312,6 +366,7 @@ def test_image_roundtrip(image: Image) -> None: def test_class_and_function_references() -> None: assert roundtrip(Outer.Inner) is Outer.Inner + assert roundtrip(Greeter) is Greeter assert roundtrip(list[float]) == list[float] assert roundtrip(dict[str, tuple[int, ...]]) == dict[str, tuple[int, ...]] assert roundtrip(int | None) == int | None @@ -430,3 +485,74 @@ def test_bad_payloads_raise_decode_error(payload: Any) -> None: def test_truncated_input_raises_decode_error() -> None: with pytest.raises(CborDecodeError): decode(encode([1, 2, 3])[:-1]) + + +@pytest.mark.parametrize( + "form, expected", + [ + (typing.List[int], list[int]), # noqa: UP006 + (typing.Dict[str, typing.Any], dict[str, typing.Any]), # noqa: UP006 + (typing.Optional[int], int | None), # noqa: UP045 + (typing.Union[int, str, None], int | str | None), + (typing.Literal["a", 1], typing.Literal["a", 1]), + (typing.Literal["a"] | None, typing.Literal["a"] | None), + (typing.Any, typing.Any), + (collections.abc.Callable[[int], str], collections.abc.Callable[[int], str]), + (Box[int], Box[int]), + (tuple[()], tuple[()]), + ], +) +def test_typing_forms_roundtrip(form: Any, expected: Any) -> None: + assert roundtrip(form) == expected + + +def test_registered_codec_payload_cannot_name_its_own_decoder() -> None: + # Before the fix this ran eval("6*7"); now only a type this process registered decodes. + for type_ref in ("builtins:eval", "builtins:exec"): + with pytest.raises(CborDecodeError, match="is not a class"): + decode(dumps(CBORTag(TAG_REGISTERED, [type_ref, "builtins", "6*7"]))) + with pytest.raises(CborDecodeError, match="no codec registered"): + decode(dumps(CBORTag(TAG_REGISTERED, ["builtins:int", "builtins", "6*7"]))) + with pytest.raises(CborDecodeError, match="cannot import"): + decode(dumps(CBORTag(TAG_REGISTERED, [f"{__name__}:Celsius", "no_such_module_x", 1.0]))) + + +@pytest.mark.parametrize( + "payload", + [ + CBORTag(TAG_UNION, [1, 2]), + CBORTag(TAG_UNION, []), + CBORTag(TAG_GENERIC_ALIAS, [1, [2]]), + CBORTag(TAG_GENERIC_ALIAS, ["not a class", []]), + ], +) +def test_type_form_payloads_must_hold_types(payload: Any) -> None: + with pytest.raises(CborDecodeError): + decode(dumps(payload)) + + +def test_missing_dataclass_field_names_class_and_field() -> None: + with pytest.raises(CborDecodeError, match=r"Frozen.*'b'"): + decode(dumps(CBORTag(TAG_DATACLASS, [f"{__name__}:Frozen", {"a": 1}]))) + + +def test_encode_only_raises_type_error() -> None: + with pytest.raises(TypeError, match="Cannot CBOR-encode"): + encode("\ud800") # a lone surrogate, which cbor2 itself rejects + with pytest.raises(TypeError, match=r"RaisingHook at value\[0\]: no owner"): + encode([RaisingHook()]) + + +def test_dataclass_derived_state_stays_behind() -> None: + # The documented rule: fields only, __post_init__ does not run again. + out = roundtrip(WithDerived(5)) + assert out.seed == 5 + assert not hasattr(out, "_double") + + +def test_pydantic_private_state_is_process_local() -> None: + model = WithPrivate(value=3) + model._derived = 99 + out = roundtrip(model) + assert out.value == 3 + assert out._derived == 6 # model_post_init ran on the receiving side diff --git a/dimos/protocol/pubsub/encoders.py b/dimos/protocol/pubsub/encoders.py index 042f75bdc50..7ef316ef100 100644 --- a/dimos/protocol/pubsub/encoders.py +++ b/dimos/protocol/pubsub/encoders.py @@ -102,7 +102,9 @@ class CborEncoderMixin(PubSubEncoderMixin[TopicT, MsgT, bytes]): """Encoder mixin using the CBOR codec (dimos.protocol.encode.cbor). Encoding an unsupported value raises TypeError out of publish(); a message - that does not decode is logged and dropped. + that does not decode is logged and dropped. RPC keeps its bodies as bytes + inside a primitives-only envelope (see pubsubrpc) so that a body this + process cannot decode is answered with the error instead of dropped here. """ def encode(self, msg: MsgT, _: TopicT) -> bytes: @@ -112,7 +114,11 @@ def decode(self, msg: bytes, topic: TopicT) -> MsgT: try: return cast("MsgT", cbor.decode(msg)) except cbor.CborDecodeError as e: - logger.error("Dropping undecodable CBOR message", topic=str(topic), error=str(e)) + logger.error( + "Dropping undecodable CBOR message: is another dimos version on the bus?", + topic=str(topic), + error=str(e), + ) raise DecodingError(str(e)) from e diff --git a/dimos/protocol/pubsub/impl/webrtc/providers/spec.py b/dimos/protocol/pubsub/impl/webrtc/providers/spec.py index 945a2d1c70e..4938f71fdbf 100644 --- a/dimos/protocol/pubsub/impl/webrtc/providers/spec.py +++ b/dimos/protocol/pubsub/impl/webrtc/providers/spec.py @@ -101,7 +101,8 @@ def shutdown_all_providers() -> None: class ProviderConfig(BaseModel): - """Picklable provider factory. Equal configs share one provider per process.""" + """Provider factory that crosses process boundaries as data. Equal configs share one + provider per process.""" model_config = {"frozen": True, "extra": "forbid"} diff --git a/dimos/protocol/rpc/pubsubrpc.py b/dimos/protocol/rpc/pubsubrpc.py index 270b84f7930..ef12616b06c 100644 --- a/dimos/protocol/rpc/pubsubrpc.py +++ b/dimos/protocol/rpc/pubsubrpc.py @@ -25,9 +25,11 @@ Generic, TypedDict, TypeVar, + cast, ) from dimos.constants import LCM_MAX_CHANNEL_NAME_LENGTH +from dimos.protocol.encode.cbor import decode, encode from dimos.protocol.pubsub.impl.lcmpubsub import CborLCM, Topic from dimos.protocol.pubsub.spec import PubSub from dimos.protocol.rpc.rpc_utils import deserialize_exception, serialize_exception @@ -50,16 +52,38 @@ MsgGen = Callable[[str, list], MsgT] # type: ignore[type-arg] +# Bump when the envelope below changes. A peer on another version is logged and +# ignored instead of being misparsed. +RPC_WIRE_VERSION = 1 + + class RPCReq(TypedDict): + v: int id: float | None name: str - args: Args + # encode((args, kwargs)). The envelope holds only primitives, so it always decodes; + # the handler decodes the body and answers a bad one with the error instead of the + # pubsub dropping the whole request. + args: bytes class RPCRes(TypedDict, total=False): + v: int id: float - res: Any - exception: dict[str, Any] | None # Contains exception info: type, message, traceback + res: bytes # encode(result) + exception: bytes # encode(serialize_exception(e)) + + +def _decode_reply(res: RPCRes) -> Any: + """The result or the remote exception, or the decode failure itself so the caller + raises instead of timing out.""" + try: + exc_data = res.get("exception") + if exc_data is not None: + return deserialize_exception(decode(exc_data)) + return decode(res["res"]) + except Exception as e: + return e class PubSubRPCMixin(RPCSpec, PubSub[TopicT, MsgT], Generic[TopicT, MsgT]): @@ -90,17 +114,17 @@ def __init__( @abstractmethod def topicgen(self, name: str, req_or_res: bool) -> TopicT: ... - def _encodeRPCReq(self, req: RPCReq) -> dict[str, Any]: - return dict(req) - - def _decodeRPCRes(self, msg: dict[Any, Any]) -> RPCRes: - return msg # type: ignore[return-value] - - def _encodeRPCRes(self, res: RPCRes) -> dict[str, Any]: - return dict(res) - - def _decodeRPCReq(self, msg: dict[Any, Any]) -> RPCReq: - return msg # type: ignore[return-value] + def _wire_version_ok(self, msg: Any, topic: TopicT) -> bool: + version = msg.get("v") if isinstance(msg, dict) else None + if version == RPC_WIRE_VERSION: + return True + logger.error( + "Ignoring RPC message from another wire version: is another dimos version on the bus?", + topic=str(topic), + version=version, + expected=RPC_WIRE_VERSION, + ) + return False def _get_call_thread_pool(self) -> ThreadPoolExecutor: return self._call_pool.get() @@ -141,12 +165,16 @@ def call_cb(self, name: str, arguments: Args, cb: Callable[..., Any]) -> Any: topic_req = self.topicgen(name, False) topic_res = self.topicgen(name, True) + # Encode before registering the callback: an unencodable argument raises to + # the caller and leaves nothing pending. + body = encode(arguments) + # Generate unique msg_id: timestamp + counter for concurrent calls with self._msg_id_lock: self._msg_id_counter += 1 msg_id = time.time() + (self._msg_id_counter / 1_000_000) - req: RPCReq = {"name": name, "args": arguments, "id": msg_id} + req: RPCReq = {"v": RPC_WIRE_VERSION, "name": name, "args": body, "id": msg_id} # Get or create shared subscription for this RPC's response topic topic_res_key = str(topic_res) @@ -155,8 +183,10 @@ def call_cb(self, name: str, arguments: Args, cb: Callable[..., Any]) -> Any: # Create shared handler that routes to callbacks by msg_id callbacks_dict: dict[float, Callable[..., Any]] = {} - def shared_response_handler(msg: MsgT, _: TopicT) -> None: - res = self._decodeRPCRes(msg) # type: ignore[arg-type] + def shared_response_handler(msg: MsgT, topic: TopicT) -> None: + if not self._wire_version_ok(msg, topic): + return + res = cast("RPCRes", msg) res_id = res.get("id") if res_id is None: return @@ -168,19 +198,7 @@ def shared_response_handler(msg: MsgT, _: TopicT) -> None: if callback is None: return # No callback registered (already handled or timed out) - # Check if response contains an exception - exc_data = res.get("exception") - if exc_data: - # Reconstruct the exception and pass it to the callback - from typing import cast - - from dimos.protocol.rpc.rpc_utils import SerializedException - - exc = deserialize_exception(cast("SerializedException", exc_data)) - callback(exc) - else: - # Normal response - pass the result - callback(res.get("res")) + callback(_decode_reply(res)) # Create single shared subscription unsub = self.subscribe(topic_res, shared_response_handler) @@ -197,10 +215,8 @@ def unsubscribe_callback() -> None: _, callbacks_dict = self._response_subs[topic_res_key] callbacks_dict.pop(msg_id, None) - # Publish request; an unencodable argument raises to the caller and leaves - # nothing pending. try: - self.publish(topic_req, self._encodeRPCReq(req)) # type: ignore[arg-type] + self.publish(topic_req, cast("MsgT", req)) except Exception: unsubscribe_callback() raise @@ -209,8 +225,8 @@ def unsubscribe_callback() -> None: def call_nowait(self, name: str, arguments: Args) -> None: topic_req = self.topicgen(name, False) - req: RPCReq = {"name": name, "args": arguments, "id": None} - self.publish(topic_req, self._encodeRPCReq(req)) # type: ignore[arg-type] + req: RPCReq = {"v": RPC_WIRE_VERSION, "name": name, "args": encode(arguments), "id": None} + self.publish(topic_req, cast("MsgT", req)) def serve_rpc(self, f: FunctionType, name: str | None = None): # type: ignore[no-untyped-def, override] if not name: @@ -219,36 +235,38 @@ def serve_rpc(self, f: FunctionType, name: str | None = None): # type: ignore[n topic_req = self.topicgen(name, False) topic_res = self.topicgen(name, True) - def receive_call(msg: MsgT, _: TopicT) -> None: - req = self._decodeRPCReq(msg) # type: ignore[arg-type] + def receive_call(msg: MsgT, topic: TopicT) -> None: + if not self._wire_version_ok(msg, topic): + return + req = cast("RPCReq", msg) if req.get("name") != name: return - args = req.get("args") - if args is None: + body = req.get("args") + if body is None: return + req_id = req.get("id") # Execute RPC handler in a separate thread to avoid deadlock when - # the handler makes nested RPC calls. + # the handler makes nested RPC calls. The body is decoded here so that a + # payload this process cannot decode is answered with the error. def execute_and_respond() -> None: try: + args = decode(body) response = f(*args[0], **args[1]) - req_id = req.get("id") if req_id is not None: - self.publish(topic_res, self._encodeRPCRes({"id": req_id, "res": response})) # type: ignore[arg-type] - + res: RPCRes = {"v": RPC_WIRE_VERSION, "id": req_id, "res": encode(response)} + self.publish(topic_res, cast("MsgT", res)) except Exception as e: logger.exception(f"Exception in RPC handler for {name}: {e}", exc_info=e) - # Send exception data to client if this was a request with an ID - req_id = req.get("id") if req_id is not None: - exc_data = serialize_exception(e) - # Type ignore: SerializedException is compatible with dict[str, Any] - self.publish( - topic_res, - self._encodeRPCRes({"id": req_id, "exception": exc_data}), # type: ignore[arg-type, typeddict-item] - ) + failure: RPCRes = { + "v": RPC_WIRE_VERSION, + "id": req_id, + "exception": encode(serialize_exception(e)), + } + self.publish(topic_res, cast("MsgT", failure)) # Always use thread pool to execute RPC handlers (prevents deadlock) self._get_call_thread_pool().submit(execute_and_respond) diff --git a/dimos/protocol/rpc/test_spec.py b/dimos/protocol/rpc/test_spec.py index f6e53bb26f3..74aea56efe1 100644 --- a/dimos/protocol/rpc/test_spec.py +++ b/dimos/protocol/rpc/test_spec.py @@ -29,8 +29,9 @@ import pytest from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.protocol.rpc.pubsubrpc import LCMRPC -from dimos.protocol.rpc.rpc_utils import RemoteError +from dimos.protocol.encode.cbor import decode, encode +from dimos.protocol.rpc.pubsubrpc import LCMRPC, RPC_WIRE_VERSION +from dimos.protocol.rpc.rpc_utils import RemoteError, deserialize_exception from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, RPCSpec from dimos.protocol.rpc.zenohrpc import ZenohRPC from dimos.protocol.service.zenohservice import ZenohSessionPool @@ -477,17 +478,78 @@ def test_unencodable_exception_argument_arrives_as_text(rpc_context, impl_name: unsub() -def test_zenoh_reply_decode_failure_raises_instead_of_timing_out(monkeypatch) -> None: +@pytest.mark.parametrize( + "rpc_context, codec_module", + [ + (lcm_rpc_context, "dimos.protocol.rpc.pubsubrpc"), + (zenoh_rpc_context, "dimos.protocol.rpc.zenohrpc"), + ], + ids=["lcm", "zenoh"], +) +def test_reply_decode_failure_raises_instead_of_timing_out( + rpc_context, codec_module: str, monkeypatch +) -> None: def corrupt(_payload: bytes) -> Any: raise ValueError("corrupt reply") - with zenoh_rpc_context() as (server, client): + with rpc_context() as (server, client): unsub = server.serve_rpc(add_function, "add") try: - monkeypatch.setattr("dimos.protocol.rpc.zenohrpc.decode", corrupt) + monkeypatch.setattr(f"{codec_module}.decode", corrupt) started = time.monotonic() with pytest.raises(ValueError, match="corrupt reply"): client.call_sync("add", ([1, 2], {}), rpc_timeout=10.0) assert time.monotonic() - started < 5.0 finally: unsub() + + +def test_lcm_undecodable_request_is_answered_with_the_error() -> None: + """The request id lives in the envelope, so a bad body gets an error reply, not silence.""" + with lcm_rpc_context() as (server, client): + unsub = server.serve_rpc(add_function, "add") + replies: list[Any] = [] + event = threading.Event() + + def on_reply(res: Any, _topic: Any) -> None: + replies.append(res) + event.set() + + unsub_reply = client.subscribe(client.topicgen("add", True), on_reply) + try: + request = {"v": RPC_WIRE_VERSION, "id": 1.5, "name": "add", "args": b"\xff\xff"} + client.publish(client.topicgen("add", False), request) + assert event.wait(5.0) + finally: + unsub_reply() + unsub() + (reply,) = replies + assert reply["id"] == 1.5 + error = deserialize_exception(decode(reply["exception"])) + assert isinstance(error, RemoteError) + assert "CborDecodeError" in error.remote_type + + +def test_lcm_request_from_another_wire_version_is_ignored() -> None: + calls: list[int] = [] + + def record() -> str: + calls.append(1) + return "ok" + + with lcm_rpc_context() as (server, client): + unsub = server.serve_rpc(record, "versioned") + try: + foreign = { + "v": RPC_WIRE_VERSION + 1, + "id": 1.0, + "name": "versioned", + "args": encode(([], {})), + } + client.publish(client.topicgen("versioned", False), foreign) + # Same topic, so this request is handled after the foreign one was dropped. + result, _ = client.call_sync("versioned", ([], {}), rpc_timeout=5.0) + assert result == "ok" + assert calls == [1] + finally: + unsub() diff --git a/dimos/protocol/service/lcmservice.py b/dimos/protocol/service/lcmservice.py index d915fd90871..2b241bae315 100644 --- a/dimos/protocol/service/lcmservice.py +++ b/dimos/protocol/service/lcmservice.py @@ -90,7 +90,8 @@ def __init__(self, **kwargs: Any) -> None: self._call_pool = CallPool(max_workers=4) # Services ride into workers as blueprint config (``pubsubs=[LCM(...)]``): only - # their settings cross, and a copy is a fresh, unstarted service. + # their settings cross, and a copy is a fresh, unstarted service. RPC services are + # never sent, so their rpc_timeouts have no place here. def cbor_encode(self) -> dict[str, Any]: return self.config.model_dump(exclude={"lcm"}) diff --git a/dimos/protocol/service/test_lcmservice.py b/dimos/protocol/service/test_lcmservice.py index 854f740a66f..5457d43adc3 100644 --- a/dimos/protocol/service/test_lcmservice.py +++ b/dimos/protocol/service/test_lcmservice.py @@ -14,6 +14,7 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager +import copy import errno import os import threading @@ -23,6 +24,7 @@ from lcm import LCM import pytest +from dimos.protocol.encode.cbor import decode, encode from dimos.protocol.pubsub.impl.lcmpubsub import Topic from dimos.protocol.service.lcmservice import ( _DEFAULT_LCM_URL, @@ -395,6 +397,23 @@ def test_start_and_stop(self) -> None: time.sleep(0.1) assert service._thread is None + def test_codec_round_trip_and_deepcopy_rebuild_an_unstarted_service(self) -> None: + with patch("dimos.protocol.service.lcmservice.lcm_mod.LCM") as mock_lcm_class: + mock_lcm_class.return_value = create_autospec(LCM, spec_set=True, instance=True) + service = LCMService(url="udpm://239.255.76.68:7799?ttl=1", ttl=1) + service.start() + try: + assert "lcm" not in service.cbor_encode() + for clone in (decode(encode(service)), copy.deepcopy(service)): + assert type(clone) is LCMService + assert clone.config == service.config + assert clone._thread is None + clone.stop() + # The service plus one native instance per clone. + assert mock_lcm_class.call_count == 3 + finally: + service.stop() + def test_stop_cleans_up_lcm_instance(self) -> None: with patch("dimos.protocol.service.lcmservice.lcm_mod.LCM") as mock_lcm_class: mock_lcm_instance = create_autospec(LCM, spec_set=True, instance=True) diff --git a/dimos/protocol/service/test_zenohservice.py b/dimos/protocol/service/test_zenohservice.py index 03784ec2dd6..b52bf57c356 100644 --- a/dimos/protocol/service/test_zenohservice.py +++ b/dimos/protocol/service/test_zenohservice.py @@ -15,6 +15,7 @@ from __future__ import annotations from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError +import copy import json import os import threading @@ -24,6 +25,8 @@ import pytest import zenoh +from dimos.protocol.encode.cbor import decode, encode +from dimos.protocol.pubsub.impl.zenohpubsub import ZenohPubSubBase from dimos.protocol.service import zenohservice from dimos.protocol.service.zenohservice import ZenohConfig, ZenohService, ZenohSessionPool @@ -158,6 +161,37 @@ def test_different_modes_produce_different_keys() -> None: assert peer.session_key != client.session_key +def test_codec_round_trip_sheds_the_session(session_pool) -> None: + """A service travels to its worker over the codec and re-acquires on the far side.""" + svc = ZenohService(session_pool=session_pool) + svc.start() + clone = decode(encode(svc)) + + assert clone._session is None + assert clone._session_pool is zenohservice.default_session_pool + assert clone.config.session_key == svc.config.session_key + + clone._session_pool = session_pool # keep the test on its own pool + clone.start() + assert clone.session is not None + clone.stop() + + +def test_deepcopy_is_a_fresh_unstarted_service(session_pool) -> None: + svc = ZenohService(session_pool=session_pool) + svc.start() + clone = copy.deepcopy(svc) + assert clone._session is None + assert clone.config == svc.config + + +def test_pubsub_round_trip_rebuilds_its_runtime(session_pool) -> None: + pubsub = decode(encode(ZenohPubSubBase(session_pool=session_pool))) + assert pubsub._publishers == {} + assert pubsub._subscribers == [] + assert pubsub._drain_stops == [] + + def test_start_creates_session(session_pool) -> None: svc = ZenohService(session_pool=session_pool) svc.start() diff --git a/dimos/protocol/service/zenohservice.py b/dimos/protocol/service/zenohservice.py index 20742c20144..9e4ef15ba91 100644 --- a/dimos/protocol/service/zenohservice.py +++ b/dimos/protocol/service/zenohservice.py @@ -405,7 +405,8 @@ def __init__(self, *, session_pool: ZenohSessionPool | None = None, **kwargs: An # Services ride into workers as blueprint config (``pubsubs=[Zenoh(...)]``): only # their settings cross, the far side uses the default session pool, and a copy - # is a fresh, unstarted service. + # is a fresh, unstarted service. RPC services are never sent, so their + # rpc_timeouts have no place here. def cbor_encode(self) -> dict[str, Any]: return self.config.model_dump() diff --git a/dimos/robot/assets/model.py b/dimos/robot/assets/model.py index f004e0ad0ee..37844c955c8 100644 --- a/dimos/robot/assets/model.py +++ b/dimos/robot/assets/model.py @@ -336,7 +336,11 @@ def _loaded(self) -> LoadedRobotModel: return LoadedRobotModel(xml, source_path, package_paths) def __getstate__(self) -> dict[str, object]: - """Exclude materialized XML when configurations cross worker processes.""" + """Drop the memoized model when copy.deepcopy clones a blueprint config value. + + Worker transfer goes through the CBOR codec, which sends dataclass fields only, + so the cache never crosses that way either. + """ state = dict(self.__dict__) state.pop("_loaded", None) return state diff --git a/dimos/robot/assets/source.py b/dimos/robot/assets/source.py index 09bf37201cd..328bd3dc125 100644 --- a/dimos/robot/assets/source.py +++ b/dimos/robot/assets/source.py @@ -77,7 +77,8 @@ def cbor_decode(cls, data: dict[str, str]) -> RobotDescriptionSource: return cls(data["url"], data["ref"]) def __reduce__(self) -> tuple[type[RobotDescriptionSource], tuple[str, str]]: - """Recreate the source without serializing checkout state.""" + """For copy.deepcopy, which clones blueprint config values before deployment: + recreate the source without its checkout state. Workers get cbor_encode().""" return (RobotDescriptionSource, (self.url, self.ref)) @@ -119,11 +120,11 @@ def __getattribute__(self, name: str) -> object: except AttributeError: return object.__getattribute__(self, name) + # The CBOR codec is not listed: it looks cbor_encode up on the class. if name.startswith("_robot_description_") or name in { "_resolve", "__reduce__", "__reduce_ex__", - "cbor_encode", }: return object.__getattribute__(self, name) @@ -183,7 +184,8 @@ def cbor_decode(cls, data: dict[str, Any]) -> RobotDescriptionPath: return cls(data["source"], Path(data["relative_path"])) def __reduce__(self) -> tuple[type[RobotDescriptionPath], tuple[RobotDescriptionSource, Path]]: - """Preserve the lazy source handle across worker serialization.""" + """For copy.deepcopy of blueprint config values: keep the lazy source handle + instead of resolving the checkout. Workers get cbor_encode().""" source: RobotDescriptionSource = object.__getattribute__(self, "_robot_description_source") relative_path: Path = object.__getattribute__(self, "_robot_description_relative_path") return (RobotDescriptionPath, (source, relative_path)) diff --git a/dimos/robot/galaxea/r1pro/blueprints/basic/r1pro_coordinator.py b/dimos/robot/galaxea/r1pro/blueprints/basic/r1pro_coordinator.py index 5a79cc60dbc..d172d5ddc22 100644 --- a/dimos/robot/galaxea/r1pro/blueprints/basic/r1pro_coordinator.py +++ b/dimos/robot/galaxea/r1pro/blueprints/basic/r1pro_coordinator.py @@ -31,8 +31,9 @@ from dimos.control.components import HardwareComponent, HardwareType, make_twist_base_joints from dimos.control.coordinator import ControlCoordinator, TaskConfig from dimos.control.tasks.trajectory_task.trajectory_task import joint_trajectory_task -from dimos.core.coordination.blueprints import Blueprint, TransportSpec, autoconnect +from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.core.global_config import global_config +from dimos.core.stream import TransportSpec from dimos.core.transport import ZenohTransport from dimos.hardware.sensors.lidar.pointlio.module import PointLio from dimos.hardware.sensors.lidar.pointlio.pointlio_blueprints import mid360_for_pointlio diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_holonomic_benchmark.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_holonomic_benchmark.py index 66d72e49295..84cfff3460c 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_holonomic_benchmark.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_holonomic_benchmark.py @@ -38,8 +38,8 @@ from typing import Any from dimos.control.benchmarking.benchmark import Benchmarker -from dimos.core.coordination.blueprints import TransportSpec, autoconnect -from dimos.core.stream import Transport +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.stream import Transport, TransportSpec from dimos.core.transport import LCMTransport from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_holonomic_controller import ( diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_rpp_benchmark.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_rpp_benchmark.py index f85b71bcd32..39d3ada54e9 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_rpp_benchmark.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_rpp_benchmark.py @@ -51,8 +51,8 @@ from typing import Any from dimos.control.benchmarking.benchmark import Benchmarker -from dimos.core.coordination.blueprints import TransportSpec, autoconnect -from dimos.core.stream import Transport +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.stream import Transport, TransportSpec from dimos.core.transport import LCMTransport from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_rpp_controller import ( diff --git a/dimos/utils/data.py b/dimos/utils/data.py index 5b5a6946b71..90075be47d5 100644 --- a/dimos/utils/data.py +++ b/dimos/utils/data.py @@ -380,12 +380,12 @@ def __getattribute__(self, name: str) -> object: return object.__getattribute__(self, name) # Copying and serializer/type introspection must inspect the lazy wrapper, - # not materialize the asset (including probes for absent attributes). + # not materialize the asset (including probes for absent attributes). The CBOR + # codec is not listed: it looks cbor_encode up on the class, never the instance. if name in ( "_lfs_filename", "_lfs_resolved_cache", "_ensure_downloaded", - "cbor_encode", "__class__", "__copy__", "__deepcopy__", diff --git a/dimos/web/cockpit.py b/dimos/web/cockpit.py index ecefb7c9ebd..34b1ed305b0 100644 --- a/dimos/web/cockpit.py +++ b/dimos/web/cockpit.py @@ -82,8 +82,9 @@ def _check_rate(name: str, value: float) -> None: class _FrozenDict(dict[str, Any]): """Read-only dict for Channel.params: a frozen dataclass must not leak a mutable mapping whose edits would change what a later cockpit() emits. - A dict subclass (not mappingproxy) so it stays picklable and - json.dumps-able.""" + A dict subclass (not mappingproxy) so it stays json.dumps-able and + copyable: __reduce__ serves copy.deepcopy (blueprint config values are + deep-copied before deployment), cbor_encode the worker pipe.""" def __reduce__(self) -> tuple[Any, ...]: return (_FrozenDict, (dict(self),)) From 8e34cc0463d0b13c08968ffcf20a98a48bfad397 Mon Sep 17 00:00:00 2001 From: Paul Nechifor Date: Sun, 25 Oct 2026 08:05:26 +0200 Subject: [PATCH 3/3] fix(rpc): keep Image lossless over RPC and make plain lifecycle overrides fail clearly Review fixes for PR #4821. Image travelled as its LCM form, which labels every RGB image rgb8 and decodes through np.frombuffer: float32 RGB images failed to reshape and uint8 images came back read-only. Image now has cbor_encode/cbor_decode built on its pickle state, so dtype, shape, format, timestamp and writability survive. Module.rpcs only sees methods carrying the @rpc marker, so a plain start() override drops out of the proxy's RPC set and the host fetched the bound method from the worker, which CBOR rejects (pickle shipped a module copy and ran start() in the parent). DemoRobot and the image encoding demo modules get @rpc, the proxy raises an AttributeError that names the missing decorator, and a codebase check fails any built-in blueprint whose module overrides an @rpc method without it. The blueprint codec test skipped on any exception; it now uses load_blueprint and skips only on OptionalDependencyError, with the self_hosted marks and robot_ips monkeypatch the sibling tests use. --- AGENTS.md | 2 - dimos/agents/skills/demo_robot.py | 3 + dimos/codebase_checks/test_rpc_overrides.py | 73 +++++++++++++++++++ dimos/core/rpc_client.py | 6 ++ dimos/core/test_rpc_client.py | 16 ++++ dimos/msgs/sensor_msgs/Image.py | 9 +++ dimos/protocol/encode/test_cbor.py | 3 + .../robot/test_blueprints_cross_processes.py | 28 ++++--- dimos/utils/demo_image_encoding.py | 5 ++ 9 files changed, 134 insertions(+), 11 deletions(-) create mode 100644 dimos/codebase_checks/test_rpc_overrides.py diff --git a/AGENTS.md b/AGENTS.md index 5c65c7743b3..f45f0f3ec3d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -321,8 +321,6 @@ class MySkillContainer(Module): If multiple modules match the spec, use `.remappings()` to resolve. Source: `dimos/spec/utils.py`, `dimos/core/coordination/blueprints.py`. -RPC arguments, results and worker-pipe messages are CBOR (`dimos/protocol/encode/cbor.py`), not pickle. JSON-like values, tuples, sets, numpy arrays, enums, dataclasses (their fields only: state that `__post_init__` derives stays behind), pydantic models (private attributes stay behind), `Path`, message types with `lcm_encode`, classes and module-level functions cross as themselves. Anything else raises `TypeError` at the call site; give such a class `cbor_encode()` and a `cbor_decode()` classmethod. - ### Adding a New Skill 1. Pick the right container (robot-specific or `dimos/agents/skills/`). diff --git a/dimos/agents/skills/demo_robot.py b/dimos/agents/skills/demo_robot.py index 9e7ac8433ba..674ffc2c18e 100644 --- a/dimos/agents/skills/demo_robot.py +++ b/dimos/agents/skills/demo_robot.py @@ -15,6 +15,7 @@ from reactivex import interval +from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import Out from dimos.mapping.models import LatLon @@ -23,10 +24,12 @@ class DemoRobot(Module): gps_location: Out[LatLon] + @rpc def start(self) -> None: super().start() self.register_disposable(interval(1.0).subscribe(lambda _: self._publish_gps_location())) + @rpc def stop(self) -> None: super().stop() diff --git a/dimos/codebase_checks/test_rpc_overrides.py b/dimos/codebase_checks/test_rpc_overrides.py new file mode 100644 index 00000000000..026a6af7784 --- /dev/null +++ b/dimos/codebase_checks/test_rpc_overrides.py @@ -0,0 +1,73 @@ +# 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. + +"""An override of an @rpc method must carry the decorator itself. + +``Module.rpcs`` reads the ``__rpc__`` marker off each attribute, so a plain override +hides the base method from the module's proxy. The host then fetches it with a +GetAttrRequest, and the worker cannot send a bound method back: ``start()`` on such a +module fails at deploy time. +""" + +import pytest + +from dimos.core.global_config import global_config +from dimos.core.module import ModuleBase +from dimos.robot.all_blueprints import all_blueprints +from dimos.robot.get_all_blueprints import OptionalDependencyError, load_blueprint +from dimos.robot.test_all_blueprints import SELF_HOSTED_BLUEPRINTS + + +def _plain_overrides(module: type[ModuleBase]) -> list[str]: + """Public methods without @rpc that some base class declares with it.""" + return sorted( + name + for name in dir(module) + if not name.startswith("_") + and callable(getattr(module, name, None)) + and not hasattr(getattr(module, name), "__rpc__") + and any(hasattr(vars(base).get(name), "__rpc__") for base in module.__mro__) + ) + + +@pytest.mark.parametrize( + "blueprint_name", + [ + pytest.param(name, marks=pytest.mark.self_hosted) + if name in SELF_HOSTED_BLUEPRINTS + else name + for name in sorted(all_blueprints) + ], +) +def test_rpc_overrides_keep_the_decorator( + blueprint_name: str, monkeypatch: pytest.MonkeyPatch +) -> None: + # The multi-robot blueprints read ROBOT_IPS at import time. + monkeypatch.setattr(global_config, "robot_ips", "192.0.2.10,192.0.2.11") + try: + blueprint = load_blueprint(blueprint_name) + except OptionalDependencyError as e: + pytest.skip(f"Skipping due to missing optional dependency: {e}") + + violations = sorted( + { + f"{atom.module.__module__}.{atom.module.__name__}.{name}" + for atom in blueprint.blueprints + for name in _plain_overrides(atom.module) + } + ) + assert not violations, ( + f"Blueprint {blueprint_name!r} has @rpc methods overridden without the decorator; " + f"add @rpc to: {violations}" + ) diff --git a/dimos/core/rpc_client.py b/dimos/core/rpc_client.py index 528a89cb8b4..5de7ccf9cab 100644 --- a/dimos/core/rpc_client.py +++ b/dimos/core/rpc_client.py @@ -223,6 +223,12 @@ def __getattr__(self, name: str): # type: ignore[no-untyped-def] f"(remote-mode), so non-@rpc attribute access is unavailable." ) + if inspect.isfunction(getattr(self.actor_class, name, None)): + raise AttributeError( + f"{self.remote_name}.{name} is a method without @rpc, so the worker cannot " + "send it back; decorate it with @rpc to call it from the host" + ) + # return super().__getattr__(name) # Try to avoid recursion by directly accessing attributes that are known result = self.actor_instance.__getattr__(name) diff --git a/dimos/core/test_rpc_client.py b/dimos/core/test_rpc_client.py index dff5f397b63..da5a93290dc 100644 --- a/dimos/core/test_rpc_client.py +++ b/dimos/core/test_rpc_client.py @@ -17,10 +17,13 @@ from IPython.core.completer import provisionalcompleter from IPython.core.interactiveshell import InteractiveShell import pytest +from pytest_mock import MockerFixture from traitlets.config import Config from dimos.core.coordination.blueprints import DisabledModuleProxy +from dimos.core.coordination.python_worker import Actor from dimos.core.demos.stress_test_module import StressTestModule +from dimos.core.module import Module from dimos.core.rpc_client import AsyncSpecProxy, RPCClient from dimos.protocol.encode.cbor import decode, encode from dimos.protocol.rpc.spec import RPCSpec @@ -58,6 +61,19 @@ def test_dir_exposes_rpcs_and_proxy_attributes_without_transport_calls(proxy): assert transport.mock_calls == [] +class PlainStart(Module): + def start(self) -> None: # no @rpc, so it is missing from PlainStart.rpcs + super().start() + + +def test_plain_method_override_is_a_clear_error(mocker: MockerFixture) -> None: + actor = Actor(None, PlainStart, worker_id=0) + client = RPCClient(actor, PlainStart, rpc=mocker.Mock(spec=RPCSpec)) + + with pytest.raises(AttributeError, match="without @rpc"): + client.start # noqa: B018 + + def test_ipython_completes_rpc_without_transport_calls(ipython, proxy): _, transport = proxy diff --git a/dimos/msgs/sensor_msgs/Image.py b/dimos/msgs/sensor_msgs/Image.py index fb79e10a44d..9f064ff8068 100644 --- a/dimos/msgs/sensor_msgs/Image.py +++ b/dimos/msgs/sensor_msgs/Image.py @@ -114,6 +114,15 @@ def __setstate__(self, state: dict[str, Any]) -> None: self.frame_id = state.get("frame_id", "") self.ts = state.get("ts", time.time()) + def cbor_encode(self) -> dict[str, Any]: + # Not the LCM form: that fixes the dtype per format (a float RGB image would come + # back as uint8) and decodes into read-only arrays. + return self.__getstate__() + + @classmethod + def cbor_decode(cls, state: dict[str, Any]) -> Image: + return cls(**state) + @property def height(self) -> int: return int(self.data.shape[0]) diff --git a/dimos/protocol/encode/test_cbor.py b/dimos/protocol/encode/test_cbor.py index 2295339fb2d..e923e432127 100644 --- a/dimos/protocol/encode/test_cbor.py +++ b/dimos/protocol/encode/test_cbor.py @@ -350,6 +350,8 @@ def test_global_config_roundtrip() -> None: Image(data=np.arange(4 * 6 * 3, dtype=np.uint8).reshape(4, 6, 3), format=ImageFormat.RGB), Image(data=np.full((2, 3), 700, dtype=np.uint16), format=ImageFormat.DEPTH16), Image(data=np.full((2, 3), 0.25, dtype=np.float32), format=ImageFormat.DEPTH), + # The LCM form would label this rgb8 and decode it as uint8. + Image(data=np.full((2, 3, 3), 0.5, dtype=np.float32), format=ImageFormat.RGB), ], ) def test_image_roundtrip(image: Image) -> None: @@ -362,6 +364,7 @@ def test_image_roundtrip(image: Image) -> None: assert out.ts == 1.5 assert out.data.dtype == image.data.dtype assert np.array_equal(out.data, image.data) + assert out.data.flags.writeable def test_class_and_function_references() -> None: diff --git a/dimos/robot/test_blueprints_cross_processes.py b/dimos/robot/test_blueprints_cross_processes.py index 3826a9acdc7..de9b2d9b7f1 100644 --- a/dimos/robot/test_blueprints_cross_processes.py +++ b/dimos/robot/test_blueprints_cross_processes.py @@ -19,20 +19,30 @@ encodable (see dimos/protocol/encode/cbor.py). """ -import importlib - import pytest +from dimos.core.global_config import global_config from dimos.protocol.encode.cbor import decode, encode from dimos.robot.all_blueprints import all_blueprints - - -@pytest.mark.parametrize("name", sorted(all_blueprints)) -def test_blueprint_roundtrips_through_the_codec(name: str) -> None: - module_path, attr = all_blueprints[name].split(":") +from dimos.robot.get_all_blueprints import OptionalDependencyError, load_blueprint +from dimos.robot.test_all_blueprints import SELF_HOSTED_BLUEPRINTS + + +@pytest.mark.parametrize( + "name", + [ + pytest.param(name, marks=pytest.mark.self_hosted) + if name in SELF_HOSTED_BLUEPRINTS + else name + for name in sorted(all_blueprints) + ], +) +def test_blueprint_roundtrips_through_the_codec(name: str, monkeypatch: pytest.MonkeyPatch) -> None: + # The multi-robot blueprints read ROBOT_IPS at import time. + monkeypatch.setattr(global_config, "robot_ips", "192.0.2.10,192.0.2.11") try: - blueprint = getattr(importlib.import_module(module_path), attr) - except Exception as error: # an optional dependency or robot setting this machine lacks + blueprint = load_blueprint(name) + except OptionalDependencyError as error: pytest.skip(f"{name} is not importable here: {error}") restored = decode(encode(blueprint)) diff --git a/dimos/utils/demo_image_encoding.py b/dimos/utils/demo_image_encoding.py index b74f227066b..97e66a2e4ae 100644 --- a/dimos/utils/demo_image_encoding.py +++ b/dimos/utils/demo_image_encoding.py @@ -31,6 +31,7 @@ from reactivex.disposable import Disposable from dimos.core.coordination.module_coordinator import ModuleCoordinator +from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In, Out from dimos.core.transport import JpegLcmTransport, LCMTransport @@ -44,12 +45,14 @@ class EmitterModule(Module): _thread: threading.Thread | None = None _stop_event: threading.Event | None = None + @rpc def start(self) -> None: super().start() self._stop_event = threading.Event() self._thread = threading.Thread(target=self._publish_image, daemon=True) self._thread.start() + @rpc def stop(self) -> None: if self._thread: self._stop_event.set() # type: ignore[union-attr] @@ -73,11 +76,13 @@ class ReceiverModule(Module): _open_file = None + @rpc def start(self) -> None: super().start() self.register_disposable(Disposable(self.image.subscribe(self._on_image))) self._open_file = open("/tmp/receiver-times", "w") + @rpc def stop(self) -> None: self._open_file.close() # type: ignore[union-attr] super().stop()