diff --git a/AGENTS.md b/AGENTS.md index 3b7e2371262..02e6ade33cb 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 (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/mcp/mcp_client.py b/dimos/agents/mcp/mcp_client.py index 6d50e23fdee..f3103ccce5d 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/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/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 d5943551848..c48cc651754 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: requirement_error = ffmpeg_requirement() 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/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/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 dbdaca766b3..92640cfa559 100644 --- a/dimos/core/coordination/blueprints.py +++ b/dimos/core/coordination/blueprints.py @@ -13,22 +13,30 @@ # 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 pydantic import BaseModel +from typing import ( + TYPE_CHECKING, + Any, + Literal, + Union, + cast, + get_args, + get_origin, + get_type_hints, +) 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 @@ -57,8 +65,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"" @@ -176,30 +188,7 @@ 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) - - -# 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 +223,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 +511,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 78c247d2458..3d7e7a52e52 100644 --- a/dimos/core/coordination/module_coordinator.py +++ b/dimos/core/coordination/module_coordinator.py @@ -28,16 +28,15 @@ 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 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.stream import Transport, TransportSpec from dimos.core.transport import ( LCMTransport, PubSubTransport, @@ -98,7 +97,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 @@ -367,8 +365,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]] @@ -846,7 +846,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 659c937e525..d596c9a262b 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,34 @@ 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 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(request) - response: WorkerResponse = self._conn.recv() + self._conn.send_bytes(payload) + reply = self._conn.recv_bytes() else: - self._conn.send(request) - response = self._conn.recv() + self._conn.send_bytes(payload) + 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) @@ -202,7 +224,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 @@ -238,9 +261,11 @@ 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) + reply = self._conn.recv_bytes() + response: WorkerResponse = decode(reply) if response.error: raise RuntimeError(f"Failed to deploy module: {response.error}") @@ -265,8 +290,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}") @@ -278,8 +303,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 @@ -287,9 +312,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.", @@ -463,11 +488,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( @@ -475,18 +502,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..ebbabf707d9 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 @@ -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 @@ -1151,11 +1154,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 73167e2b46b..bef9fa1ce8a 100644 --- a/dimos/core/coordination/test_worker.py +++ b/dimos/core/coordination/test_worker.py @@ -19,12 +19,14 @@ import pytest from dimos.agents.annotation import skill +from dimos.core.coordination.python_worker import Actor, MethodCallProxy 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 @@ -54,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 @@ -146,29 +148,29 @@ def test_worker_manager_basic(create_worker_manager): module.stop() -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 @@ -468,3 +470,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..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 + 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,14 +42,20 @@ 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. + + ``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 direction: StreamDirection type_name: str channel: str - transport: bytes + type: type + transport: TransportSpec @dataclass diff --git a/dimos/core/module.py b/dimos/core/module.py index 07d345c0c4f..919fe814197 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 ( @@ -40,8 +39,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.rpc_client import RpcCall -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 @@ -82,14 +80,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 "" @@ -156,7 +153,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_thread_timeout: float = 2.0 @@ -264,33 +260,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", 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 = 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 { @@ -386,7 +355,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=TransportSpec.of(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 bef546e2664..0e9dc82fdcb 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).""" @@ -141,7 +129,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__ @@ -178,14 +166,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)) @@ -229,13 +222,20 @@ 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) - # 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) @@ -251,7 +251,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: @@ -276,14 +276,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 ebe906b5a8e..ea74ab2cb90 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, @@ -36,7 +37,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 @@ -106,6 +107,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 @@ -173,18 +215,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: @@ -201,11 +237,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]: @@ -246,10 +306,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) @@ -282,10 +344,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..da5a93290dc 100644 --- a/dimos/core/test_rpc_client.py +++ b/dimos/core/test_rpc_client.py @@ -17,10 +17,15 @@ 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.rpc_client import RPCClient +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 @@ -56,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 @@ -64,3 +82,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 9fab86480b8..af5f1d4e50c 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 +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 @@ -279,3 +279,17 @@ def test_hot_getter(dimos) -> None: robot.stop() subscriber.stop_rpc_client() robot.stop_rpc_client() + + +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 + + # 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 b25292a3b2b..b87071face0 100644 --- a/dimos/core/test_transport.py +++ b/dimos/core/test_transport.py @@ -12,24 +12,38 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pickle + +from dataclasses import dataclass import pytest +from dimos.core.stream import TransportSpec from dimos.core.transport import ( + DDS_AVAILABLE, + DDSTransport, JpegLcmTransport, LCMTransport, pLCMTransport, ) 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_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 +52,54 @@ 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() + + +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/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..33edbcd8f10 100644 --- a/dimos/core/transport.py +++ b/dimos/core/transport.py @@ -14,7 +14,8 @@ from __future__ import annotations -from functools import partial +from abc import abstractmethod +import copy import importlib.util import threading import time @@ -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,12 +92,29 @@ def channel(self) -> str: """The channel string this transport publishes and subscribes on.""" return str(self.topic) + # 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 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 + def cbor_decode(cls, data: dict[str, Any]) -> PubSubTransport[Any]: + 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: + # blueprint config values are deep-copied before deployment. + return cast( + "PubSubTransport[T]", type(self).cbor_decode(copy.deepcopy(self.cbor_encode(), memo)) + ) - return TransportSpec(cls, args, kwargs) + 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]): @@ -106,11 +124,8 @@ 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"})} def broadcast(self, _: Stream[T] | None, msg: T) -> None: if not self._started: @@ -150,11 +165,12 @@ 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"}), + } def broadcast(self, _, msg) -> None: # type: ignore[no-untyped-def] if not self._started: @@ -179,12 +195,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 @@ -197,15 +207,14 @@ 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 __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} def broadcast(self, _: Stream[DimosMsg] | None, msg: DimosMsg) -> None: if self._ros is None: @@ -225,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: @@ -240,20 +249,27 @@ 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() - 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(), + } def start(self) -> None: with self._start_lock: @@ -283,12 +299,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 +311,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 +337,12 @@ 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} def broadcast(self, _: Stream[M] | None, msg: M) -> None: if not self._started: @@ -405,12 +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) - @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 + # Constructor kwargs, like PubSubTransport.cbor_encode. + def cbor_encode(self) -> dict[str, Any]: + return {"config": self._config} - return TransportSpec(cls, args, kwargs) + @classmethod + 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) @@ -512,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() @@ -528,8 +546,8 @@ 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} def start(self) -> None: with self._start_lock: @@ -565,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() @@ -580,8 +598,8 @@ 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} def start(self) -> None: with self._start_lock: 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 750bde20dfb..3f56d9ec934 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() @@ -111,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) @@ -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/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/spec/models.py b/dimos/manipulation/planning/spec/models.py index 125b5939fb7..7dcf37f68b0 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) """ @@ -120,11 +122,24 @@ 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 - # pickled across worker RPC, and an ndarray field breaks all three. - 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 + 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/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..aa8c9549251 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 @@ -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 a pickled " - 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_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/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/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/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/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/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/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..68dcc27ccbb 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,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 = pickle.loads(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 0db92e544e0..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"]) @@ -441,8 +449,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..92cabaf3bb2 --- /dev/null +++ b/dimos/protocol/encode/cbor.py @@ -0,0 +1,516 @@ +# 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 (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 + +from collections.abc import Callable +import dataclasses +import enum +import functools +import importlib +import operator +import pathlib +import sys +import types +import typing +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 +TAG_LITERAL = 4060 + +_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``. + The codec looks both hooks up on the class, never on the instance. + """ + + def cbor_encode(self) -> Any: ... + + @classmethod + def cbor_decode(cls, data: Any) -> Any: ... + + +# 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, 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. + """ + _ref(decode, repr(decode)) + _REGISTRY[type_] = (encode, decode, decode.__module__) + + +def encode(obj: Any) -> bytes: + """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: + """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). + + 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 + 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 = _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, _, module) in _REGISTRY.items(): + if isinstance(obj, registered_type): + 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): + # 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}") + 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), _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): + 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 _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)] + + +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 _import(module_name: str, what: str) -> types.ModuleType: + try: + return importlib.import_module(module_name) + except ImportError as e: + 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 + + +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 _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) + + +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, member_value = value + return _resolve_class(ref, enum.Enum)(member_value) + + +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. 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 + + +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) or cls is typing.Any): # Any is not a class on 3.10 + 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: + 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: + return dict(value) + + +def _decode_generic_alias(value: Any, _: bool) -> Any: + origin, args = value + 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 + + +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_LITERAL: _decode_literal, + 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..e923e432127 --- /dev/null +++ b/dimos/protocol/encode/test_cbor.py @@ -0,0 +1,561 @@ +# 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 + +import collections.abc +from dataclasses import dataclass, field +import enum +import functools +import math +from pathlib import Path +import threading +from types import MappingProxyType +import typing +from typing import Any, Generic, NamedTuple, TypeVar + +from cbor2 import CBORTag, dumps +import numpy as np +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr +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, + TAG_GENERIC_ALIAS, + TAG_REGISTERED, + TAG_UNION, + CborDecodeError, + decode, + encode, + register_codec, +) +from dimos.types.robot_location import RobotLocation + +T = TypeVar("T") + + +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" + + +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 + 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, + Permission.R | Permission.W, + Bits.A | Bits.B, + 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), + # 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: + 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) + assert out.data.flags.writeable + + +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 + 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]) + + +@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/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..7ef316ef100 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,30 @@ 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. 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: + 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: is another dimos version on the bus?", + 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..4938f71fdbf 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. """ @@ -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/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..ef12616b06c 100644 --- a/dimos/protocol/rpc/pubsubrpc.py +++ b/dimos/protocol/rpc/pubsubrpc.py @@ -25,10 +25,12 @@ Generic, TypedDict, TypeVar, + cast, ) from dimos.constants import LCM_MAX_CHANNEL_NAME_LENGTH -from dimos.protocol.pubsub.impl.lcmpubsub import PickleLCM, Topic +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 from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, DEFAULT_RPC_TIMEOUTS, Args, RPCSpec @@ -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]): @@ -87,45 +111,20 @@ 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: ... - 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() @@ -166,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) @@ -180,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 @@ -193,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) @@ -215,22 +208,25 @@ 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) + try: + self.publish(topic_req, cast("MsgT", req)) + except Exception: + unsubscribe_callback() + raise + return unsubscribe_callback 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: @@ -239,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) @@ -276,7 +274,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 +283,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..74aea56efe1 100644 --- a/dimos/protocol/rpc/test_spec.py +++ b/dimos/protocol/rpc/test_spec.py @@ -19,14 +19,19 @@ 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.protocol.rpc.pubsubrpc import LCMRPC -from dimos.protocol.rpc.rpc_utils import RemoteError +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +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 @@ -382,3 +387,169 @@ 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() + + +@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 rpc_context() as (server, client): + unsub = server.serve_rpc(add_function, "add") + try: + 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/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..2b241bae315 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,18 @@ 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. 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"}) + + @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 +113,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..5457d43adc3 100644 --- a/dimos/protocol/service/test_lcmservice.py +++ b/dimos/protocol/service/test_lcmservice.py @@ -14,9 +14,9 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager +import copy import errno import os -import pickle import threading import time from unittest.mock import MagicMock, create_autospec, patch @@ -24,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, @@ -396,50 +397,22 @@ def test_start_and_stop(self) -> None: time.sleep(0.1) assert service._thread is None - def test_getstate_excludes_unpicklable_attrs(self) -> 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_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() + 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: diff --git a/dimos/protocol/service/test_zenohservice.py b/dimos/protocol/service/test_zenohservice.py index 103238aabb2..b52bf57c356 100644 --- a/dimos/protocol/service/test_zenohservice.py +++ b/dimos/protocol/service/test_zenohservice.py @@ -15,9 +15,9 @@ from __future__ import annotations from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError +import copy import json import os -import pickle import threading from types import SimpleNamespace from typing import Any @@ -25,8 +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.rpc.zenohrpc import ZenohRPC from dimos.protocol.service import zenohservice from dimos.protocol.service.zenohservice import ZenohConfig, ZenohService, ZenohSessionPool @@ -161,36 +161,36 @@ 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.""" - +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 = pickle.loads(pickle.dumps(svc)) + 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_pickled_pubsub_and_rpc_rebuild_their_runtime(session_pool) -> None: - pubsub = pickle.loads(pickle.dumps(ZenohPubSubBase(session_pool=session_pool))) +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 == [] - 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) diff --git a/dimos/protocol/service/zenohservice.py b/dimos/protocol/service/zenohservice.py index cc0f1b18ab9..9e4ef15ba91 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,19 @@ 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. 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() - 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/model.py b/dimos/robot/assets/model.py index ea481ce95d9..bb9433059cf 100644 --- a/dimos/robot/assets/model.py +++ b/dimos/robot/assets/model.py @@ -308,7 +308,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 c0bff41af60..328bd3dc125 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,8 +69,16 @@ 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.""" + """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)) @@ -112,6 +120,7 @@ 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__", @@ -164,8 +173,19 @@ 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.""" + """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/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/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/test_blueprints_cross_processes.py b/dimos/robot/test_blueprints_cross_processes.py new file mode 100644 index 00000000000..de9b2d9b7f1 --- /dev/null +++ b/dimos/robot/test_blueprints_cross_processes.py @@ -0,0 +1,53 @@ +# 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 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 +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 = load_blueprint(name) + except OptionalDependencyError as error: + 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_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_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/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 257d3e881f3..90075be47d5 100644 --- a/dimos/utils/data.py +++ b/dimos/utils/data.py @@ -380,7 +380,8 @@ 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", @@ -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/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() 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..34b1ed305b0 100644 --- a/dimos/web/cockpit.py +++ b/dimos/web/cockpit.py @@ -82,12 +82,20 @@ 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),)) + 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 ddc485c1e82..f286bf1f7bc 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 a3202a33d82..e5c22072fb2 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" @@ -1842,6 +1874,7 @@ dependencies = [ { name = "annotation-protocol" }, { name = "attrs" }, { name = "bleak" }, + { name = "cbor2" }, { name = "cmeel-tinyxml2" }, { name = "cryptography" }, { name = "dimos-lcm" }, @@ -2534,6 +2567,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" },