diff --git a/.gitignore b/.gitignore index f97d9f906a..9d9d85690f 100644 --- a/.gitignore +++ b/.gitignore @@ -71,3 +71,4 @@ CLAUDE.MD /assets/teleop_certs/ /.mcp.json +*.speedscope.json diff --git a/README.md b/README.md index 92bbed926e..e07d623300 100644 --- a/README.md +++ b/README.md @@ -204,7 +204,10 @@ See below a simple robot connection module that sends streams of continuous `cmd ```py import threading, time, numpy as np -from dimos.core import In, Module, Out, rpc, autoconnect +from dimos.core.blueprints import autoconnect +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import Twist from dimos.msgs.sensor_msgs import Image, ImageFormat @@ -249,7 +252,8 @@ Blueprints can be composed, remapped, and have transports overridden if `autocon A blueprint example that connects the image stream from a robot to an LLM Agent for reasoning and action execution. ```py -from dimos.core import autoconnect, LCMTransport +from dimos.core.blueprints import autoconnect +from dimos.core.transport import LCMTransport from dimos.msgs.sensor_msgs import Image from dimos.robot.unitree.go2.connection import go2_connection from dimos.agents.agent import agent diff --git a/dimos/agents/agent.py b/dimos/agents/agent.py index 98f23d7e8d..37e1a4757c 100644 --- a/dimos/agents/agent.py +++ b/dimos/agents/agent.py @@ -19,7 +19,6 @@ from typing import TYPE_CHECKING, Any, Protocol import uuid -from langchain.agents import create_agent from langchain_core.messages import HumanMessage from langchain_core.messages.base import BaseMessage from langchain_core.tools import StructuredTool @@ -104,6 +103,9 @@ def on_system_modules(self, modules: list[RPCClient]) -> None: model = MockModel(json_path=self.config.model_fixture) with self._lock: + # Here to prevent unwanted imports in the file. + from langchain.agents import create_agent + self._state_graph = create_agent( model=model, tools=_get_tools_from_modules(self, modules, self.rpc), diff --git a/dimos/agents/conftest.py b/dimos/agents/conftest.py index 1be2aadc0c..7f34679a00 100644 --- a/dimos/agents/conftest.py +++ b/dimos/agents/conftest.py @@ -42,7 +42,6 @@ def fn( *, blueprints, messages: list[BaseMessage], - dask: bool = False, system_prompt: str | None = None, fixture: str | None = None, ) -> list[BaseMessage]: @@ -79,10 +78,7 @@ def on_message(msg: BaseMessage) -> None: AgentTestRunner.blueprint(messages=messages), ) - global_config.update( - viewer_backend="none", - dask=dask, - ) + global_config.update(viewer_backend="none") nonlocal coordinator coordinator = blueprint.build() diff --git a/dimos/agents/demo_agent.py b/dimos/agents/demo_agent.py index b3250fba5b..bd69fc6cae 100644 --- a/dimos/agents/demo_agent.py +++ b/dimos/agents/demo_agent.py @@ -20,13 +20,18 @@ demo_agent = autoconnect(Agent.blueprint()) + +def _create_webcam() -> Webcam: + return Webcam( + camera_index=0, + fps=15, + camera_info=zed.CameraInfo.SingleWebcam, + ) + + demo_agent_camera = autoconnect( Agent.blueprint(), camera_module( - hardware=lambda: Webcam( - camera_index=0, - fps=15, - camera_info=zed.CameraInfo.SingleWebcam, - ), + hardware=_create_webcam, ), ) diff --git a/dimos/agents/mcp/conftest.py b/dimos/agents/mcp/conftest.py index 532ef16592..27c20f054b 100644 --- a/dimos/agents/mcp/conftest.py +++ b/dimos/agents/mcp/conftest.py @@ -43,7 +43,6 @@ def fn( *, blueprints, messages: list[BaseMessage], - dask: bool = False, system_prompt: str | None = None, fixture: str | None = None, ) -> list[BaseMessage]: @@ -78,10 +77,7 @@ def on_message(msg: BaseMessage) -> None: AgentTestRunner.blueprint(messages=messages), ) - global_config.update( - viewer_backend="none", - dask=dask, - ) + global_config.update(viewer_backend="none") nonlocal coordinator coordinator = blueprint.build() diff --git a/dimos/agents/mcp/mcp_server.py b/dimos/agents/mcp/mcp_server.py index 1f8ce92888..87c27302db 100644 --- a/dimos/agents/mcp/mcp_server.py +++ b/dimos/agents/mcp/mcp_server.py @@ -28,11 +28,12 @@ logger = setup_logger() -from dimos.core import Module, rpc # noqa: I001 -from dimos.core.rpc_client import RpcCall, RPCClient - from starlette.requests import Request # noqa: TC002 +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.rpc_client import RpcCall, RPCClient + if TYPE_CHECKING: import concurrent.futures diff --git a/dimos/agents/mcp/test_mcp_client.py b/dimos/agents/mcp/test_mcp_client.py index 946bdc4eb8..16427103e4 100644 --- a/dimos/agents/mcp/test_mcp_client.py +++ b/dimos/agents/mcp/test_mcp_client.py @@ -30,12 +30,10 @@ def add(self, x: int, y: int) -> str: @pytest.mark.slow -@pytest.mark.parametrize("dask", [False, True]) -def test_can_call_tool(dask, agent_setup): +def test_can_call_tool(agent_setup): history = agent_setup( blueprints=[Adder.blueprint()], messages=[HumanMessage("What is 33333 + 100? Use the tool.")], - dask=dask, ) assert "33433" in history[-1].content @@ -67,8 +65,7 @@ def register_user(self, name: str) -> str: @pytest.mark.slow -@pytest.mark.parametrize("dask", [False, True]) -def test_can_call_again_on_error(dask, agent_setup): +def test_can_call_again_on_error(agent_setup): history = agent_setup( blueprints=[UserRegistration.blueprint()], messages=[ @@ -76,7 +73,6 @@ def test_can_call_again_on_error(dask, agent_setup): "Register a user named 'Paul'. If there are errors, just try again until you succeed." ) ], - dask=dask, ) assert any(message.content == "User name registered successfully." for message in history) diff --git a/dimos/agents/skills/navigation.py b/dimos/agents/skills/navigation.py index 322a09c2bb..b02ff3a446 100644 --- a/dimos/agents/skills/navigation.py +++ b/dimos/agents/skills/navigation.py @@ -21,8 +21,7 @@ from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In -from dimos.models.qwen.video_query import BBox -from dimos.models.vl.qwen import QwenVlModel +from dimos.models.qwen.bbox import BBox from dimos.msgs.geometry_msgs import PoseStamped, Quaternion, Vector3 from dimos.msgs.geometry_msgs.Vector3 import make_vector3 from dimos.msgs.sensor_msgs import Image @@ -59,6 +58,10 @@ class NavigationSkillContainer(Module): def __init__(self) -> None: super().__init__() self._skill_started = False + + # Here to prevent unwanted imports in the file. + from dimos.models.vl.qwen import QwenVlModel + self._vl_model = QwenVlModel() @rpc diff --git a/dimos/agents/skills/person_follow.py b/dimos/agents/skills/person_follow.py index 641055e6f6..4bb42b2970 100644 --- a/dimos/agents/skills/person_follow.py +++ b/dimos/agents/skills/person_follow.py @@ -26,8 +26,7 @@ from dimos.core.global_config import GlobalConfig from dimos.core.module import Module from dimos.core.stream import In, Out -from dimos.models.qwen.video_query import BBox -from dimos.models.segmentation.edge_tam import EdgeTAMProcessor +from dimos.models.qwen.bbox import BBox from dimos.models.vl.qwen import QwenVlModel from dimos.msgs.geometry_msgs import Twist from dimos.msgs.sensor_msgs import CameraInfo, Image, PointCloud2 @@ -37,6 +36,7 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: + from dimos.models.segmentation.edge_tam import EdgeTAMProcessor from dimos.models.vl.base import VlModel logger = setup_logger() @@ -176,6 +176,9 @@ def _follow_person(self, query: str, initial_bbox: BBox) -> str: with self._lock: if self._tracker is None: + # Here to prevent unwanted imports in the file. + from dimos.models.segmentation.edge_tam import EdgeTAMProcessor + self._tracker = EdgeTAMProcessor() tracker = self._tracker latest_image = self._latest_image @@ -202,7 +205,7 @@ def _follow_person(self, query: str, initial_bbox: BBox) -> str: "the 'stop_following' tool." ) - def _follow_loop(self, tracker: EdgeTAMProcessor, query: str) -> None: + def _follow_loop(self, tracker: "EdgeTAMProcessor", query: str) -> None: lost_count = 0 period = 1.0 / self._frequency next_time = time.monotonic() diff --git a/dimos/agents/test_agent.py b/dimos/agents/test_agent.py index cd571a56ae..2464e622ca 100644 --- a/dimos/agents/test_agent.py +++ b/dimos/agents/test_agent.py @@ -30,12 +30,10 @@ def add(self, x: int, y: int) -> str: @pytest.mark.slow -@pytest.mark.parametrize("dask", [False, True]) -def test_can_call_tool(dask, agent_setup): +def test_can_call_tool(agent_setup): history = agent_setup( blueprints=[Adder.blueprint()], messages=[HumanMessage("What is 33333 + 100? Use the tool.")], - dask=dask, ) assert "33433" in history[-1].content @@ -69,8 +67,7 @@ def register_user(self, name: str) -> str: @pytest.mark.slow -@pytest.mark.parametrize("dask", [False, True]) -def test_can_call_again_on_error(dask, agent_setup): +def test_can_call_again_on_error(agent_setup): history = agent_setup( blueprints=[UserRegistration.blueprint()], messages=[ @@ -78,7 +75,6 @@ def test_can_call_again_on_error(dask, agent_setup): "Register a user named 'Paul'. If there are errors, just try again until you succeed." ) ], - dask=dask, ) assert any(message.content == "User name registered successfully." for message in history) diff --git a/dimos/agents/vlm_agent.py b/dimos/agents/vlm_agent.py index c99f8afa49..ec0aec1442 100644 --- a/dimos/agents/vlm_agent.py +++ b/dimos/agents/vlm_agent.py @@ -19,8 +19,8 @@ from langchain_core.messages import AIMessage, HumanMessage, SystemMessage from dimos.agents.system_prompt import SYSTEM_PROMPT -from dimos.core import Module, rpc -from dimos.core.module import ModuleConfig +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.msgs.sensor_msgs import Image from dimos.utils.logging_config import setup_logger diff --git a/dimos/agents/vlm_stream_tester.py b/dimos/agents/vlm_stream_tester.py index 79bb802a03..4126c6b3a0 100644 --- a/dimos/agents/vlm_stream_tester.py +++ b/dimos/agents/vlm_stream_tester.py @@ -17,7 +17,8 @@ from langchain_core.messages import AIMessage, HumanMessage -from dimos.core import Module, rpc +from dimos.core.core import rpc +from dimos.core.module import Module from dimos.core.stream import In, Out from dimos.msgs.sensor_msgs import Image from dimos.utils.logging_config import setup_logger diff --git a/dimos/agents/web_human_input.py b/dimos/agents/web_human_input.py index 09d5400cdc..22fdb231b3 100644 --- a/dimos/agents/web_human_input.py +++ b/dimos/agents/web_human_input.py @@ -18,10 +18,10 @@ import reactivex as rx import reactivex.operators as ops -from dimos.core import Module, rpc +from dimos.core.core import rpc +from dimos.core.module import Module from dimos.core.transport import pLCMTransport from dimos.stream.audio.node_normalizer import AudioNormalizer -from dimos.stream.audio.stt.node_whisper import WhisperNode from dimos.utils.logging_config import setup_logger from dimos.web.robot_web_interface import RobotWebInterface @@ -51,6 +51,10 @@ def start(self) -> None: ) normalizer = AudioNormalizer() + + # Here to prevent unwanted imports in the file. + from dimos.stream.audio.stt.node_whisper import WhisperNode + stt_node = WhisperNode() # Connect audio pipeline: browser audio → normalizer → whisper diff --git a/dimos/agents_deprecated/memory/spatial_vector_db.py b/dimos/agents_deprecated/memory/spatial_vector_db.py index c482076325..7d0c8eb2f7 100644 --- a/dimos/agents_deprecated/memory/spatial_vector_db.py +++ b/dimos/agents_deprecated/memory/spatial_vector_db.py @@ -21,7 +21,6 @@ from typing import Any -import chromadb import numpy as np from dimos.agents_deprecated.memory.visual_memory import VisualMemory @@ -57,6 +56,9 @@ def __init__( # type: ignore[no-untyped-def] """ self.collection_name = collection_name + # Here to prevent unwanted imports in the file. + import chromadb + # Use provided client or create in-memory client self.client = chroma_client if chroma_client is not None else chromadb.Client() diff --git a/dimos/agents_deprecated/modules/base_agent.py b/dimos/agents_deprecated/modules/base_agent.py index efe81fd90b..18ac15b317 100644 --- a/dimos/agents_deprecated/modules/base_agent.py +++ b/dimos/agents_deprecated/modules/base_agent.py @@ -20,7 +20,9 @@ from dimos.agents_deprecated.agent_message import AgentMessage from dimos.agents_deprecated.agent_types import AgentResponse from dimos.agents_deprecated.memory.base import AbstractAgentSemanticMemory -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.skills.skills import AbstractSkill, SkillLibrary from dimos.utils.logging_config import setup_logger diff --git a/dimos/conftest.py b/dimos/conftest.py index 93c0e91306..701b915dcf 100644 --- a/dimos/conftest.py +++ b/dimos/conftest.py @@ -19,6 +19,7 @@ from dotenv import load_dotenv import pytest +from dimos.core.module_coordinator import ModuleCoordinator from dimos.protocol.service.lcmservice import autoconf load_dotenv() @@ -86,9 +87,8 @@ def _autoconf(request): @pytest.fixture(scope="module") def dimos_cluster(): - from dimos.core import start - - dimos = start(4) + dimos = ModuleCoordinator() + dimos.start() try: yield dimos finally: diff --git a/dimos/control/coordinator.py b/dimos/control/coordinator.py index c9182e6aa8..6cd96d0000 100644 --- a/dimos/control/coordinator.py +++ b/dimos/control/coordinator.py @@ -25,9 +25,8 @@ - Aggregated preemption notifications """ -from __future__ import annotations - from dataclasses import dataclass, field +from pathlib import Path import threading import time from typing import TYPE_CHECKING, Any @@ -43,28 +42,28 @@ from dimos.control.hardware_interface import ConnectedHardware, ConnectedTwistBase from dimos.control.task import ControlTask from dimos.control.tick_loop import TickLoop -from dimos.core import In, Module, Out, rpc -from dimos.core.module import ModuleConfig +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out from dimos.hardware.drive_trains.spec import ( TwistBaseAdapter, ) +from dimos.hardware.manipulators.spec import ManipulatorAdapter from dimos.msgs.geometry_msgs import ( - PoseStamped, # noqa: TC001 - needed at runtime for In[PoseStamped] - Twist, # noqa: TC001 - needed at runtime for In[Twist] + PoseStamped, + Twist, ) from dimos.msgs.sensor_msgs import ( JointState, ) from dimos.teleop.quest.quest_types import ( - Buttons, # noqa: TC001 - needed at runtime for In[Buttons] + Buttons, ) from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from collections.abc import Callable - from pathlib import Path - from dimos.hardware.manipulators.spec import ManipulatorAdapter logger = setup_logger() diff --git a/dimos/core/__init__.py b/dimos/core/__init__.py index 2b6296b623..e69de29bb2 100644 --- a/dimos/core/__init__.py +++ b/dimos/core/__init__.py @@ -1,278 +0,0 @@ -from __future__ import annotations - -import multiprocessing as mp -import time -from typing import TYPE_CHECKING, cast - -import lazy_loader as lazy -from rich.console import Console - -from dimos.core.core import rpc -from dimos.utils.logging_config import setup_logger - -if TYPE_CHECKING: - # Avoid runtime import to prevent circular import; ruff's TC001 would otherwise move it. - from dask.distributed import LocalCluster - - from dimos.core._dask_exports import DimosCluster - from dimos.core.module import Module - from dimos.core.rpc_client import ModuleProxy - -logger = setup_logger() - -__getattr__, __dir__, __all__ = lazy.attach( - __name__, - submodules=["colors"], - submod_attrs={ - "blueprints": ["autoconnect", "Blueprint"], - "_dask_exports": ["DimosCluster"], - "_protocol_exports": ["LCMRPC", "RPCSpec", "LCMTF", "TF", "PubSubTF", "TFConfig", "TFSpec"], - "module": ["Module", "ModuleBase", "ModuleConfig", "ModuleConfigT"], - "stream": ["In", "Out", "RemoteIn", "RemoteOut", "Transport"], - "transport": [ - "LCMTransport", - "SHMTransport", - "ZenohTransport", - "pLCMTransport", - "pSHMTransport", - ], - }, -) -__all__ += ["DimosCluster", "Module", "rpc", "start", "wait_exit"] - - -class CudaCleanupPlugin: - """Dask worker plugin to cleanup CUDA resources on shutdown.""" - - def setup(self, worker) -> None: # type: ignore[no-untyped-def] - """Called when worker starts.""" - pass - - def teardown(self, worker) -> None: # type: ignore[no-untyped-def] - """Clean up CUDA resources when worker shuts down.""" - try: - import sys - - if "cupy" in sys.modules: - import cupy as cp # type: ignore[import-not-found, import-untyped] - - # Clear memory pools - mempool = cp.get_default_memory_pool() - pinned_mempool = cp.get_default_pinned_memory_pool() - mempool.free_all_blocks() - pinned_mempool.free_all_blocks() - cp.cuda.Stream.null.synchronize() - mempool.free_all_blocks() - pinned_mempool.free_all_blocks() - except Exception: - pass - - -def patch_actor(actor, cls) -> None: ... # type: ignore[no-untyped-def] - - -def patchdask(dask_client: DimosCluster, local_cluster: LocalCluster) -> DimosCluster: - from dimos.core.rpc_client import RPCClient - from dimos.utils.actor_registry import ActorRegistry - - def deploy( # type: ignore[no-untyped-def] - actor_class: type[Module], - *args, - **kwargs, - ) -> ModuleProxy: - from dimos.core.docker_runner import DockerModule, is_docker_module - - # Check if this module should run in Docker (based on its default_config) - if is_docker_module(actor_class): - logger.info("Deploying module in Docker.", module=actor_class.__name__) - dm = DockerModule(actor_class, *args, **kwargs) - dm.start() # Explicit start - follows create -> configure -> start lifecycle - dask_client._docker_modules.append(dm) # type: ignore[attr-defined] - return dm # type: ignore[return-value] - - logger.info("Deploying module.", module=actor_class.__name__) - actor = dask_client.submit( # type: ignore[no-untyped-call] - actor_class, - *args, - **kwargs, - actor=True, - ).result() - - worker = actor.set_ref(actor).result() - logger.info("Deployed module.", module=actor._cls.__name__, worker_id=worker) - - # Register actor deployment in shared memory - ActorRegistry.update(str(actor), str(worker)) - - return cast("ModuleProxy", RPCClient(actor, actor_class)) - - def check_worker_memory() -> None: - """Check memory usage of all workers.""" - info = dask_client.scheduler_info() - - console = Console() - total_workers = len(info.get("workers", {})) - total_memory_used = 0 - total_memory_limit = 0 - - for worker_addr, worker_info in info.get("workers", {}).items(): - metrics = worker_info.get("metrics", {}) - memory_used = metrics.get("memory", 0) - memory_limit = worker_info.get("memory_limit", 0) - - cpu_percent = metrics.get("cpu", 0) - managed_bytes = metrics.get("managed_bytes", 0) - spilled = metrics.get("spilled_bytes", {}).get("memory", 0) - worker_status = worker_info.get("status", "unknown") - worker_id = worker_info.get("id", "?") - - memory_used_gb = memory_used / 1e9 - memory_limit_gb = memory_limit / 1e9 - managed_gb = managed_bytes / 1e9 - spilled / 1e9 - - total_memory_used += memory_used - total_memory_limit += memory_limit - - percentage = (memory_used_gb / memory_limit_gb * 100) if memory_limit_gb > 0 else 0 - - if worker_status == "paused": - status = "[red]PAUSED" - elif percentage >= 95: - status = "[red]CRITICAL" - elif percentage >= 80: - status = "[yellow]WARNING" - else: - status = "[green]OK" - - console.print( - f"Worker-{worker_id} {worker_addr}: " - f"{memory_used_gb:.2f}/{memory_limit_gb:.2f}GB ({percentage:.1f}%) " - f"CPU:{cpu_percent:.0f}% Managed:{managed_gb:.2f}GB " - f"{status}" - ) - - if total_workers > 0: - total_used_gb = total_memory_used / 1e9 - total_limit_gb = total_memory_limit / 1e9 - total_percentage = (total_used_gb / total_limit_gb * 100) if total_limit_gb > 0 else 0 - console.print( - f"[bold]Total: {total_used_gb:.2f}/{total_limit_gb:.2f}GB ({total_percentage:.1f}%) across {total_workers} workers[/bold]" - ) - - def close_all() -> None: - # Prevents multiple calls to close_all - if hasattr(dask_client, "_closed") and dask_client._closed: - return - dask_client._closed = True # type: ignore[attr-defined] - - # Stop all Docker modules (in reverse order of deployment) - for dm in reversed(dask_client._docker_modules): # type: ignore[attr-defined] - try: - dm.stop() - except Exception: - pass - dask_client._docker_modules.clear() # type: ignore[attr-defined] - - # Stop all SharedMemory transports before closing Dask - # This prevents the "leaked shared_memory objects" warning and hangs - try: - import gc - - from dimos.protocol.pubsub.impl import shmpubsub - - for obj in gc.get_objects(): - if isinstance(obj, shmpubsub.SharedMemoryPubSubBase): - try: - obj.stop() - except Exception: - pass - except Exception: - pass - - # Get the event loop before shutting down - loop = dask_client.loop - - # Clear the actor registry - ActorRegistry.clear() - - # Close cluster and client with reasonable timeout - # The CudaCleanupPlugin will handle CUDA cleanup on each worker - try: - local_cluster.close(timeout=5) - except Exception: - pass - - try: - dask_client.close(timeout=5) # type: ignore[no-untyped-call] - except Exception: - pass - - if loop and hasattr(loop, "add_callback") and hasattr(loop, "stop"): - try: - loop.add_callback(loop.stop) - except Exception: - pass - - # Note: We do NOT shutdown the _offload_executor here because it's a global - # module-level ThreadPoolExecutor shared across all Dask clients in the process. - # Shutting it down here would break subsequent Dask client usage (e.g., in tests). - # The executor will be cleaned up when the Python process exits. - - # Give threads time to clean up - # Dask's IO loop and Profile threads are daemon threads - # that will be cleaned up when the process exits - # This is needed, solves race condition in CI thread check - time.sleep(0.1) - - dask_client._docker_modules = [] # type: ignore[attr-defined] - dask_client.deploy = deploy # type: ignore[attr-defined] - dask_client.check_worker_memory = check_worker_memory # type: ignore[attr-defined] - dask_client.stop = lambda: dask_client.close() # type: ignore[attr-defined, no-untyped-call] - dask_client.close_all = close_all # type: ignore[attr-defined] - return dask_client # type: ignore[return-value] - - -def start(n: int | None = None, memory_limit: str = "auto") -> DimosCluster: - """Start a Dask LocalCluster with specified workers and memory limits. - - Args: - n: Number of workers (defaults to CPU count) - memory_limit: Memory limit per worker (e.g., '4GB', '2GiB', or 'auto' for Dask's default) - - Returns: - DimosCluster: A patched Dask client with deploy(), check_worker_memory(), stop(), and close_all() methods - """ - - from dask.distributed import Client, LocalCluster - - console = Console() - if not n: - n = mp.cpu_count() - with console.status( - f"[green]Initializing dimos local cluster with [bright_blue]{n} workers", spinner="arc" - ): - cluster = LocalCluster( # type: ignore[no-untyped-call] - n_workers=n, - threads_per_worker=4, - memory_limit=memory_limit, - plugins=[CudaCleanupPlugin()], # Register CUDA cleanup plugin - ) - client = Client(cluster) # type: ignore[no-untyped-call] - - console.print( - f"[green]Initialized dimos local cluster with [bright_blue]{n} workers, memory limit: {memory_limit}" - ) - - patched_client = patchdask(client, cluster) - - return patched_client - - -def wait_exit() -> None: - while True: - try: - time.sleep(1) - except KeyboardInterrupt: - print("exiting...") - return diff --git a/dimos/core/_protocol_exports.py b/dimos/core/_protocol_exports.py deleted file mode 100644 index be77fd8323..0000000000 --- a/dimos/core/_protocol_exports.py +++ /dev/null @@ -1,19 +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 dimos.protocol.rpc import LCMRPC -from dimos.protocol.rpc.spec import RPCSpec -from dimos.protocol.tf import LCMTF, TF, PubSubTF, TFConfig, TFSpec - -__all__ = ["LCMRPC", "LCMTF", "TF", "PubSubTF", "RPCSpec", "TFConfig", "TFSpec"] diff --git a/dimos/core/blueprints.py b/dimos/core/blueprints.py index d2179d16a5..4c45827850 100644 --- a/dimos/core/blueprints.py +++ b/dimos/core/blueprints.py @@ -490,6 +490,7 @@ def build( self, cli_config_overrides: Mapping[str, Any] | None = None, ) -> ModuleCoordinator: + logger.info("Building the blueprint") global_config.update(**dict(self.global_config_overrides)) if cli_config_overrides: global_config.update(**dict(cli_config_overrides)) @@ -498,6 +499,7 @@ def build( self._check_requirements() self._verify_no_name_conflicts() + logger.info("Starting the modules") module_coordinator = ModuleCoordinator(cfg=global_config) module_coordinator.start() diff --git a/dimos/core/_dask_exports.py b/dimos/core/conftest.py similarity index 72% rename from dimos/core/_dask_exports.py rename to dimos/core/conftest.py index cb257e7804..71f4d3bad5 100644 --- a/dimos/core/_dask_exports.py +++ b/dimos/core/conftest.py @@ -12,6 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dask.distributed import Client as DimosCluster +import pytest -__all__ = ["DimosCluster"] +from dimos.core.module_coordinator import ModuleCoordinator + + +@pytest.fixture +def dimos(): + client = ModuleCoordinator() + client.start() + try: + yield client + finally: + client.stop() diff --git a/dimos/core/docker_runner.py b/dimos/core/docker_runner.py index 9be2ff6012..ee56163ca6 100644 --- a/dimos/core/docker_runner.py +++ b/dimos/core/docker_runner.py @@ -269,7 +269,7 @@ def tail_logs(self, n: int = 200) -> str: return _tail_logs(self._config, self._container_name, n=n) def set_transport(self, stream_name: str, transport: Any) -> bool: - """Configure stream transport in container. Mirrors DaskModule.set_transport() for autoconnect().""" + """Configure stream transport in container. Mirrors Module.set_transport() for autoconnect().""" topic = getattr(transport, "topic", None) if topic is None: return False diff --git a/dimos/core/global_config.py b/dimos/core/global_config.py index 080c2c8bbc..f629122cee 100644 --- a/dimos/core/global_config.py +++ b/dimos/core/global_config.py @@ -31,7 +31,7 @@ class GlobalConfig(BaseSettings): simulation: bool = False replay: bool = False viewer_backend: ViewerBackend = "rerun-web" - n_dask_workers: int = 2 + n_workers: int = 2 memory_limit: str = "auto" mujoco_camera_position: str | None = None mujoco_room: str | None = None @@ -45,7 +45,6 @@ class GlobalConfig(BaseSettings): robot_rotation_diameter: float = 0.6 planner_strategy: NavigationStrategy = "simple" planner_robot_speed: float | None = None - dask: bool = True model_config = SettingsConfigDict( env_file=".env", diff --git a/dimos/core/introspection/blueprint/dot.py b/dimos/core/introspection/blueprint/dot.py index c60ad06fc8..ea66401033 100644 --- a/dimos/core/introspection/blueprint/dot.py +++ b/dimos/core/introspection/blueprint/dot.py @@ -48,7 +48,6 @@ class LayoutAlgo(Enum): DEFAULT_IGNORED_MODULES = { "WebsocketVisModule", - "UtilizationModule", # "FoxgloveBridge", } diff --git a/dimos/core/introspection/module/ansi.py b/dimos/core/introspection/module/ansi.py index 6e835d63d3..f6be616f47 100644 --- a/dimos/core/introspection/module/ansi.py +++ b/dimos/core/introspection/module/ansi.py @@ -14,7 +14,6 @@ """ANSI terminal renderer for module IO diagrams.""" -from dimos.core import colors from dimos.core.introspection.module.info import ( ModuleInfo, ParamInfo, @@ -22,6 +21,7 @@ SkillInfo, StreamInfo, ) +from dimos.utils import colors def render(info: ModuleInfo, color: bool = True) -> str: diff --git a/dimos/core/module.py b/dimos/core/module.py index d6089a8f0a..8f2dd916c9 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -11,9 +11,8 @@ # 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 asyncio +from collections.abc import Callable from dataclasses import dataclass from functools import partial import inspect @@ -32,26 +31,23 @@ from typing_extensions import TypeVar as TypeVarExtension if TYPE_CHECKING: - from collections.abc import Callable - from dimos.core.introspection.module import ModuleInfo from dimos.core.rpc_client import RPCClient from typing import TypeVar -from dask.distributed import Actor, get_worker from langchain_core.tools import tool from reactivex.disposable import CompositeDisposable -from dimos.core import colors from dimos.core.core import T, rpc from dimos.core.introspection.module import extract_module_info, render_module_io from dimos.core.resource import Resource -from dimos.core.rpc_client import RpcCall # noqa: TC001 -from dimos.core.stream import In, Out, RemoteIn, RemoteOut, Transport +from dimos.core.rpc_client import RpcCall +from dimos.core.stream import In, Out, RemoteOut, Transport from dimos.protocol.rpc import LCMRPC, RPCSpec from dimos.protocol.service import Configurable # type: ignore[attr-defined] from dimos.protocol.tf import LCMTF, TFSpec +from dimos.utils import colors from dimos.utils.generic import classproperty @@ -63,21 +59,6 @@ class SkillInfo: def get_loop() -> tuple[asyncio.AbstractEventLoop, threading.Thread | None]: - # we are actually instantiating a new loop here - # to not interfere with an existing dask loop - - # try: - # # here we attempt to figure out if we are running on a dask worker - # # if so we use the dask worker _loop as ours, - # # and we register our RPC server - # worker = get_worker() - # if worker.loop: - # print("using dask worker loop") - # return worker.loop.asyncio_loop - - # except ValueError: - # ... - try: running_loop = asyncio.get_running_loop() return running_loop, None @@ -108,6 +89,8 @@ class ModuleBase(Configurable[ModuleConfigT], Resource): _loop_thread: threading.Thread | None _disposables: CompositeDisposable _bound_rpc_calls: dict[str, RpcCall] = {} + _module_closed: bool = False + _module_closed_lock: threading.Lock rpc_calls: list[str] = [] @@ -115,13 +98,10 @@ class ModuleBase(Configurable[ModuleConfigT], Resource): def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] super().__init__(*args, **kwargs) + self._module_closed_lock = threading.Lock() self._loop, self._loop_thread = get_loop() self._disposables = CompositeDisposable() - # we can completely override comms protocols if we want try: - # here we attempt to figure out if we are running on a dask worker - # if so we use the dask worker _loop as ours, - # and we register our RPC server self.rpc = self.config.rpc_transport() self.rpc.serve_module_rpc(self) self.rpc.start() # type: ignore[attr-defined] @@ -144,6 +124,11 @@ def stop(self) -> None: self._close_module() def _close_module(self) -> None: + with self._module_closed_lock: + if self._module_closed: + return + self._module_closed = True + self._close_rpc() # Save into local variables to avoid race when stopping concurrently @@ -175,6 +160,7 @@ def __getstate__(self): # type: ignore[no-untyped-def] 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) @@ -186,6 +172,7 @@ def __setstate__(self, state) -> None: # type: ignore[no-untyped-def] self.__dict__.update(state) # Reinitialize runtime attributes self._disposables = CompositeDisposable() + self._module_closed_lock = threading.Lock() self._loop = None self._loop_thread = None self._rpc = None @@ -286,7 +273,7 @@ class _io_descriptor: """Descriptor that makes io() work on both class and instance.""" def __get__( - self, obj: ModuleBase | None, objtype: type[ModuleBase] + self, obj: "ModuleBase | None", objtype: "type[ModuleBase]" ) -> Callable[[bool], str]: if obj is None: return objtype._io_class @@ -295,7 +282,7 @@ def __get__( io = _io_descriptor() @classmethod - def _module_info_class(cls) -> ModuleInfo: + def _module_info_class(cls) -> "ModuleInfo": """Class-level module_info() - returns ModuleInfo from annotations.""" hints = get_type_hints(cls) @@ -331,8 +318,8 @@ class _module_info_descriptor: """Descriptor that makes module_info() work on both class and instance.""" def __get__( - self, obj: ModuleBase | None, objtype: type[ModuleBase] - ) -> Callable[[], ModuleInfo]: + self, obj: "ModuleBase | None", objtype: "type[ModuleBase]" + ) -> "Callable[[], ModuleInfo]": if obj is None: return objtype._module_info_class # For instances, extract from actual streams @@ -362,7 +349,7 @@ def set_rpc_method(self, method: str, callable: RpcCall) -> None: self._bound_rpc_calls[method] = callable @rpc - def set_module_ref(self, name: str, module_ref: RPCClient) -> None: + def set_module_ref(self, name: str, module_ref: "RPCClient") -> None: setattr(self, name, module_ref) @overload @@ -396,9 +383,6 @@ def get_skills(self) -> list[SkillInfo]: class Module(ModuleBase[ModuleConfigT]): - ref: Actor - worker: int - def __init_subclass__(cls, **kwargs: Any) -> None: """Set class-level None attributes for In/Out type annotations. @@ -431,8 +415,6 @@ def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] # Get type hints with proper namespace resolution for subclasses # Collect namespaces from all classes in the MRO chain - import sys - globalns = {} for cls in self.__class__.__mro__: if cls.__module__ in sys.modules: @@ -457,12 +439,6 @@ def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] setattr(self, name, stream) super().__init__(*args, **kwargs) - def set_ref(self, ref) -> int: # type: ignore[no-untyped-def] - worker = get_worker() - self.ref = ref - self.worker = worker.name - return worker.name # type: ignore[no-any-return] - def __str__(self) -> str: return f"{self.__class__.__name__}" @@ -498,14 +474,8 @@ def connect_stream(self, input_name: str, remote_stream: RemoteOut[T]): # type: raise TypeError(f"Input {input_name} is not a valid stream") input_stream.connection = remote_stream - def dask_receive_msg(self, input_name: str, msg: Any) -> None: - getattr(self, input_name).transport.dask_receive_msg(msg) - - def dask_register_subscriber(self, output_name: str, subscriber: RemoteIn[T]) -> None: - getattr(self, output_name).transport.dask_register_subscriber(subscriber) - -ModuleT = TypeVar("ModuleT", bound="Module") +ModuleT = TypeVar("ModuleT", bound="Module[Any]") def is_module_type(value: Any) -> bool: diff --git a/dimos/core/module_coordinator.py b/dimos/core/module_coordinator.py index c6d975731d..d816a60cb4 100644 --- a/dimos/core/module_coordinator.py +++ b/dimos/core/module_coordinator.py @@ -16,19 +16,20 @@ import time from typing import TYPE_CHECKING, Any -from dimos import core -from dimos.core import DimosCluster from dimos.core.global_config import GlobalConfig, global_config from dimos.core.module import Module, ModuleT from dimos.core.resource import Resource from dimos.core.worker_manager import WorkerManager +from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from dimos.core.rpc_client import ModuleProxy +logger = setup_logger() + class ModuleCoordinator(Resource): # type: ignore[misc] - _client: DimosCluster | WorkerManager | None = None + _client: WorkerManager | None = None _global_config: GlobalConfig _n: int | None = None _memory_limit: str = "auto" @@ -39,26 +40,30 @@ def __init__( n: int | None = None, cfg: GlobalConfig = global_config, ) -> None: - self._n = n if n is not None else cfg.n_dask_workers + self._n = n if n is not None else cfg.n_workers self._memory_limit = cfg.memory_limit self._global_config = cfg self._deployed_modules = {} def start(self) -> None: - if self._global_config.dask: - self._client = core.start(self._n, self._memory_limit) - else: - self._client = WorkerManager() + n = self._n if self._n is not None else 2 + self._client = WorkerManager(n_workers=n) + self._client.start() def stop(self) -> None: - for module in reversed(self._deployed_modules.values()): - module.stop() + for module_class, module in reversed(self._deployed_modules.items()): + logger.info("Stopping module...", module=module_class.__name__) + try: + module.stop() + except Exception: + logger.error("Error stopping module", module=module_class.__name__, exc_info=True) + logger.info("Module stopped.", module=module_class.__name__) self._client.close_all() # type: ignore[union-attr] def deploy(self, module_class: type[ModuleT], *args, **kwargs) -> "ModuleProxy": # type: ignore[no-untyped-def] if not self._client: - raise ValueError("Trying to dimos.deploy before dask client has started") + raise ValueError("Trying to dimos.deploy before the client has started") module: ModuleProxy = self._client.deploy(module_class, *args, **kwargs) # type: ignore[union-attr, attr-defined, assignment] self._deployed_modules[module_class] = module @@ -70,16 +75,10 @@ def deploy_parallel( if not self._client: raise ValueError("Not started") - if isinstance(self._client, WorkerManager): - modules = self._client.deploy_parallel(module_specs) - for (module_class, _, _), module in zip(module_specs, modules, strict=True): - self._deployed_modules[module_class] = module # type: ignore[assignment] - return modules # type: ignore[return-value] - else: - return [ - self.deploy(module_class, *args, **kwargs) - for module_class, args, kwargs in module_specs - ] + modules = self._client.deploy_parallel(module_specs) + for (module_class, _, _), module in zip(module_specs, modules, strict=True): + self._deployed_modules[module_class] = module # type: ignore[assignment] + return modules # type: ignore[return-value] def start_all_modules(self) -> None: modules = list(self._deployed_modules.values()) diff --git a/dimos/core/rpc_client.py b/dimos/core/rpc_client.py index 30cc4f3017..e46124469c 100644 --- a/dimos/core/rpc_client.py +++ b/dimos/core/rpc_client.py @@ -15,6 +15,8 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any +from dimos.core.stream import RemoteStream +from dimos.core.worker import MethodCallProxy from dimos.protocol.rpc import LCMRPC, RPCSpec from dimos.utils.logging_config import setup_logger @@ -136,7 +138,15 @@ def __getattr__(self, name: str): # type: ignore[no-untyped-def] # return super().__getattr__(name) # Try to avoid recursion by directly accessing attributes that are known - return self.actor_instance.__getattr__(name) + 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. + if isinstance(result, RemoteStream): + result.owner = MethodCallProxy(self.actor_instance) + + return result if TYPE_CHECKING: diff --git a/dimos/core/stream.py b/dimos/core/stream.py index 77edf45417..7791968a29 100644 --- a/dimos/core/stream.py +++ b/dimos/core/stream.py @@ -22,13 +22,12 @@ TypeVar, ) -from dask.distributed import Actor import reactivex as rx from reactivex import operators as ops from reactivex.disposable import Disposable -import dimos.core.colors as colors from dimos.core.resource import Resource +from dimos.utils import colors from dimos.utils.logging_config import setup_logger import dimos.utils.reactive as reactive from dimos.utils.reactive import backpressure @@ -132,11 +131,7 @@ def __str__(self) -> str: + " " + self._color_fn()(f"{self.name}[{self.type_name}]") + " @ " - + ( - colors.orange(self.owner) # type: ignore[arg-type] - if isinstance(self.owner, Actor) - else colors.green(self.owner) # type: ignore[arg-type] - ) + + colors.green(self.owner) # type: ignore[arg-type] + ("" if not self._transport else " via " + str(self._transport)) ) diff --git a/dimos/core/test_core.py b/dimos/core/test_core.py index 3866d55bdb..197539ef67 100644 --- a/dimos/core/test_core.py +++ b/dimos/core/test_core.py @@ -17,21 +17,15 @@ import pytest from reactivex.disposable import Disposable -from dimos.core import ( - In, - LCMTransport, - Module, - Out, - pLCMTransport, - rpc, -) -from dimos.core.testing import MockRobotClient, dimos +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out +from dimos.core.testing import MockRobotClient +from dimos.core.transport import LCMTransport, pLCMTransport from dimos.msgs.geometry_msgs import Vector3 from dimos.msgs.sensor_msgs import PointCloud2 from dimos.robot.unitree.type.odometry import Odometry -assert dimos - class Navigation(Module): mov: Out[Vector3] @@ -127,17 +121,7 @@ def test_basic_deployment(dimos) -> None: nav.start() time.sleep(1) - robot.stop() - - print("robot.mov_msg_count", robot.mov_msg_count) - print("nav.odom_msg_count", nav.odom_msg_count) - print("nav.lidar_msg_count", nav.lidar_msg_count) assert robot.mov_msg_count >= 8 assert nav.odom_msg_count >= 8 assert nav.lidar_msg_count >= 8 - - nav.stop() - nav.stop_rpc_client() - robot.stop_rpc_client() - dimos.close_all() diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index a022be0685..1c37c3a9d6 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -25,10 +25,10 @@ import pytest -from dimos.core import DimosCluster from dimos.core.blueprints import autoconnect from dimos.core.core import rpc from dimos.core.module import Module +from dimos.core.module_coordinator import ModuleCoordinator from dimos.core.native_module import LogFormat, NativeModule, NativeModuleConfig from dimos.core.stream import In, Out from dimos.core.transport import LCMTransport @@ -110,15 +110,16 @@ def test_process_crash_triggers_stop() -> None: assert mod._process is None, f"Watchdog did not clean up after process {pid} died" -def test_manual(dimos_cluster: DimosCluster, args_file: str) -> None: +@pytest.mark.slow +def test_manual(dimos_cluster: ModuleCoordinator, args_file: str) -> None: native_module = dimos_cluster.deploy( # type: ignore[attr-defined] StubNativeModule, some_param=2.5, output_file=args_file, ) - native_module.pointcloud.transport = LCMTransport("/my/custom/lidar", PointCloud2) - native_module.cmd_vel.transport = LCMTransport("/cmd_vel", Twist) + native_module.set_transport("pointcloud", LCMTransport("/my/custom/lidar", PointCloud2)) + native_module.set_transport("cmd_vel", LCMTransport("/cmd_vel", Twist)) native_module.start() time.sleep(1) native_module.stop() diff --git a/dimos/core/test_rpcstress.py b/dimos/core/test_rpcstress.py index 1d09f3e210..acfdb1b86d 100644 --- a/dimos/core/test_rpcstress.py +++ b/dimos/core/test_rpcstress.py @@ -15,7 +15,11 @@ import threading import time -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.module_coordinator import ModuleCoordinator +from dimos.core.stream import In, Out +from dimos.core.transport import pLCMTransport class Counter(Module): @@ -123,11 +127,8 @@ def get_stats(self): if __name__ == "__main__": - import dimos.core as core - from dimos.core import pLCMTransport - - # Start dimos with 2 workers - client = core.start(2) + client = ModuleCoordinator() + client.start() # Deploy counter module counter = client.deploy(Counter) @@ -168,10 +169,6 @@ def get_stats(self): if stats["missing_numbers"]: print(f" - First missing numbers: {stats['missing_numbers']}") - # Stop modules validator.stop() - # Shutdown dimos - client.shutdown() - - print("[MAIN] Test complete.") + client.stop() diff --git a/dimos/core/test_stream.py b/dimos/core/test_stream.py index 7a594f64e4..a7c949b33a 100644 --- a/dimos/core/test_stream.py +++ b/dimos/core/test_stream.py @@ -18,19 +18,14 @@ import pytest -from dimos.core import ( - In, - LCMTransport, - Module, - pLCMTransport, - rpc, -) -from dimos.core.testing import MockRobotClient, dimos +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In +from dimos.core.testing import MockRobotClient +from dimos.core.transport import LCMTransport, pLCMTransport from dimos.msgs.sensor_msgs import PointCloud2 from dimos.robot.unitree.type.odometry import Odometry -assert dimos - class SubscriberBase(Module): sub1_msgs: list[Odometry] = None @@ -212,7 +207,6 @@ def test_subscription(dimos, subscriber_class) -> None: robot.stop() subscriber.stop_rpc_client() robot.stop_rpc_client() - dimos.close_all() @pytest.mark.slow @@ -245,7 +239,6 @@ def test_get_next(dimos) -> None: robot.stop() subscriber.stop_rpc_client() robot.stop_rpc_client() - dimos.close_all() @pytest.mark.slow @@ -285,4 +278,3 @@ def test_hot_getter(dimos) -> None: robot.stop() subscriber.stop_rpc_client() robot.stop_rpc_client() - dimos.close_all() diff --git a/dimos/core/test_worker.py b/dimos/core/test_worker.py index 6892d226fd..8c75f41222 100644 --- a/dimos/core/test_worker.py +++ b/dimos/core/test_worker.py @@ -14,7 +14,9 @@ import pytest -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.core.worker_manager import WorkerManager from dimos.msgs.geometry_msgs import Vector3 @@ -74,16 +76,24 @@ def get_multiplier(self) -> int: @pytest.fixture -def worker_manager(): - manager = WorkerManager() - try: - yield manager - finally: +def create_worker_manager(): + manager = None + + def _create(n_workers): + nonlocal manager + manager = WorkerManager(n_workers=n_workers) + manager.start() + return manager + + yield _create + + if manager is not None: manager.close_all() @pytest.mark.slow -def test_worker_manager_basic(worker_manager): +def test_worker_manager_basic(create_worker_manager): + worker_manager = create_worker_manager(n_workers=2) module = worker_manager.deploy(SimpleModule) module.start() @@ -100,7 +110,8 @@ def test_worker_manager_basic(worker_manager): @pytest.mark.slow -def test_worker_manager_multiple_different_modules(worker_manager): +def test_worker_manager_multiple_different_modules(create_worker_manager): + worker_manager = create_worker_manager(n_workers=2) module1 = worker_manager.deploy(SimpleModule) module2 = worker_manager.deploy(AnotherModule) @@ -121,7 +132,8 @@ def test_worker_manager_multiple_different_modules(worker_manager): @pytest.mark.slow -def test_worker_manager_parallel_deployment(worker_manager): +def test_worker_manager_parallel_deployment(create_worker_manager): + worker_manager = create_worker_manager(n_workers=2) modules = worker_manager.deploy_parallel( [ (SimpleModule, (), {}), @@ -151,3 +163,28 @@ def test_worker_manager_parallel_deployment(worker_manager): module1.stop() module2.stop() module3.stop() + + +@pytest.mark.slow +def test_worker_pool_modules_share_workers(create_worker_manager): + manager = create_worker_manager(n_workers=1) + module1 = manager.deploy(SimpleModule) + module2 = manager.deploy(AnotherModule) + + module1.start() + module2.start() + + # Verify isolated state + module1.increment() + module1.increment() + module2.add(10) + + assert module1.get_counter() == 2 + assert module2.get_value() == 110 + + # Verify only 1 worker process was used + assert len(manager._workers) == 1 + assert manager._workers[0].module_count == 2 + + module1.stop() + module2.stop() diff --git a/dimos/core/testing.py b/dimos/core/testing.py index dee25aaa45..6431c09dbd 100644 --- a/dimos/core/testing.py +++ b/dimos/core/testing.py @@ -15,9 +15,9 @@ from threading import Event, Thread import time -import pytest # type: ignore[import-not-found] - -from dimos.core import In, Module, Out, rpc, start +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import Vector3 from dimos.msgs.sensor_msgs import PointCloud2 from dimos.robot.unitree.type.lidar import pointcloud2_from_webrtc_lidar @@ -25,14 +25,6 @@ from dimos.utils.testing import SensorReplay -@pytest.fixture -def dimos(): # type: ignore[no-untyped-def] - """Fixture to create a Dimos client for testing.""" - client = start(2) - yield client - client.stop() # type: ignore[attr-defined] - - class MockRobotClient(Module): odometry: Out[Odometry] lidar: Out[PointCloud2] @@ -40,14 +32,14 @@ class MockRobotClient(Module): mov_msg_count = 0 - def mov_callback(self, msg) -> None: # type: ignore[no-untyped-def] - self.mov_msg_count += 1 - def __init__(self) -> None: super().__init__() self._stop_event = Event() self._thread = None + def mov_callback(self, msg) -> None: # type: ignore[no-untyped-def] + self.mov_msg_count += 1 + @rpc def start(self) -> None: super().start() diff --git a/dimos/core/transport.py b/dimos/core/transport.py index 2586706feb..2300122c52 100644 --- a/dimos/core/transport.py +++ b/dimos/core/transport.py @@ -21,9 +21,9 @@ TypeVar, ) -import dimos.core.colors as colors from dimos.core.stream import In, Out, Stream, Transport from dimos.msgs.protocol import DimosMsg +from dimos.utils import colors try: import cyclonedds as _cyclonedds # noqa: F401 @@ -31,6 +31,7 @@ DDS_AVAILABLE = True except ImportError: DDS_AVAILABLE = False + from dimos.protocol.pubsub.impl.jpeg_shm import JpegSharedMemory from dimos.protocol.pubsub.impl.lcmpubsub import LCM, JpegLCM, PickleLCM, Topic as LCMTopic from dimos.protocol.pubsub.impl.rospubsub import DimosROS, ROSTopic diff --git a/dimos/core/worker.py b/dimos/core/worker.py index d6ff71918c..bb2ede4e03 100644 --- a/dimos/core/worker.py +++ b/dimos/core/worker.py @@ -12,17 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import multiprocessing as mp -from multiprocessing.connection import Connection +import threading import traceback -from typing import Any +from typing import TYPE_CHECKING, Any -from dimos.core.module import ModuleT -from dimos.core.rpc_client import RPCClient -from dimos.utils.actor_registry import ActorRegistry from dimos.utils.logging_config import setup_logger from dimos.utils.sequential_ids import SequentialIds +if TYPE_CHECKING: + from multiprocessing.connection import Connection + + from dimos.core.module import ModuleT + logger = setup_logger() @@ -36,25 +40,65 @@ def result(self, _timeout: float | None = None) -> Any: return self._value +class MethodCallProxy: + """Proxy that wraps an Actor to support method calls returning ActorFuture. + + Used as the owner of RemoteOut/RemoteIn on the parent side so that calls like + `owner.set_transport(name, value).result()` work through the pipe to the worker. + """ + + def __init__(self, actor: Actor) -> None: + self._actor = actor + + def __reduce__(self) -> tuple[type, tuple[Actor]]: + return (MethodCallProxy, (self._actor,)) + + def __getattr__(self, name: str) -> Any: + # Don't intercept private/dunder attributes - they must follow normal lookup. + if name.startswith("_"): + raise AttributeError(name) + + def _call(*args: Any, **kwargs: Any) -> ActorFuture: + result = self._actor._send_request_to_worker( + {"type": "call_method", "name": name, "args": args, "kwargs": kwargs} + ) + return ActorFuture(result) + + return _call + + class Actor: """Proxy that forwards method calls to the worker process.""" def __init__( - self, conn: Connection | None, module_class: type[ModuleT], worker_id: int + self, + conn: Connection | None, + module_class: type[ModuleT], + worker_id: int, + module_id: int = 0, + lock: threading.Lock | None = None, ) -> None: self._conn = conn self._cls = module_class self._worker_id = worker_id + self._module_id = module_id + self._lock = lock - def __reduce__(self) -> tuple[type, tuple[None, type, int]]: - """Exclude the connection when pickling - it can't be used in other processes.""" - return (Actor, (None, self._cls, self._worker_id)) + def __reduce__(self) -> tuple[type, tuple[None, type, int, int, None]]: + """Exclude the connection and lock when pickling.""" + return (Actor, (None, self._cls, self._worker_id, self._module_id, None)) def _send_request_to_worker(self, request: dict[str, Any]) -> Any: if self._conn is None: raise RuntimeError("Actor connection not available - cannot send requests") - self._conn.send(request) - response = self._conn.recv() + request["module_id"] = self._module_id + if self._lock is not None: + with self._lock: + self._conn.send(request) + response = self._conn.recv() + else: + self._conn.send(request) + response = self._conn.recv() if response.get("error"): if "AttributeError" in response["error"]: # TODO: better error handling raise AttributeError(response["error"]) @@ -92,24 +136,28 @@ def reset_forkserver_context() -> None: _forkserver_ctx = None -_seq_ids = SequentialIds() +_worker_ids = SequentialIds() +_module_ids = SequentialIds() class Worker: - def __init__( - self, - module_class: type[ModuleT], - args: tuple[Any, ...] = (), - kwargs: dict[Any, Any] | None = None, - ) -> None: - self._module_class: type[ModuleT] = module_class - self._args: tuple[Any, ...] = args - self._kwargs: dict[Any, Any] = kwargs or {} + """Generic worker process that can host multiple modules.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._modules: dict[int, Actor] = {} + self._reserved: int = 0 self._process: Any = None self._conn: Connection | None = None - self._actor: Actor | None = None - self._worker_id: int = _seq_ids.next() - self._ready: bool = False + self._worker_id: int = _worker_ids.next() + + @property + def module_count(self) -> int: + return len(self._modules) + self._reserved + + def reserve_slot(self) -> None: + """Reserve a slot so _select_worker() sees the pending load.""" + self._reserved += 1 def start_process(self) -> None: ctx = get_forkserver_context() @@ -118,49 +166,76 @@ def start_process(self) -> None: self._process = ctx.Process( target=_worker_entrypoint, - args=(child_conn, self._module_class, self._args, self._kwargs, self._worker_id), + args=(child_conn, self._worker_id), daemon=True, ) self._process.start() - self._actor = Actor(parent_conn, self._module_class, self._worker_id) - def wait_until_ready(self) -> None: - if self._ready: - return - if self._actor is None: + def deploy_module( + self, + module_class: type[ModuleT], + args: tuple[Any, ...] = (), + kwargs: dict[Any, Any] | None = None, + ) -> Actor: + if self._conn is None: raise RuntimeError("Worker process not started") - worker_id = self._actor.set_ref(self._actor).result() - ActorRegistry.update(str(self._actor), str(worker_id)) - self._ready = True + kwargs = kwargs or {} + module_id = _module_ids.next() + + # Send deploy_module request to the worker process + request = { + "type": "deploy_module", + "module_id": module_id, + "module_class": module_class, + "args": args, + "kwargs": kwargs, + } + with self._lock: + self._conn.send(request) + response = self._conn.recv() - logger.info( - "Deployed module.", module=self._module_class.__name__, worker_id=self._worker_id - ) + if response.get("error"): + raise RuntimeError(f"Failed to deploy module: {response['error']}") - def deploy(self) -> None: - self.start_process() - self.wait_until_ready() + actor = Actor(self._conn, module_class, self._worker_id, module_id, self._lock) + actor.set_ref(actor).result() - def get_instance(self) -> RPCClient: - if self._actor is None: - raise RuntimeError("Worker not deployed") - return RPCClient(self._actor, self._module_class) + self._modules[module_id] = actor + self._reserved = max(0, self._reserved - 1) + logger.info( + "Deployed module.", + module=module_class.__name__, + worker_id=self._worker_id, + module_id=module_id, + ) + return actor def shutdown(self) -> None: if self._conn is not None: try: - self._conn.send({"type": "shutdown"}) - self._conn.recv() - except (BrokenPipeError, EOFError): + with self._lock: + self._conn.send({"type": "shutdown"}) + if self._conn.poll(timeout=5): + self._conn.recv() + else: + logger.warning( + "Worker did not respond to shutdown within 5s, closing pipe.", + worker_id=self._worker_id, + ) + except (BrokenPipeError, EOFError, ConnectionResetError): pass finally: self._conn.close() self._conn = None if self._process is not None: - self._process.join(timeout=2) + self._process.join(timeout=5) if self._process.is_alive(): + logger.warning( + "Worker still alive after 5s, terminating.", + worker_id=self._worker_id, + ) self._process.terminate() self._process.join(timeout=1) self._process = None @@ -168,29 +243,43 @@ def shutdown(self) -> None: def _worker_entrypoint( conn: Connection, - module_class: type[ModuleT], - args: tuple[Any, ...], - kwargs: dict[Any, Any], worker_id: int, ) -> None: - instance = None + instances: dict[int, Any] = {} try: - instance = module_class(*args, **kwargs) - instance.worker = worker_id - - _worker_loop(conn, instance, worker_id) + _worker_loop(conn, instances, worker_id) + except KeyboardInterrupt: + logger.info("Worker got KeyboardInterrupt.", worker_id=worker_id) except Exception as e: logger.error(f"Worker process error: {e}", exc_info=True) finally: - if instance is not None: + for module_id, instance in reversed(list(instances.items())): try: + logger.info( + "Worker stopping module...", + module=type(instance).__name__, + worker_id=worker_id, + module_id=module_id, + ) instance.stop() + logger.info( + "Worker module stopped.", + module=type(instance).__name__, + worker_id=worker_id, + module_id=module_id, + ) + except KeyboardInterrupt: + logger.warning( + "KeyboardInterrupt during worker stop", + module=type(instance).__name__, + worker_id=worker_id, + ) except Exception: logger.error("Error during worker shutdown", exc_info=True) -def _worker_loop(conn: Connection, instance: Any, worker_id: int) -> None: +def _worker_loop(conn: Connection, instances: dict[int, Any], worker_id: int) -> None: while True: try: if not conn.poll(timeout=0.1): @@ -203,12 +292,29 @@ def _worker_loop(conn: Connection, instance: Any, worker_id: int) -> None: try: req_type = request.get("type") - if req_type == "set_ref": - instance.ref = request.get("ref") + if req_type == "deploy_module": + module_class = request["module_class"] + args = request.get("args", ()) + kwargs = request.get("kwargs", {}) + module_id = request["module_id"] + instance = module_class(*args, **kwargs) + instances[module_id] = instance + response["result"] = module_id + + elif req_type == "set_ref": + module_id = request["module_id"] + instances[module_id].ref = request.get("ref") response["result"] = worker_id elif req_type == "getattr": - response["result"] = getattr(instance, request["name"]) + module_id = request["module_id"] + response["result"] = getattr(instances[module_id], request["name"]) + + elif req_type == "call_method": + module_id = request["module_id"] + method = getattr(instances[module_id], request["name"]) + result = method(*request.get("args", ()), **request.get("kwargs", {})) + response["result"] = result elif req_type == "shutdown": response["result"] = True diff --git a/dimos/core/worker_manager.py b/dimos/core/worker_manager.py index 175b650fd2..d80b431c50 100644 --- a/dimos/core/worker_manager.py +++ b/dimos/core/worker_manager.py @@ -12,30 +12,48 @@ # See the License for the specific language governing permissions and # limitations under the License. +from concurrent.futures import ThreadPoolExecutor from typing import Any from dimos.core.module import ModuleT from dimos.core.rpc_client import RPCClient from dimos.core.worker import Worker -from dimos.utils.actor_registry import ActorRegistry from dimos.utils.logging_config import setup_logger logger = setup_logger() class WorkerManager: - def __init__(self) -> None: + def __init__(self, n_workers: int = 2) -> None: + self._n_workers = n_workers self._workers: list[Worker] = [] self._closed = False + self._started = False + + def start(self) -> None: + if self._started: + return + self._started = True + for _ in range(self._n_workers): + worker = Worker() + worker.start_process() + self._workers.append(worker) + logger.info("Worker pool started.", n_workers=self._n_workers) + + def _select_worker(self) -> Worker: + return min(self._workers, key=lambda w: w.module_count) def deploy(self, module_class: type[ModuleT], *args: Any, **kwargs: Any) -> RPCClient: if self._closed: raise RuntimeError("WorkerManager is closed") - worker = Worker(module_class, args=args, kwargs=kwargs) - worker.deploy() - self._workers.append(worker) - return worker.get_instance() + # Auto-start for backward compatibility + if not self._started: + self.start() + + worker = self._select_worker() + actor = worker.deploy_module(module_class, args=args, kwargs=kwargs) + return RPCClient(actor, module_class) def deploy_parallel( self, module_specs: list[tuple[type[ModuleT], tuple[Any, ...], dict[Any, Any]]] @@ -43,17 +61,30 @@ def deploy_parallel( if self._closed: raise RuntimeError("WorkerManager is closed") - workers: list[Worker] = [] + # Auto-start for backward compatibility + if not self._started: + self.start() + + # Pre-assign workers sequentially (so least-loaded accounting is + # correct), then deploy concurrently via threads. The per-worker lock + # serializes deploys that land on the same worker process. + assignments: list[tuple[Worker, type[ModuleT], tuple[Any, ...], dict[Any, Any]]] = [] for module_class, args, kwargs in module_specs: - worker = Worker(module_class, args=args, kwargs=kwargs) - worker.start_process() - workers.append(worker) + worker = self._select_worker() + worker.reserve_slot() + assignments.append((worker, module_class, args, kwargs)) - for worker in workers: - worker.wait_until_ready() - self._workers.append(worker) + def _deploy( + item: tuple[Worker, type[ModuleT], tuple[Any, ...], dict[Any, Any]], + ) -> RPCClient: + worker, module_class, args, kwargs = item + actor = worker.deploy_module(module_class, args=args, kwargs=kwargs) + return RPCClient(actor, module_class) + + with ThreadPoolExecutor(max_workers=len(assignments)) as pool: + results = list(pool.map(_deploy, assignments)) - return [worker.get_instance() for worker in workers] + return results def close_all(self) -> None: if self._closed: @@ -69,6 +100,5 @@ def close_all(self) -> None: logger.error(f"Error shutting down worker: {e}", exc_info=True) self._workers.clear() - ActorRegistry.clear() logger.info("All workers shut down") diff --git a/dimos/hardware/manipulators/README.md b/dimos/hardware/manipulators/README.md index 60d3c94567..08c679726e 100644 --- a/dimos/hardware/manipulators/README.md +++ b/dimos/hardware/manipulators/README.md @@ -96,7 +96,9 @@ class MyArmAdapter: # No inheritance needed - just match the Protocol 2. **Create the driver** (`arm.py`): ```python -from dimos.core import Module, ModuleConfig, In, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out from .adapter import MyArmAdapter class MyArm(Module[MyArmConfig]): diff --git a/dimos/hardware/sensors/camera/gstreamer/gstreamer_camera_test_script.py b/dimos/hardware/sensors/camera/gstreamer/gstreamer_camera_test_script.py index cc0e3424a5..8785a9260b 100755 --- a/dimos/hardware/sensors/camera/gstreamer/gstreamer_camera_test_script.py +++ b/dimos/hardware/sensors/camera/gstreamer/gstreamer_camera_test_script.py @@ -18,7 +18,8 @@ import logging import time -from dimos import core +from dimos.core.module_coordinator import ModuleCoordinator +from dimos.core.transport import LCMTransport from dimos.hardware.sensors.camera.gstreamer.gstreamer_camera import GstreamerCameraModule from dimos.msgs.sensor_msgs import Image from dimos.protocol import pubsub @@ -61,8 +62,8 @@ def main() -> None: pubsub.lcm.autoconf() # type: ignore[attr-defined] # Start dimos - logger.info("Starting dimos...") - dimos = core.start(8) + dimos = ModuleCoordinator() + dimos.start() # Deploy the GStreamer camera module logger.info(f"Deploying GStreamer TCP camera module (connecting to {args.host}:{args.port})...") @@ -75,7 +76,7 @@ def main() -> None: ) # Set up LCM transport for the video output - camera.video.transport = core.LCMTransport("/zed/video", Image) + camera.video.transport = LCMTransport("/zed/video", Image) # Counter for received frames frame_count = [0] @@ -123,9 +124,7 @@ def on_frame(msg) -> None: # type: ignore[no-untyped-def] while True: time.sleep(1) except KeyboardInterrupt: - logger.info("Shutting down...") - camera.stop() - logger.info("Stopped.") + dimos.stop() if __name__ == "__main__": diff --git a/dimos/hardware/sensors/camera/test_webcam.py b/dimos/hardware/sensors/camera/test_webcam.py deleted file mode 100644 index e40a73acc9..0000000000 --- a/dimos/hardware/sensors/camera/test_webcam.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright 2025-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. - -import time - -import pytest - -from dimos import core -from dimos.hardware.sensors.camera import zed -from dimos.hardware.sensors.camera.module import CameraModule -from dimos.hardware.sensors.camera.webcam import Webcam -from dimos.msgs.geometry_msgs import Quaternion, Transform, Vector3 -from dimos.msgs.sensor_msgs import CameraInfo, Image - - -@pytest.fixture -def dimos(): - dimos_instance = core.start(1) - yield dimos_instance - dimos_instance.stop() - - -@pytest.mark.tool -def test_streaming_single(dimos) -> None: - camera = dimos.deploy( - CameraModule, - transform=Transform( - translation=Vector3(0.05, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="sensor", - child_frame_id="camera_link", - ), - hardware=lambda: Webcam( - camera_index=0, - fps=0.0, # full speed but set something to test sharpness barrier - camera_info=zed.CameraInfo.SingleWebcam, - ), - ) - - camera.color_image.transport = core.LCMTransport("/color_image", Image) - camera.camera_info.transport = core.LCMTransport("/camera_info", CameraInfo) - camera.start() - - try: - while True: - time.sleep(1) - except KeyboardInterrupt: - camera.stop() - dimos.stop() diff --git a/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py b/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py index 05801729e3..b1a6baef44 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py +++ b/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py @@ -26,7 +26,7 @@ "world/lidar": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), } ), -).global_config(n_dask_workers=2, robot_model="mid360_fastlio2") +).global_config(n_workers=2, robot_model="mid360_fastlio2") mid360_fastlio_voxels = autoconnect( FastLio2.blueprint(), @@ -37,7 +37,7 @@ "world/lidar": None, } ), -).global_config(n_dask_workers=3, robot_model="mid360_fastlio2_voxels") +).global_config(n_workers=3, robot_model="mid360_fastlio2_voxels") mid360_fastlio_voxels_native = autoconnect( FastLio2.blueprint(voxel_size=voxel_size, map_voxel_size=voxel_size, map_freq=3.0), @@ -47,4 +47,4 @@ "world/global_map": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), } ), -).global_config(n_dask_workers=2, robot_model="mid360_fastlio2") +).global_config(n_workers=2, robot_model="mid360_fastlio2") diff --git a/dimos/hardware/sensors/lidar/fastlio2/module.py b/dimos/hardware/sensors/lidar/fastlio2/module.py index ee9a0783a0..fb894ddce5 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/module.py +++ b/dimos/hardware/sensors/lidar/fastlio2/module.py @@ -34,8 +34,8 @@ from pathlib import Path from typing import TYPE_CHECKING -from dimos.core import Out # noqa: TC001 from dimos.core.native_module import NativeModule, NativeModuleConfig +from dimos.core.stream import Out # noqa: TC001 from dimos.hardware.sensors.lidar.livox.ports import ( SDK_CMD_DATA_PORT, SDK_HOST_CMD_DATA_PORT, diff --git a/dimos/hardware/sensors/lidar/livox/livox_blueprints.py b/dimos/hardware/sensors/lidar/livox/livox_blueprints.py index 0cda912b73..9ded4578ba 100644 --- a/dimos/hardware/sensors/lidar/livox/livox_blueprints.py +++ b/dimos/hardware/sensors/lidar/livox/livox_blueprints.py @@ -19,4 +19,4 @@ mid360 = autoconnect( Mid360.blueprint(), rerun_bridge(), -).global_config(n_dask_workers=2, robot_model="mid360") +).global_config(n_workers=2, robot_model="mid360") diff --git a/dimos/hardware/sensors/lidar/livox/module.py b/dimos/hardware/sensors/lidar/livox/module.py index 672968a0eb..2e470b21ef 100644 --- a/dimos/hardware/sensors/lidar/livox/module.py +++ b/dimos/hardware/sensors/lidar/livox/module.py @@ -29,8 +29,8 @@ from dataclasses import dataclass from typing import TYPE_CHECKING -from dimos.core import Out # noqa: TC001 from dimos.core.native_module import NativeModule, NativeModuleConfig +from dimos.core.stream import Out # noqa: TC001 from dimos.hardware.sensors.lidar.livox.ports import ( SDK_CMD_DATA_PORT, SDK_HOST_CMD_DATA_PORT, diff --git a/dimos/manipulation/control/dual_trajectory_setter.py b/dimos/manipulation/control/dual_trajectory_setter.py index 4f8a8802e1..05793eeb76 100644 --- a/dimos/manipulation/control/dual_trajectory_setter.py +++ b/dimos/manipulation/control/dual_trajectory_setter.py @@ -33,7 +33,7 @@ import sys import time -from dimos import core +from dimos.core.transport import LCMTransport from dimos.manipulation.planning.trajectory_generator.joint_trajectory_generator import ( JointTrajectoryGenerator, ) @@ -85,20 +85,16 @@ def __init__( self.right = ArmState(name="right") # Publishers for trajectories - self.left_trajectory_pub: core.LCMTransport[JointTrajectory] = core.LCMTransport( + self.left_trajectory_pub: LCMTransport[JointTrajectory] = LCMTransport( left_trajectory_topic, JointTrajectory ) - self.right_trajectory_pub: core.LCMTransport[JointTrajectory] = core.LCMTransport( + self.right_trajectory_pub: LCMTransport[JointTrajectory] = LCMTransport( right_trajectory_topic, JointTrajectory ) # Subscribers for joint states - self.left_joint_sub: core.LCMTransport[JointState] = core.LCMTransport( - left_joint_topic, JointState - ) - self.right_joint_sub: core.LCMTransport[JointState] = core.LCMTransport( - right_joint_topic, JointState - ) + self.left_joint_sub: LCMTransport[JointState] = LCMTransport(left_joint_topic, JointState) + self.right_joint_sub: LCMTransport[JointState] = LCMTransport(right_joint_topic, JointState) print("DualTrajectorySetter initialized") print(f" Left arm: {left_joint_topic} -> {left_trajectory_topic}") diff --git a/dimos/manipulation/control/servo_control/README.md b/dimos/manipulation/control/servo_control/README.md index fb11fdb2a4..f20ba40b9a 100644 --- a/dimos/manipulation/control/servo_control/README.md +++ b/dimos/manipulation/control/servo_control/README.md @@ -172,7 +172,7 @@ arm_driver = XArmDriver(config=XArmDriverConfig(ip_address="192.168.1.235")) controller = CartesianMotionController(arm_driver=arm_driver) # Set up LCM transports for target_pose and current_pose -from dimos.core import LCMTransport +from dimos.core.transport import LCMTransport controller.target_pose.connection = LCMTransport("/target_pose", PoseStamped) controller.current_pose.connection = LCMTransport("/xarm/current_pose", PoseStamped) diff --git a/dimos/manipulation/control/servo_control/cartesian_motion_controller.py b/dimos/manipulation/control/servo_control/cartesian_motion_controller.py index f5a0810803..2c11b0cc10 100644 --- a/dimos/manipulation/control/servo_control/cartesian_motion_controller.py +++ b/dimos/manipulation/control/servo_control/cartesian_motion_controller.py @@ -32,8 +32,9 @@ import time from typing import Any -from dimos.core import In, Module, Out, rpc -from dimos.core.module import ModuleConfig +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import Pose, PoseStamped, Quaternion, Twist, Vector3 from dimos.msgs.sensor_msgs import JointCommand, JointState, RobotState from dimos.utils.logging_config import setup_logger diff --git a/dimos/manipulation/control/target_setter.py b/dimos/manipulation/control/target_setter.py index 1a937d12bb..f54a6af2f0 100644 --- a/dimos/manipulation/control/target_setter.py +++ b/dimos/manipulation/control/target_setter.py @@ -24,7 +24,7 @@ import sys import time -from dimos import core +from dimos.core.transport import LCMTransport from dimos.msgs.geometry_msgs import PoseStamped, Quaternion, Vector3 @@ -40,12 +40,10 @@ class TargetSetter: def __init__(self) -> None: """Initialize the target setter.""" # Create LCM transport for publishing targets - self.target_pub: core.LCMTransport[PoseStamped] = core.LCMTransport( - "/target_pose", PoseStamped - ) + self.target_pub: LCMTransport[PoseStamped] = LCMTransport("/target_pose", PoseStamped) # Subscribe to current pose from controller - self.current_pose_sub: core.LCMTransport[PoseStamped] = core.LCMTransport( + self.current_pose_sub: LCMTransport[PoseStamped] = LCMTransport( "/xarm/current_pose", PoseStamped ) self.latest_current_pose: PoseStamped | None = None diff --git a/dimos/manipulation/control/trajectory_controller/joint_trajectory_controller.py b/dimos/manipulation/control/trajectory_controller/joint_trajectory_controller.py index 6ecdff1714..1ce3149dd2 100644 --- a/dimos/manipulation/control/trajectory_controller/joint_trajectory_controller.py +++ b/dimos/manipulation/control/trajectory_controller/joint_trajectory_controller.py @@ -34,8 +34,9 @@ import time from typing import Any -from dimos.core import In, Module, Out, rpc -from dimos.core.module import ModuleConfig +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out from dimos.msgs.sensor_msgs import JointCommand, JointState, RobotState from dimos.msgs.trajectory_msgs import JointTrajectory, TrajectoryState, TrajectoryStatus from dimos.utils.logging_config import setup_logger diff --git a/dimos/manipulation/control/trajectory_controller/spec.py b/dimos/manipulation/control/trajectory_controller/spec.py index 3da272a5b9..e11da91847 100644 --- a/dimos/manipulation/control/trajectory_controller/spec.py +++ b/dimos/manipulation/control/trajectory_controller/spec.py @@ -29,7 +29,7 @@ from typing import TYPE_CHECKING, Protocol if TYPE_CHECKING: - from dimos.core import In, Out + from dimos.core.stream import In, Out from dimos.msgs.sensor_msgs import JointCommand, JointState, RobotState from dimos.msgs.trajectory_msgs import JointTrajectory as JointTrajectoryMsg, TrajectoryState diff --git a/dimos/manipulation/control/trajectory_setter.py b/dimos/manipulation/control/trajectory_setter.py index bad3854521..a5baa512b5 100644 --- a/dimos/manipulation/control/trajectory_setter.py +++ b/dimos/manipulation/control/trajectory_setter.py @@ -32,7 +32,7 @@ import sys import time -from dimos import core +from dimos.core.transport import LCMTransport from dimos.manipulation.planning.trajectory_generator.joint_trajectory_generator import ( JointTrajectoryGenerator, ) @@ -64,15 +64,13 @@ def __init__(self, arm_type: str = "xarm"): self.arm_type = arm_type.lower() # Publisher for trajectories - self.trajectory_pub: core.LCMTransport[JointTrajectory] = core.LCMTransport( + self.trajectory_pub: LCMTransport[JointTrajectory] = LCMTransport( "/trajectory", JointTrajectory ) # Subscribe to arm-specific joint state topic joint_state_topic = f"/{self.arm_type}/joint_states" - self.joint_state_sub: core.LCMTransport[JointState] = core.LCMTransport( - joint_state_topic, JointState - ) + self.joint_state_sub: LCMTransport[JointState] = LCMTransport(joint_state_topic, JointState) self.latest_joint_state: JointState | None = None # Will be set dynamically from joint_state diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 310b77d766..86b14fdf8d 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -20,21 +20,21 @@ - @skill (long-horizon): Multi-step composed behaviors (pick, place, place_back, pick_and_place) """ -from __future__ import annotations - from dataclasses import dataclass, field from enum import Enum import math from pathlib import Path import threading import time -from typing import TYPE_CHECKING, Any, TypeAlias +from typing import Any, TypeAlias from dimos.agents.annotation import skill from dimos.constants import DIMOS_PROJECT_ROOT -from dimos.core import In, Module, rpc +from dimos.core.core import rpc from dimos.core.docker_runner import DockerModule as DockerRunner -from dimos.core.module import ModuleConfig +from dimos.core.module import Module, ModuleConfig +from dimos.core.rpc_client import RPCClient +from dimos.core.stream import In from dimos.manipulation.grasping.graspgen_module import GraspGenModule from dimos.manipulation.planning import ( JointPath, @@ -50,21 +50,13 @@ create_planner, ) from dimos.manipulation.planning.monitor import WorldMonitor -from dimos.msgs.geometry_msgs import Pose, Quaternion, Vector3 - -# These must be imported at runtime (not TYPE_CHECKING) for In/Out port creation -from dimos.msgs.sensor_msgs import JointState +from dimos.msgs.geometry_msgs import Pose, PoseArray, Quaternion, Vector3 +from dimos.msgs.sensor_msgs import JointState, PointCloud2 from dimos.msgs.trajectory_msgs import JointTrajectory from dimos.perception.detection.type.detection3d.object import Object as DetObject from dimos.utils.data import get_data from dimos.utils.logging_config import setup_logger -if TYPE_CHECKING: - from dimos.core.rpc_client import RPCClient - from dimos.msgs.geometry_msgs import PoseArray - from dimos.msgs.sensor_msgs import PointCloud2 - from dimos.perception.detection.type.detection3d.object import Object as DetObject - logger = setup_logger() # Composite type aliases for readability (using semantic IDs from planning.spec) @@ -234,7 +226,7 @@ def _initialize_planning(self) -> None: # Start TF publishing thread if any robot has tf_extra_links if any(c.tf_extra_links for _, c, _ in self._robots.values()): - _ = self.tf # Eager init — lazy init blocks in Dask workers + _ = self.tf # Eager init self._tf_stop_event.clear() self._tf_thread = threading.Thread( target=self._tf_publish_loop, name="ManipTFThread", daemon=True diff --git a/dimos/mapping/costmapper.py b/dimos/mapping/costmapper.py index 70cd770777..fa0ce826f2 100644 --- a/dimos/mapping/costmapper.py +++ b/dimos/mapping/costmapper.py @@ -17,9 +17,10 @@ from reactivex import operators as ops -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc from dimos.core.global_config import GlobalConfig, global_config -from dimos.core.module import ModuleConfig +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out from dimos.mapping.pointclouds.occupancy import ( OCCUPANCY_ALGOS, HeightCostConfig, diff --git a/dimos/mapping/pointclouds/test_occupancy.py b/dimos/mapping/pointclouds/test_occupancy.py index 2e301c772d..d265800f24 100644 --- a/dimos/mapping/pointclouds/test_occupancy.py +++ b/dimos/mapping/pointclouds/test_occupancy.py @@ -18,7 +18,7 @@ from open3d.geometry import PointCloud import pytest -from dimos.core import LCMTransport +from dimos.core.transport import LCMTransport from dimos.mapping.occupancy.visualizations import visualize_occupancy_grid from dimos.mapping.pointclouds.occupancy import ( height_cost_occupancy, diff --git a/dimos/mapping/test_voxels.py b/dimos/mapping/test_voxels.py index 8fdb1f2827..3fdc2dd102 100644 --- a/dimos/mapping/test_voxels.py +++ b/dimos/mapping/test_voxels.py @@ -18,7 +18,7 @@ import numpy as np import pytest -from dimos.core import LCMTransport +from dimos.core.transport import LCMTransport from dimos.mapping.voxels import VoxelGridMapper from dimos.msgs.sensor_msgs import PointCloud2 from dimos.utils.data import get_data diff --git a/dimos/mapping/voxels.py b/dimos/mapping/voxels.py index 4c1805e059..df0e730868 100644 --- a/dimos/mapping/voxels.py +++ b/dimos/mapping/voxels.py @@ -22,9 +22,10 @@ from reactivex.disposable import Disposable from reactivex.subject import Subject -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc from dimos.core.global_config import GlobalConfig, global_config -from dimos.core.module import ModuleConfig +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out from dimos.msgs.sensor_msgs import PointCloud2 from dimos.utils.decorators import simple_mcache from dimos.utils.logging_config import setup_logger diff --git a/dimos/memory/embedding.py b/dimos/memory/embedding.py index 20dd82422c..4627ecfc35 100644 --- a/dimos/memory/embedding.py +++ b/dimos/memory/embedding.py @@ -20,8 +20,9 @@ from reactivex import operators as ops from reactivex.observable import Observable -from dimos.core import In, rpc +from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In from dimos.models.embedding.base import Embedding, EmbeddingModel from dimos.models.embedding.clip import CLIPModel from dimos.msgs.geometry_msgs import PoseStamped diff --git a/dimos/models/__init__.py b/dimos/models/__init__.py index d8e2e14341..e69de29bb2 100644 --- a/dimos/models/__init__.py +++ b/dimos/models/__init__.py @@ -1,3 +0,0 @@ -from dimos.models.base import HuggingFaceModel, LocalModel - -__all__ = ["HuggingFaceModel", "LocalModel"] diff --git a/dimos/models/qwen/bbox.py b/dimos/models/qwen/bbox.py new file mode 100644 index 0000000000..bddee8308f --- /dev/null +++ b/dimos/models/qwen/bbox.py @@ -0,0 +1 @@ +BBox = tuple[float, float, float, float] # (x1, y1, x2, y2) diff --git a/dimos/models/qwen/video_query.py b/dimos/models/qwen/video_query.py index 7ba80ae069..05b93028d7 100644 --- a/dimos/models/qwen/video_query.py +++ b/dimos/models/qwen/video_query.py @@ -10,10 +10,9 @@ from dimos.agents_deprecated.agent import OpenAIAgent from dimos.agents_deprecated.tokenizer.huggingface_tokenizer import HuggingFaceTokenizer +from dimos.models.qwen.bbox import BBox from dimos.utils.threadpool import get_scheduler -BBox = tuple[float, float, float, float] # (x1, y1, x2, y2) - def query_single_frame_observable( video_observable: Observable, # type: ignore[type-arg] diff --git a/dimos/models/vl/base.py b/dimos/models/vl/base.py index 93caba4de7..237feb1d1b 100644 --- a/dimos/models/vl/base.py +++ b/dimos/models/vl/base.py @@ -1,17 +1,22 @@ +from __future__ import annotations + from abc import ABC, abstractmethod from dataclasses import dataclass import json import logging +from typing import TYPE_CHECKING import warnings from dimos.core.resource import Resource from dimos.msgs.sensor_msgs import Image -from dimos.perception.detection.type import Detection2DBBox, Detection2DPoint, ImageDetections2D from dimos.protocol.service import Configurable # type: ignore[attr-defined] from dimos.utils.data import get_data from dimos.utils.decorators import retry from dimos.utils.llm_utils import extract_json +if TYPE_CHECKING: + from dimos.perception.detection.type import Detection2DBBox, Detection2DPoint, ImageDetections2D + logger = logging.getLogger(__name__) @@ -64,6 +69,9 @@ def vlm_detection_to_detection2d( Returns: Detection2DBBox instance or None if invalid """ + # Here to prevent unwanted imports in the file. + from dimos.perception.detection.type import Detection2DBBox + # Validate list/tuple structure if not isinstance(vlm_detection, (list, tuple)): logger.debug(f"VLM detection is not a list/tuple: {type(vlm_detection)}") @@ -119,6 +127,8 @@ def vlm_point_to_detection2d_point( Returns: Detection2DPoint instance or None if invalid """ + from dimos.perception.detection.type import Detection2DPoint + # Validate list/tuple structure if not isinstance(vlm_point, (list, tuple)): logger.debug(f"VLM point is not a list/tuple: {type(vlm_point)}") @@ -245,6 +255,9 @@ def query_json(self, image: Image, query: str) -> dict: # type: ignore[type-arg def query_detections( self, image: Image, query: str, **kwargs: object ) -> ImageDetections2D[Detection2DBBox]: + # Here to prevent unwanted imports in the file. + from dimos.perception.detection.type import ImageDetections2D + full_query = f"""show me bounding boxes in pixels for this query: `{query}` format should be: @@ -303,6 +316,9 @@ def query_points( Returns: ImageDetections2D containing Detection2DPoint instances """ + # Here to prevent unwanted imports in the file. + from dimos.perception.detection.type import ImageDetections2D + full_query = f"""Show me point coordinates in pixels for this query: `{query}` The format should be: diff --git a/dimos/models/vl/test_base.py b/dimos/models/vl/test_base.py index 2e4229d944..0cc5c90d0e 100644 --- a/dimos/models/vl/test_base.py +++ b/dimos/models/vl/test_base.py @@ -3,7 +3,7 @@ from dimos_lcm.foxglove_msgs.ImageAnnotations import ImageAnnotations import pytest -from dimos.core import LCMTransport +from dimos.core.transport import LCMTransport from dimos.models.vl.moondream import MoondreamVlModel from dimos.models.vl.qwen import QwenVlModel from dimos.msgs.sensor_msgs import Image, ImageFormat diff --git a/dimos/models/vl/test_vlm.py b/dimos/models/vl/test_vlm.py index 54ceddadc5..43dad0ef94 100644 --- a/dimos/models/vl/test_vlm.py +++ b/dimos/models/vl/test_vlm.py @@ -7,7 +7,7 @@ ) import pytest -from dimos.core import LCMTransport +from dimos.core.transport import LCMTransport from dimos.models.vl.moondream import MoondreamVlModel from dimos.models.vl.moondream_hosted import MoondreamHostedVlModel from dimos.models.vl.qwen import QwenVlModel diff --git a/dimos/navigation/bbox_navigation.py b/dimos/navigation/bbox_navigation.py index 4f4aff3d16..e0752dfd00 100644 --- a/dimos/navigation/bbox_navigation.py +++ b/dimos/navigation/bbox_navigation.py @@ -17,7 +17,9 @@ from dimos_lcm.sensor_msgs import CameraInfo from reactivex.disposable import Disposable -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import PoseStamped, Quaternion, Vector3 from dimos.msgs.vision_msgs import Detection2DArray from dimos.utils.logging_config import setup_logger diff --git a/dimos/navigation/demo_ros_navigation.py b/dimos/navigation/demo_ros_navigation.py index 4919ab0efd..4d57867d59 100644 --- a/dimos/navigation/demo_ros_navigation.py +++ b/dimos/navigation/demo_ros_navigation.py @@ -14,18 +14,19 @@ import time -from dimos import core +from dimos.core.module_coordinator import ModuleCoordinator from dimos.msgs.geometry_msgs import PoseStamped, Quaternion, Vector3 from dimos.navigation import rosnav -from dimos.protocol import pubsub +from dimos.protocol.service.lcmservice import autoconf from dimos.utils.logging_config import setup_logger logger = setup_logger() def main() -> None: - pubsub.lcm.autoconf() # type: ignore[attr-defined] - dimos = core.start(2) + autoconf() + dimos = ModuleCoordinator() + dimos.start() ros_nav = rosnav.deploy(dimos) @@ -52,8 +53,7 @@ def main() -> None: while True: time.sleep(1) except KeyboardInterrupt: - logger.info("\nShutting down...") - ros_nav.stop() + dimos.stop() if __name__ == "__main__": diff --git a/dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py b/dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py index 3adfc1c598..6e598e8316 100644 --- a/dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py +++ b/dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py @@ -29,7 +29,9 @@ from reactivex.disposable import Disposable from dimos.agents.annotation import skill -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.mapping.occupancy.inflation import simple_inflate from dimos.msgs.geometry_msgs import PoseStamped, Vector3 from dimos.msgs.nav_msgs import CostValues, OccupancyGrid diff --git a/dimos/navigation/replanning_a_star/module.py b/dimos/navigation/replanning_a_star/module.py index d1d87cbbf6..60dfad73ae 100644 --- a/dimos/navigation/replanning_a_star/module.py +++ b/dimos/navigation/replanning_a_star/module.py @@ -17,8 +17,10 @@ from dimos_lcm.std_msgs import Bool, String from reactivex.disposable import Disposable -from dimos.core import In, Module, Out, rpc +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 import PoseStamped, Twist from dimos.msgs.nav_msgs import OccupancyGrid, Path from dimos.navigation.base import NavigationInterface, NavigationState diff --git a/dimos/navigation/rosnav.py b/dimos/navigation/rosnav.py index 8efabaebee..89b299ae5b 100644 --- a/dimos/navigation/rosnav.py +++ b/dimos/navigation/rosnav.py @@ -28,10 +28,9 @@ from dimos import spec from dimos.agents.annotation import skill -from dimos.core import DimosCluster, In, LCMTransport, Module, Out, rpc -from dimos.core._dask_exports import DimosCluster from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig +from dimos.core.module_coordinator import ModuleCoordinator from dimos.core.stream import In, Out from dimos.core.transport import LCMTransport, ROSTransport from dimos.msgs.geometry_msgs import ( @@ -378,7 +377,7 @@ def stop(self) -> None: ros_nav = ROSNav.blueprint -def deploy(dimos: DimosCluster): # type: ignore[no-untyped-def] +def deploy(dimos: ModuleCoordinator): # type: ignore[no-untyped-def] nav = dimos.deploy(ROSNav) # type: ignore[attr-defined] # Existing ports on LCM transports diff --git a/dimos/navigation/visual/query.py b/dimos/navigation/visual/query.py index 2e0951951e..37b743506a 100644 --- a/dimos/navigation/visual/query.py +++ b/dimos/navigation/visual/query.py @@ -13,7 +13,7 @@ # limitations under the License. -from dimos.models.qwen.video_query import BBox +from dimos.models.qwen.bbox import BBox from dimos.models.vl.base import VlModel from dimos.msgs.sensor_msgs import Image from dimos.utils.generic import extract_json_from_llm_response diff --git a/dimos/perception/detection/conftest.py b/dimos/perception/detection/conftest.py index 3b24422c47..70338bd74f 100644 --- a/dimos/perception/detection/conftest.py +++ b/dimos/perception/detection/conftest.py @@ -21,7 +21,7 @@ from dimos_lcm.visualization_msgs.MarkerArray import MarkerArray import pytest -from dimos.core import LCMTransport +from dimos.core.transport import LCMTransport from dimos.msgs.geometry_msgs import Transform from dimos.msgs.sensor_msgs import CameraInfo, Image, PointCloud2 from dimos.msgs.vision_msgs import Detection2DArray diff --git a/dimos/perception/detection/detectors/test_bbox_detectors.py b/dimos/perception/detection/detectors/test_bbox_detectors.py index 32a509061a..2e69016eb5 100644 --- a/dimos/perception/detection/detectors/test_bbox_detectors.py +++ b/dimos/perception/detection/detectors/test_bbox_detectors.py @@ -16,7 +16,7 @@ import pytest from reactivex.disposable import CompositeDisposable -from dimos.core import LCMTransport +from dimos.core.transport import LCMTransport from dimos.msgs.sensor_msgs import Image from dimos.perception.detection.type import Detection2D, ImageDetections2D diff --git a/dimos/perception/detection/module2D.py b/dimos/perception/detection/module2D.py index cfca3b2192..f86794a1f7 100644 --- a/dimos/perception/detection/module2D.py +++ b/dimos/perception/detection/module2D.py @@ -23,8 +23,10 @@ from reactivex.subject import Subject from dimos import spec -from dimos.core import DimosCluster, In, Module, Out, rpc -from dimos.core.module import ModuleConfig +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.module_coordinator import ModuleCoordinator +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import Transform, Vector3 from dimos.msgs.sensor_msgs import CameraInfo, Image from dimos.msgs.sensor_msgs.Image import sharpness_barrier @@ -158,12 +160,12 @@ def stop(self) -> None: def deploy( # type: ignore[no-untyped-def] - dimos: DimosCluster, + dimos: ModuleCoordinator, camera: spec.Camera, prefix: str = "/detector2d", **kwargs, ) -> Detection2DModule: - from dimos.core import LCMTransport + from dimos.core.transport import LCMTransport detector = Detection2DModule(**kwargs) detector.color_image.connect(camera.color_image) diff --git a/dimos/perception/detection/module3D.py b/dimos/perception/detection/module3D.py index d275fbc85f..96ae4e8297 100644 --- a/dimos/perception/detection/module3D.py +++ b/dimos/perception/detection/module3D.py @@ -13,7 +13,7 @@ # limitations under the License. -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from dimos_lcm.foxglove_msgs.ImageAnnotations import ( ImageAnnotations, @@ -25,7 +25,9 @@ from dimos import spec from dimos.agents.annotation import skill from dimos.core.core import rpc +from dimos.core.module_coordinator import ModuleCoordinator from dimos.core.stream import In, Out +from dimos.core.transport import LCMTransport from dimos.msgs.geometry_msgs import PoseStamped, Quaternion, Transform, Vector3 from dimos.msgs.sensor_msgs import Image, PointCloud2 from dimos.msgs.vision_msgs import Detection2DArray @@ -37,9 +39,7 @@ from dimos.utils.reactive import backpressure if TYPE_CHECKING: - from dask.distributed import Client as DimosCluster -else: - DimosCluster = Any + from dimos.core.rpc_client import ModuleProxy class Detection3DModule(Detection2DModule): @@ -202,14 +202,12 @@ def _publish_detections(self, detections: ImageDetections3DPC) -> None: def deploy( # type: ignore[no-untyped-def] - dimos: DimosCluster, + dimos: ModuleCoordinator, lidar: spec.Pointcloud, camera: spec.Camera, prefix: str = "/detector3d", **kwargs, -) -> Detection3DModule: - from dimos.core import LCMTransport - +) -> "ModuleProxy": detector = dimos.deploy(Detection3DModule, camera_info=camera.hardware_camera_info, **kwargs) # type: ignore[attr-defined] detector.image.connect(camera.color_image) @@ -229,7 +227,7 @@ def deploy( # type: ignore[no-untyped-def] detector.start() - return detector # type: ignore[no-any-return] + return detector detection3d_module = Detection3DModule.blueprint diff --git a/dimos/perception/detection/person_tracker.py b/dimos/perception/detection/person_tracker.py index 50082742f0..913043f312 100644 --- a/dimos/perception/detection/person_tracker.py +++ b/dimos/perception/detection/person_tracker.py @@ -18,7 +18,9 @@ from reactivex import operators as ops from reactivex.observable import Observable -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import PoseStamped, Transform, Vector3 from dimos.msgs.sensor_msgs import CameraInfo, Image from dimos.msgs.vision_msgs import Detection2DArray diff --git a/dimos/perception/detection/reid/test_module.py b/dimos/perception/detection/reid/test_module.py index d962da6b6c..f5672c1f67 100644 --- a/dimos/perception/detection/reid/test_module.py +++ b/dimos/perception/detection/reid/test_module.py @@ -14,7 +14,7 @@ import pytest -from dimos.core import LCMTransport +from dimos.core.transport import LCMTransport from dimos.msgs.foxglove_msgs import ImageAnnotations from dimos.perception.detection.reid.embedding_id_system import EmbeddingIDSystem from dimos.perception.detection.reid.module import ReidModule diff --git a/dimos/perception/detection/type/__init__.py b/dimos/perception/detection/type/__init__.py index 00cf943db3..b14464d4fa 100644 --- a/dimos/perception/detection/type/__init__.py +++ b/dimos/perception/detection/type/__init__.py @@ -3,12 +3,20 @@ __getattr__, __dir__, __all__ = lazy.attach( __name__, submod_attrs={ - "detection2d": [ + "detection2d.base": [ "Detection2D", + "Filter2D", + ], + "detection2d.bbox": [ "Detection2DBBox", + ], + "detection2d.person": [ "Detection2DPerson", + ], + "detection2d.point": [ "Detection2DPoint", - "Filter2D", + ], + "detection2d.imageDetections2D": [ "ImageDetections2D", ], "detection3d": [ diff --git a/dimos/perception/detection/type/detection2d/__init__.py b/dimos/perception/detection/type/detection2d/__init__.py index dc81916679..e69de29bb2 100644 --- a/dimos/perception/detection/type/detection2d/__init__.py +++ b/dimos/perception/detection/type/detection2d/__init__.py @@ -1,30 +0,0 @@ -# Copyright 2025-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 dimos.perception.detection.type.detection2d.base import Detection2D, Filter2D -from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox -from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D -from dimos.perception.detection.type.detection2d.person import Detection2DPerson -from dimos.perception.detection.type.detection2d.point import Detection2DPoint -from dimos.perception.detection.type.detection2d.seg import Detection2DSeg - -__all__ = [ - "Detection2D", - "Detection2DBBox", - "Detection2DPerson", - "Detection2DPoint", - "Detection2DSeg", - "Filter2D", - "ImageDetections2D", -] diff --git a/dimos/perception/detection/type/detection3d/base.py b/dimos/perception/detection/type/detection3d/base.py index b036584f3e..a5dbb742b8 100644 --- a/dimos/perception/detection/type/detection3d/base.py +++ b/dimos/perception/detection/type/detection3d/base.py @@ -19,7 +19,7 @@ from typing import TYPE_CHECKING from dimos.msgs.geometry_msgs import Transform -from dimos.perception.detection.type.detection2d import Detection2DBBox +from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox if TYPE_CHECKING: from dimos_lcm.sensor_msgs import CameraInfo diff --git a/dimos/perception/detection/type/detection3d/bbox.py b/dimos/perception/detection/type/detection3d/bbox.py index cf7f4ea3cc..bdf2d27a7c 100644 --- a/dimos/perception/detection/type/detection3d/bbox.py +++ b/dimos/perception/detection/type/detection3d/bbox.py @@ -23,7 +23,7 @@ from dimos.msgs.geometry_msgs import Pose, PoseStamped, Quaternion, Transform, Vector3 from dimos.msgs.std_msgs import Header from dimos.msgs.vision_msgs import Detection3D -from dimos.perception.detection.type.detection2d import Detection2DBBox +from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox @dataclass diff --git a/dimos/perception/detection/type/detection3d/object.py b/dimos/perception/detection/type/detection3d/object.py index 43702917e1..ae7a208baa 100644 --- a/dimos/perception/detection/type/detection3d/object.py +++ b/dimos/perception/detection/type/detection3d/object.py @@ -34,7 +34,7 @@ if TYPE_CHECKING: from dimos_lcm.sensor_msgs import CameraInfo - from dimos.perception.detection.type.detection2d import ImageDetections2D + from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @dataclass(kw_only=True) diff --git a/dimos/perception/detection/type/detection3d/pointcloud.py b/dimos/perception/detection/type/detection3d/pointcloud.py index 7edceb17a5..5abc03e0c8 100644 --- a/dimos/perception/detection/type/detection3d/pointcloud.py +++ b/dimos/perception/detection/type/detection3d/pointcloud.py @@ -47,7 +47,7 @@ if TYPE_CHECKING: from dimos_lcm.sensor_msgs import CameraInfo - from dimos.perception.detection.type.detection2d import Detection2DBBox + from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox @dataclass diff --git a/dimos/perception/detection/type/detection3d/pointcloud_filters.py b/dimos/perception/detection/type/detection3d/pointcloud_filters.py index 984e04bc99..59ad6200d9 100644 --- a/dimos/perception/detection/type/detection3d/pointcloud_filters.py +++ b/dimos/perception/detection/type/detection3d/pointcloud_filters.py @@ -20,7 +20,7 @@ from dimos.msgs.geometry_msgs import Transform from dimos.msgs.sensor_msgs import PointCloud2 -from dimos.perception.detection.type.detection2d import Detection2DBBox +from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox # Filters take Detection2DBBox, PointCloud2, CameraInfo, Transform and return filtered PointCloud2 or None PointCloudFilter = Callable[ diff --git a/dimos/perception/experimental/temporal_memory/temporal_memory_deploy.py b/dimos/perception/experimental/temporal_memory/temporal_memory_deploy.py index ab3cc7a0f5..b4aaa6ac94 100644 --- a/dimos/perception/experimental/temporal_memory/temporal_memory_deploy.py +++ b/dimos/perception/experimental/temporal_memory/temporal_memory_deploy.py @@ -19,7 +19,7 @@ import os from typing import TYPE_CHECKING -from dimos.core._dask_exports import DimosCluster +from dimos.core.module_coordinator import ModuleCoordinator from dimos.models.vl.base import VlModel from dimos.spec import Camera as CameraSpec @@ -30,7 +30,7 @@ def deploy( - dimos: DimosCluster, + dimos: ModuleCoordinator, camera: CameraSpec, vlm: VlModel | None = None, config: TemporalMemoryConfig | None = None, diff --git a/dimos/perception/experimental/temporal_memory/temporal_memory_example.py b/dimos/perception/experimental/temporal_memory/temporal_memory_example.py index 8ba28bb174..2c8d79ab57 100644 --- a/dimos/perception/experimental/temporal_memory/temporal_memory_example.py +++ b/dimos/perception/experimental/temporal_memory/temporal_memory_example.py @@ -23,20 +23,28 @@ """ from pathlib import Path +from typing import TYPE_CHECKING, cast from dotenv import load_dotenv -from dimos import core +from dimos.core.module_coordinator import ModuleCoordinator from dimos.hardware.sensors.camera.module import CameraModule from dimos.hardware.sensors.camera.webcam import Webcam from .temporal_memory import TemporalMemoryConfig from .temporal_memory_deploy import deploy +if TYPE_CHECKING: + from dimos.spec import Camera + # Load environment variables load_dotenv() +def _create_webcam() -> Webcam: + return Webcam(camera_index=0) + + def example_usage() -> None: """Example of how to use TemporalMemory.""" # Initialize variables to None for cleanup @@ -46,16 +54,17 @@ def example_usage() -> None: try: # Create Dimos cluster - dimos = core.start(1) + dimos = ModuleCoordinator() + dimos.start() # Deploy camera module - camera = dimos.deploy(CameraModule, hardware=lambda: Webcam(camera_index=0)) # type: ignore[attr-defined] + camera = dimos.deploy(CameraModule, hardware=_create_webcam) # type: ignore[attr-defined] camera.start() # Deploy temporal memory using the deploy function output_dir = Path("./temporal_memory_output") temporal_memory = deploy( dimos, - camera, + cast("Camera", camera), vlm=None, # Will auto-create OpenAIVlModel if None config=TemporalMemoryConfig( fps=1.0, # Process 1 frame per second @@ -130,7 +139,7 @@ def example_usage() -> None: if camera is not None: camera.stop() if dimos is not None: - dimos.close_all() # type: ignore[attr-defined] + dimos.stop() if __name__ == "__main__": diff --git a/dimos/perception/experimental/temporal_memory/test_temporal_memory_module.py b/dimos/perception/experimental/temporal_memory/test_temporal_memory_module.py index ef584a2527..9a0b9ab156 100644 --- a/dimos/perception/experimental/temporal_memory/test_temporal_memory_module.py +++ b/dimos/perception/experimental/temporal_memory/test_temporal_memory_module.py @@ -15,15 +15,17 @@ import asyncio import os import pathlib -import tempfile import time from dotenv import load_dotenv import pytest from reactivex import operators as ops -from dimos import core -from dimos.core import Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.module_coordinator import ModuleCoordinator +from dimos.core.stream import Out +from dimos.core.transport import LCMTransport from dimos.models.vl.openai import OpenAIVlModel from dimos.msgs.sensor_msgs import Image from dimos.perception.experimental.temporal_memory import TemporalMemory, TemporalMemoryConfig @@ -80,18 +82,15 @@ def stop(self) -> None: @pytest.mark.skipif_no_openai @pytest.mark.slow class TestTemporalMemoryModule: - @pytest.fixture(scope="function") - def temp_dir(self): - """Create a temporary directory for test data.""" - temp_dir = tempfile.mkdtemp(prefix="temporal_memory_test_") - yield temp_dir - @pytest.fixture(scope="function") def dimos_cluster(self): """Create and cleanup Dimos cluster.""" - dimos = core.start(1) - yield dimos - dimos.close_all() + dimos = ModuleCoordinator() + dimos.start() + try: + yield dimos + finally: + dimos.stop() @pytest.fixture(scope="function") def video_module(self, dimos_cluster): @@ -99,7 +98,7 @@ def video_module(self, dimos_cluster): data_path = get_data("unitree_office_walk") video_path = os.path.join(data_path, "video") video_module = dimos_cluster.deploy(VideoReplayModule, video_path) - video_module.video_out.transport = core.LCMTransport("/test_video", Image) + video_module.video_out.transport = LCMTransport("/test_video", Image) yield video_module try: video_module.stop() @@ -107,9 +106,9 @@ def video_module(self, dimos_cluster): logger.warning(f"Failed to stop video_module: {e}") @pytest.fixture(scope="function") - def temporal_memory(self, dimos_cluster, temp_dir): + def temporal_memory(self, dimos_cluster, tmp_path): """Create and cleanup temporal memory module.""" - output_dir = os.path.join(temp_dir, "temporal_memory_output") + output_dir = os.path.join(tmp_path, "temporal_memory_output") # Create OpenAIVlModel instance api_key = os.getenv("OPENAI_API_KEY") vlm = OpenAIVlModel(api_key=api_key) @@ -134,7 +133,7 @@ def temporal_memory(self, dimos_cluster, temp_dir): @pytest.mark.asyncio async def test_temporal_memory_module_with_replay( - self, dimos_cluster, video_module, temporal_memory, temp_dir + self, dimos_cluster, video_module, temporal_memory, tmp_path ): """Test TemporalMemory module with TimedSensorReplay inputs.""" # Connect streams @@ -213,7 +212,7 @@ async def test_temporal_memory_module_with_replay( await asyncio.sleep(0.5) # Verify files were created - stop() already saved them - output_dir = os.path.join(temp_dir, "temporal_memory_output") + output_dir = os.path.join(tmp_path, "temporal_memory_output") output_path = pathlib.Path(output_dir) assert output_path.exists(), f"Output directory should exist: {output_dir}" assert (output_path / "state.json").exists(), "state.json should exist" diff --git a/dimos/perception/object_scene_registration.py b/dimos/perception/object_scene_registration.py index cfc4ab8d3e..ee7b87b534 100644 --- a/dimos/perception/object_scene_registration.py +++ b/dimos/perception/object_scene_registration.py @@ -21,8 +21,9 @@ import open3d as o3d # type: ignore[import-untyped] from dimos.agents.annotation import skill -from dimos.core import In, Out, rpc +from dimos.core.core import rpc from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.msgs.foxglove_msgs import ImageAnnotations from dimos.msgs.sensor_msgs import CameraInfo, Image, PointCloud2 from dimos.msgs.sensor_msgs.Image import ImageFormat diff --git a/dimos/perception/object_tracker_3d.py b/dimos/perception/object_tracker_3d.py index f8143dc861..da35577d0d 100644 --- a/dimos/perception/object_tracker_3d.py +++ b/dimos/perception/object_tracker_3d.py @@ -22,7 +22,8 @@ ) import numpy as np -from dimos.core import In, Out, rpc +from dimos.core.core import rpc +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import Pose, Quaternion, Transform, Vector3 from dimos.msgs.sensor_msgs import Image, ImageFormat from dimos.msgs.std_msgs import Header diff --git a/dimos/perception/spatial_perception.py b/dimos/perception/spatial_perception.py index d7f27c31dc..bf62d50bcf 100644 --- a/dimos/perception/spatial_perception.py +++ b/dimos/perception/spatial_perception.py @@ -32,8 +32,10 @@ from dimos.agents_deprecated.memory.spatial_vector_db import SpatialVectorDB from dimos.agents_deprecated.memory.visual_memory import VisualMemory from dimos.constants import DIMOS_PROJECT_ROOT -from dimos.core import DimosCluster, In, rpc +from dimos.core.core import rpc from dimos.core.module import Module +from dimos.core.module_coordinator import ModuleCoordinator +from dimos.core.stream import In from dimos.msgs.sensor_msgs import Image from dimos.types.robot_location import RobotLocation from dimos.utils.logging_config import setup_logger @@ -53,7 +55,7 @@ class SpatialMemory(Module): """ - A Dask module for building and querying Robot spatial memory. + A Dimos module for building and querying Robot spatial memory. This module processes video frames and odometry data from LCM streams, associates them with XY locations, and stores them in a vector database @@ -198,12 +200,12 @@ def set_video(image_msg: Image) -> None: else: logger.warning("Received image message without data attribute") - unsub = self.color_image.subscribe(set_video) - self._disposables.add(Disposable(unsub)) + self._disposables.add(Disposable(self.color_image.subscribe(set_video))) # Start periodic processing using interval - unsub = interval(self._process_interval).subscribe(lambda _: self._process_frame()) # type: ignore[assignment] - self._disposables.add(unsub) + self._disposables.add( + interval(self._process_interval).subscribe(lambda _: self._process_frame()) + ) @rpc def stop(self) -> None: @@ -578,7 +580,7 @@ def query_tagged_location(self, query: str) -> RobotLocation | None: def deploy( # type: ignore[no-untyped-def] - dimos: DimosCluster, + dimos: ModuleCoordinator, camera: spec.Camera, ): spatial_memory = dimos.deploy(SpatialMemory, db_path="/tmp/spatial_memory_db") # type: ignore[attr-defined] diff --git a/dimos/perception/test_spatial_memory_module.py b/dimos/perception/test_spatial_memory_module.py index a8c42d4f0e..ac9b132a69 100644 --- a/dimos/perception/test_spatial_memory_module.py +++ b/dimos/perception/test_spatial_memory_module.py @@ -14,14 +14,16 @@ import asyncio import os -import tempfile import time import pytest from reactivex import operators as ops -from dimos import core -from dimos.core import Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.module_coordinator import ModuleCoordinator +from dimos.core.stream import Out +from dimos.core.transport import LCMTransport from dimos.msgs.geometry_msgs import Transform from dimos.msgs.sensor_msgs import Image from dimos.perception.spatial_perception import SpatialMemory @@ -109,113 +111,97 @@ def stop(self) -> None: logger.info("OdometryReplayModule stopped") -@pytest.mark.slow -@pytest.mark.skipif_in_ci -class TestSpatialMemoryModule: - @pytest.fixture(scope="function") - def temp_dir(self): - """Create a temporary directory for test data.""" - # Use standard tempfile module to ensure proper permissions - temp_dir = tempfile.mkdtemp(prefix="spatial_memory_test_") - - yield temp_dir - - @pytest.mark.asyncio - async def test_spatial_memory_module_with_replay(self, temp_dir): - """Test SpatialMemory module with TimedSensorReplay inputs.""" - - # Start Dask - dimos = core.start(1) - - try: - # Get test data paths - data_path = get_data("unitree_office_walk") - video_path = os.path.join(data_path, "video") - odom_path = os.path.join(data_path, "odom") - - # Deploy modules - # Video replay module - video_module = dimos.deploy(VideoReplayModule, video_path) - video_module.video_out.transport = core.LCMTransport("/test_video", Image) - - # Odometry replay module (publishes to tf system directly) - odom_module = dimos.deploy(OdometryReplayModule, odom_path) - - # Spatial memory module - spatial_memory = dimos.deploy( - SpatialMemory, - collection_name="test_spatial_memory", - embedding_model="clip", - embedding_dimensions=512, - min_distance_threshold=0.5, # 0.5m for test - min_time_threshold=1.0, # 1 second - db_path=os.path.join(temp_dir, "chroma_db"), - visual_memory_path=os.path.join(temp_dir, "visual_memory.pkl"), - new_memory=True, - output_dir=os.path.join(temp_dir, "images"), - ) - - # Connect video stream - spatial_memory.color_image.connect(video_module.video_out) - - # Start all modules - video_module.start() - odom_module.start() - spatial_memory.start() - logger.info("All modules started, processing in background...") - - # Wait for frames to be processed with timeout - timeout = 10.0 # 10 second timeout - start_time = time.time() - - # Keep checking stats while modules are running - while (time.time() - start_time) < timeout: - stats = spatial_memory.get_stats() - if stats["frame_count"] > 0 and stats["stored_frame_count"] > 0: - logger.info( - f"Frames processing - Frame count: {stats['frame_count']}, Stored: {stats['stored_frame_count']}" - ) - break - await asyncio.sleep(0.5) - else: - # Timeout reached - stats = spatial_memory.get_stats() - logger.error( - f"Timeout after {timeout}s - Frame count: {stats['frame_count']}, Stored: {stats['stored_frame_count']}" - ) - raise AssertionError(f"No frames processed within {timeout} seconds") - - await asyncio.sleep(2) - - mid_stats = spatial_memory.get_stats() - logger.info( - f"Mid-test stats - Frame count: {mid_stats['frame_count']}, Stored: {mid_stats['stored_frame_count']}" - ) - assert mid_stats["frame_count"] >= stats["frame_count"], ( - "Frame count should increase or stay same" - ) +@pytest.fixture() +def dimos(): + dimos = ModuleCoordinator() + dimos.start() + try: + yield dimos + finally: + dimos.stop() - # Test query while modules are still running - try: - text_results = spatial_memory.query_by_text("office") - logger.info(f"Query by text 'office' returned {len(text_results)} results") - assert len(text_results) > 0, "Should have at least one result" - except Exception as e: - logger.warning(f"Query by text failed: {e}") - final_stats = spatial_memory.get_stats() +@pytest.mark.slow +@pytest.mark.skipif_in_ci +@pytest.mark.asyncio +async def test_spatial_memory_module_with_replay(dimos, tmp_path): + """Test SpatialMemory module with TimedSensorReplay inputs.""" + # Get test data paths + data_path = get_data("unitree_office_walk") + video_path = os.path.join(data_path, "video") + odom_path = os.path.join(data_path, "odom") + + # Deploy modules + # Video replay module + video_module = dimos.deploy(VideoReplayModule, video_path) + video_module.video_out.transport = LCMTransport("/test_video", Image) + + # Odometry replay module (publishes to tf system directly) + odom_module = dimos.deploy(OdometryReplayModule, odom_path) + + # Spatial memory module + spatial_memory = dimos.deploy( + SpatialMemory, + collection_name="test_spatial_memory", + embedding_model="clip", + embedding_dimensions=512, + min_distance_threshold=0.5, # 0.5m for test + min_time_threshold=1.0, # 1 second + db_path=str(tmp_path / "chroma_db"), + visual_memory_path=str(tmp_path / "visual_memory.pkl"), + new_memory=True, + output_dir=str(tmp_path / "images"), + ) + + # Connect video stream + spatial_memory.color_image.connect(video_module.video_out) + + # Start all modules + video_module.start() + odom_module.start() + spatial_memory.start() + logger.info("All modules started, processing in background...") + + # Wait for frames to be processed with timeout + timeout = 10.0 # 10 second timeout + start_time = time.time() + + # Keep checking stats while modules are running + while (time.time() - start_time) < timeout: + stats = spatial_memory.get_stats() + if stats["frame_count"] > 0 and stats["stored_frame_count"] > 0: logger.info( - f"Final stats - Frame count: {final_stats['frame_count']}, Stored: {final_stats['stored_frame_count']}" + f"Frames processing - Frame count: {stats['frame_count']}, Stored: {stats['stored_frame_count']}" ) - - video_module.stop() - odom_module.stop() - spatial_memory.stop() - logger.info("Stopped all modules") - - logger.info("All spatial memory module tests passed!") - - finally: - # Cleanup - if "dimos" in locals(): - dimos.close_all() + break + await asyncio.sleep(0.5) + else: + # Timeout reached + stats = spatial_memory.get_stats() + logger.error( + f"Timeout after {timeout}s - Frame count: {stats['frame_count']}, Stored: {stats['stored_frame_count']}" + ) + raise AssertionError(f"No frames processed within {timeout} seconds") + + await asyncio.sleep(2) + + mid_stats = spatial_memory.get_stats() + logger.info( + f"Mid-test stats - Frame count: {mid_stats['frame_count']}, Stored: {mid_stats['stored_frame_count']}" + ) + assert mid_stats["frame_count"] >= stats["frame_count"], ( + "Frame count should increase or stay same" + ) + + # Test query while modules are still running + try: + text_results = spatial_memory.query_by_text("office") + logger.info(f"Query by text 'office' returned {len(text_results)} results") + assert len(text_results) > 0, "Should have at least one result" + except Exception as e: + logger.warning(f"Query by text failed: {e}") + + final_stats = spatial_memory.get_stats() + logger.info( + f"Final stats - Frame count: {final_stats['frame_count']}, Stored: {final_stats['stored_frame_count']}" + ) diff --git a/dimos/protocol/pubsub/shm/ipc_factory.py b/dimos/protocol/pubsub/shm/ipc_factory.py index 5f0b20165e..fbf98d379e 100644 --- a/dimos/protocol/pubsub/shm/ipc_factory.py +++ b/dimos/protocol/pubsub/shm/ipc_factory.py @@ -15,6 +15,7 @@ # frame_ipc.py # Python 3.9+ from abc import ABC, abstractmethod +from multiprocessing import resource_tracker from multiprocessing.shared_memory import SharedMemory import os import time @@ -24,6 +25,20 @@ _UNLINK_ON_GC = os.getenv("DIMOS_IPC_UNLINK_ON_GC", "0").lower() not in ("0", "false", "no") +def _unregister(shm: SharedMemory) -> SharedMemory: + """Remove a SharedMemory segment from the resource tracker. + + We manage lifecycle explicitly via close()/unlink(), so the resource + tracker must not attempt cleanup on process exit — that causes KeyError + spam when multiple processes share the same named segment. + """ + try: + resource_tracker.unregister(shm._name, "shared_memory") # type: ignore[attr-defined] + except Exception: + pass + return shm + + def _open_shm_with_retry(name: str) -> SharedMemory: tries = int(os.getenv("DIMOS_IPC_ATTACH_RETRIES", "40")) # ~40 tries base_ms = float(os.getenv("DIMOS_IPC_ATTACH_BACKOFF_MS", "5")) # 5 ms @@ -31,7 +46,7 @@ def _open_shm_with_retry(name: str) -> SharedMemory: last = None for i in range(tries): try: - return SharedMemory(name=name) + return _unregister(SharedMemory(name=name)) except FileNotFoundError as e: last = e # exponential backoff, capped @@ -39,11 +54,6 @@ def _open_shm_with_retry(name: str) -> SharedMemory: raise FileNotFoundError(f"SHM not found after {tries} retries: {name}") from last -def _sanitize_shm_name(name: str) -> str: - # Python's SharedMemory expects names like 'psm_abc', without leading '/' - return name.lstrip("/") if isinstance(name, str) else name - - # --------------------------- # 1) Abstract interface # --------------------------- @@ -100,7 +110,6 @@ def close(self) -> None: ... -from multiprocessing.shared_memory import SharedMemory import os import weakref @@ -108,7 +117,8 @@ def close(self) -> None: def _safe_unlink(name: str) -> None: try: shm = SharedMemory(name=name) - shm.unlink() + shm.unlink() # unlink() calls resource_tracker.unregister() + shm.close() except FileNotFoundError: pass except Exception: @@ -135,15 +145,18 @@ def __init__( # type: ignore[no-untyped-def] def _create_or_open(name: str, size: int): # type: ignore[no-untyped-def] try: + # Owner: leave registered because unlink() will unregister, and + # the tracker serves as safety net if the process crashes. shm = SharedMemory(create=True, size=size, name=name) owner = True except FileExistsError: - shm = SharedMemory(name=name) # attach existing + # Reader: unregister because we only close(), never unlink(). + shm = _unregister(SharedMemory(name=name)) owner = False return shm, owner if data_name is None or ctrl_name is None: - # fallback: random names (old behavior) + # Fallback: random names (old behavior) -> always owner self._shm_data = SharedMemory(create=True, size=2 * self._nbytes) self._shm_ctrl = SharedMemory(create=True, size=24) self._is_owner = True diff --git a/dimos/protocol/tf/test_tf.py b/dimos/protocol/tf/test_tf.py index bdbd808cbb..c1f0b13fa2 100644 --- a/dimos/protocol/tf/test_tf.py +++ b/dimos/protocol/tf/test_tf.py @@ -19,9 +19,8 @@ import pytest -from dimos.core import TF from dimos.msgs.geometry_msgs import PoseStamped, Quaternion, Transform, Vector3 -from dimos.protocol.tf import MultiTBuffer, TBuffer +from dimos.protocol.tf import TF, MultiTBuffer, TBuffer # from https://foxglove.dev/blog/understanding-ros-transforms diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index ed90982801..0fb09fb696 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -134,7 +134,6 @@ "temporal-memory": "dimos.perception.experimental.temporal_memory.temporal_memory", "twist-teleop-module": "dimos.teleop.quest.quest_extensions", "unitree-skills": "dimos.robot.unitree.unitree_skill_container", - "utilization": "dimos.utils.monitoring", "visualizing-teleop-module": "dimos.teleop.quest.quest_extensions", "vlm-agent": "dimos.agents.vlm_agent", "vlm-stream-tester": "dimos.agents.vlm_stream_tester", diff --git a/dimos/robot/cli/dimos.py b/dimos/robot/cli/dimos.py index 03e9a91349..22a9e234d9 100644 --- a/dimos/robot/cli/dimos.py +++ b/dimos/robot/cli/dimos.py @@ -20,6 +20,9 @@ import typer from dimos.core.global_config import GlobalConfig, global_config +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() main = typer.Typer( help="Dimensional CLI", @@ -101,6 +104,8 @@ def run( robot_types: list[str] = typer.Argument(..., help="Blueprints or modules to run"), ) -> None: """Start a robot blueprint""" + logger.info("Starting DimOS") + from dimos.core.blueprints import autoconnect from dimos.robot.get_all_blueprints import get_by_name from dimos.utils.logging_config import setup_exception_handler diff --git a/dimos/robot/drone/camera_module.py b/dimos/robot/drone/camera_module.py index 248b1ceb6e..63389aa358 100644 --- a/dimos/robot/drone/camera_module.py +++ b/dimos/robot/drone/camera_module.py @@ -23,7 +23,9 @@ from dimos_lcm.sensor_msgs import CameraInfo -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import PoseStamped from dimos.msgs.sensor_msgs import Image from dimos.msgs.std_msgs import Header diff --git a/dimos/robot/drone/connection_module.py b/dimos/robot/drone/connection_module.py index db5c4ca4cc..4e5559f220 100644 --- a/dimos/robot/drone/connection_module.py +++ b/dimos/robot/drone/connection_module.py @@ -25,7 +25,9 @@ from reactivex.disposable import CompositeDisposable, Disposable from dimos.agents.annotation import skill -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.mapping.types import LatLon from dimos.msgs.geometry_msgs import PoseStamped, Quaternion, Transform, Twist, Vector3 from dimos.msgs.sensor_msgs import Image diff --git a/dimos/robot/drone/drone.py b/dimos/robot/drone/drone.py index 6b9500804f..718e6b4aca 100644 --- a/dimos/robot/drone/drone.py +++ b/dimos/robot/drone/drone.py @@ -20,18 +20,19 @@ import functools import logging import os -from typing import Any +from typing import TYPE_CHECKING, Any from dimos_lcm.sensor_msgs import CameraInfo from dimos_lcm.std_msgs import String from reactivex import Observable -from dimos import core from dimos.agents.agent import agent from dimos.agents.skills.google_maps_skill_container import GoogleMapsSkillContainer from dimos.agents.skills.osm import OsmSkill from dimos.agents.web_human_input import web_input from dimos.core.blueprints import Blueprint, autoconnect +from dimos.core.module_coordinator import ModuleCoordinator +from dimos.core.transport import LCMTransport, pLCMTransport from dimos.mapping.types import LatLon from dimos.msgs.geometry_msgs import PoseStamped, Twist, Vector3 from dimos.msgs.sensor_msgs import Image @@ -47,6 +48,9 @@ from dimos.utils.logging_config import setup_logger from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule +if TYPE_CHECKING: + from dimos.core.rpc_client import ModuleProxy + logger = setup_logger() @@ -88,12 +92,12 @@ def __init__( RobotCapability.VISION, ] - self.dimos: core.DimosCluster | None = None - self.connection: DroneConnectionModule | None = None - self.camera: DroneCameraModule | None = None - self.tracking: DroneTrackingModule | None = None + self.dimos: ModuleCoordinator | None = None + self.connection: ModuleProxy | None = None + self.camera: ModuleProxy | None = None + self.tracking: ModuleProxy | None = None self.foxglove_bridge: FoxgloveBridge | None = None - self.websocket_vis: WebsocketVisModule | None = None + self.websocket_vis: ModuleProxy | None = None self._setup_directories() @@ -107,7 +111,8 @@ def start(self) -> None: logger.info("Starting Drone robot system...") # Start DimOS cluster - self.dimos = core.start(4) + self.dimos = ModuleCoordinator() + self.dimos.start() # Deploy modules self._deploy_connection() @@ -127,7 +132,7 @@ def _deploy_connection(self) -> None: assert self.dimos is not None logger.info("Deploying connection module...") - self.connection = self.dimos.deploy( # type: ignore[attr-defined] + self.connection = self.dimos.deploy( DroneConnectionModule, # connection_string="replay", connection_string=self.connection_string, @@ -136,19 +141,17 @@ def _deploy_connection(self) -> None: ) # Configure LCM transports - self.connection.odom.transport = core.LCMTransport("/drone/odom", PoseStamped) - self.connection.gps_location.transport = core.pLCMTransport("/gps_location") - self.connection.gps_goal.transport = core.pLCMTransport("/gps_goal") - self.connection.status.transport = core.LCMTransport("/drone/status", String) - self.connection.telemetry.transport = core.LCMTransport("/drone/telemetry", String) - self.connection.video.transport = core.LCMTransport("/drone/video", Image) - self.connection.follow_object_cmd.transport = core.LCMTransport( + self.connection.odom.transport = LCMTransport("/drone/odom", PoseStamped) + self.connection.gps_location.transport = pLCMTransport("/gps_location") + self.connection.gps_goal.transport = pLCMTransport("/gps_goal") + self.connection.status.transport = LCMTransport("/drone/status", String) + self.connection.telemetry.transport = LCMTransport("/drone/telemetry", String) + self.connection.video.transport = LCMTransport("/drone/video", Image) + self.connection.follow_object_cmd.transport = LCMTransport( "/drone/follow_object_cmd", String ) - self.connection.movecmd.transport = core.LCMTransport("/drone/cmd_vel", Vector3) - self.connection.movecmd_twist.transport = core.LCMTransport( - "/drone/tracking/cmd_vel", Twist - ) + self.connection.movecmd.transport = LCMTransport("/drone/cmd_vel", Vector3) + self.connection.movecmd_twist.transport = LCMTransport("/drone/tracking/cmd_vel", Twist) logger.info("Connection module deployed") @@ -163,9 +166,9 @@ def _deploy_camera(self) -> None: ) # Configure LCM transports - self.camera.color_image.transport = core.LCMTransport("/drone/color_image", Image) - self.camera.camera_info.transport = core.LCMTransport("/drone/camera_info", CameraInfo) - self.camera.camera_pose.transport = core.LCMTransport("/drone/camera_pose", PoseStamped) + self.camera.color_image.transport = LCMTransport("/drone/color_image", Image) + self.camera.camera_info.transport = LCMTransport("/drone/camera_info", CameraInfo) + self.camera.camera_pose.transport = LCMTransport("/drone/camera_pose", PoseStamped) # Connect video from connection module to camera module self.camera.video.connect(self.connection.video) @@ -183,13 +186,9 @@ def _deploy_tracking(self) -> None: outdoor=self.outdoor, ) - self.tracking.tracking_overlay.transport = core.LCMTransport( - "/drone/tracking_overlay", Image - ) - self.tracking.tracking_status.transport = core.LCMTransport( - "/drone/tracking_status", String - ) - self.tracking.cmd_vel.transport = core.LCMTransport("/drone/tracking/cmd_vel", Twist) + self.tracking.tracking_overlay.transport = LCMTransport("/drone/tracking_overlay", Image) + self.tracking.tracking_status.transport = LCMTransport("/drone/tracking_status", String) + self.tracking.cmd_vel.transport = LCMTransport("/drone/tracking/cmd_vel", Twist) self.tracking.video_input.connect(self.connection.video) self.tracking.follow_object_cmd.connect(self.connection.follow_object_cmd) @@ -204,11 +203,11 @@ def _deploy_visualization(self) -> None: assert self.dimos is not None assert self.connection is not None self.websocket_vis = self.dimos.deploy(WebsocketVisModule) # type: ignore[attr-defined] - # self.websocket_vis.click_goal.transport = core.LCMTransport("/goal_request", PoseStamped) - self.websocket_vis.gps_goal.transport = core.pLCMTransport("/gps_goal") - # self.websocket_vis.explore_cmd.transport = core.LCMTransport("/explore_cmd", Bool) - # self.websocket_vis.stop_explore_cmd.transport = core.LCMTransport("/stop_explore_cmd", Bool) - self.websocket_vis.cmd_vel.transport = core.LCMTransport("/cmd_vel", Twist) + # self.websocket_vis.click_goal.transport = LCMTransport("/goal_request", PoseStamped) + self.websocket_vis.gps_goal.transport = pLCMTransport("/gps_goal") + # self.websocket_vis.explore_cmd.transport = LCMTransport("/explore_cmd", Bool) + # self.websocket_vis.stop_explore_cmd.transport = LCMTransport("/stop_explore_cmd", Bool) + self.websocket_vis.cmd_vel.transport = LCMTransport("/cmd_vel", Twist) self.websocket_vis.odom.connect(self.connection.odom) self.websocket_vis.gps_location.connect(self.connection.gps_location) @@ -232,22 +231,9 @@ def _start_modules(self) -> None: assert self.foxglove_bridge is not None logger.info("Starting modules...") - # Start connection first - result = self.connection.start() - if not result: - logger.warning("Connection module failed to start (no drone connected?)") - - # Start camera - result = self.camera.start() - if not result: - logger.warning("Camera module failed to start") - - result = self.tracking.start() - if result: - logger.info("Tracking module started successfully") - else: - logger.warning("Tracking module failed to start") - + self.connection.start() + self.camera.start() + self.tracking.start() self.websocket_vis.start() # Start Foxglove @@ -377,9 +363,6 @@ def cleanup(self) -> None: self.stop() def stop(self) -> None: - """Stop the drone system.""" - logger.info("Stopping drone system...") - if self.connection: self.connection.stop() @@ -390,9 +373,7 @@ def stop(self) -> None: self.foxglove_bridge.stop() if self.dimos: - self.dimos.close_all() # type: ignore[attr-defined] - - logger.info("Drone system stopped") + self.dimos.stop() DRONE_SYSTEM_PROMPT = """\ diff --git a/dimos/robot/drone/drone_tracking_module.py b/dimos/robot/drone/drone_tracking_module.py index e1b633a05b..276b636633 100644 --- a/dimos/robot/drone/drone_tracking_module.py +++ b/dimos/robot/drone/drone_tracking_module.py @@ -25,7 +25,9 @@ import numpy as np from numpy.typing import NDArray -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.models.qwen.video_query import get_bbox_from_qwen_frame from dimos.msgs.geometry_msgs import Twist, Vector3 from dimos.msgs.sensor_msgs import Image, ImageFormat diff --git a/dimos/robot/drone/test_drone.py b/dimos/robot/drone/test_drone.py index 7381359f5a..0afa8bef35 100644 --- a/dimos/robot/drone/test_drone.py +++ b/dimos/robot/drone/test_drone.py @@ -428,9 +428,9 @@ def tearDown(self) -> None: self.pubsub_patch.stop() self.foxglove_patch.stop() - @patch("dimos.robot.drone.drone.core.start") + @patch("dimos.robot.drone.drone.ModuleCoordinator") @patch("dimos.utils.testing.TimedSensorReplay") - def test_full_system_with_replay(self, mock_replay, mock_core_start) -> None: + def test_full_system_with_replay(self, mock_replay, mock_coordinator_class) -> None: """Test full drone system initialization and operation with replay mode.""" # Set up mock replay data mavlink_messages = [ @@ -477,8 +477,8 @@ def replay_side_effect(store_name: str): mock_replay.side_effect = replay_side_effect - # Mock DimOS core - mock_core_start.return_value = self.mock_dimos + # Mock ModuleCoordinator + mock_coordinator_class.return_value = self.mock_dimos # Create drone in replay mode drone = Drone(connection_string="replay", video_port=5600) @@ -557,7 +557,7 @@ def deploy_side_effect(module_class, **kwargs): # Verify cleanup was called mock_connection.stop.assert_called_once() mock_camera.stop.assert_called_once() - self.mock_dimos.close_all.assert_called_once() + self.mock_dimos.stop.assert_called_once() class TestDroneControlCommands(unittest.TestCase): diff --git a/dimos/robot/foxglove_bridge.py b/dimos/robot/foxglove_bridge.py index 529a14c838..6700c38901 100644 --- a/dimos/robot/foxglove_bridge.py +++ b/dimos/robot/foxglove_bridge.py @@ -21,11 +21,14 @@ FoxgloveBridge as LCMFoxgloveBridge, ) -from dimos.core import DimosCluster, Module, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.module_coordinator import ModuleCoordinator from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from dimos.core.global_config import GlobalConfig + from dimos.core.rpc_client import ModuleProxy logging.getLogger("lcm_foxglove_bridge").setLevel(logging.ERROR) logging.getLogger("FoxgloveServer").setLevel(logging.ERROR) @@ -97,9 +100,9 @@ def stop(self) -> None: def deploy( - dimos: DimosCluster, + dimos: ModuleCoordinator, shm_channels: list[str] | None = None, -) -> FoxgloveBridge: +) -> "ModuleProxy": if shm_channels is None: shm_channels = [ "/image#sensor_msgs.Image", @@ -111,7 +114,7 @@ def deploy( shm_channels=shm_channels, ) foxglove_bridge.start() - return foxglove_bridge # type: ignore[no-any-return] + return foxglove_bridge foxglove_bridge = FoxgloveBridge.blueprint diff --git a/dimos/robot/unitree/b1/connection.py b/dimos/robot/unitree/b1/connection.py index bae4bc0844..4279f78399 100644 --- a/dimos/robot/unitree/b1/connection.py +++ b/dimos/robot/unitree/b1/connection.py @@ -24,7 +24,9 @@ from reactivex.disposable import Disposable -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import PoseStamped, Twist, TwistStamped from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.std_msgs import Int32 diff --git a/dimos/robot/unitree/b1/joystick_module.py b/dimos/robot/unitree/b1/joystick_module.py index bb07094973..0a72f81617 100644 --- a/dimos/robot/unitree/b1/joystick_module.py +++ b/dimos/robot/unitree/b1/joystick_module.py @@ -25,7 +25,9 @@ import time -from dimos.core import Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import Out from dimos.msgs.geometry_msgs import Twist, TwistStamped, Vector3 from dimos.msgs.std_msgs import Int32 diff --git a/dimos/robot/unitree/b1/unitree_b1.py b/dimos/robot/unitree/b1/unitree_b1.py index a2dd6c718d..2c0c918942 100644 --- a/dimos/robot/unitree/b1/unitree_b1.py +++ b/dimos/robot/unitree/b1/unitree_b1.py @@ -23,10 +23,9 @@ import logging import os -from dimos import core from dimos.core.module_coordinator import ModuleCoordinator from dimos.core.resource import Resource -from dimos.core.transport import ROSTransport +from dimos.core.transport import LCMTransport, ROSTransport from dimos.msgs.geometry_msgs import PoseStamped, TwistStamped from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.std_msgs import Int32 @@ -98,10 +97,10 @@ def start(self) -> None: self.connection = self._dimos.deploy(B1ConnectionModule, self.ip, self.port) # type: ignore[assignment] # Configure LCM transports for connection (matching G1 pattern) - self.connection.cmd_vel.transport = core.LCMTransport("/cmd_vel", TwistStamped) # type: ignore[attr-defined] - self.connection.mode_cmd.transport = core.LCMTransport("/b1/mode", Int32) # type: ignore[attr-defined] - self.connection.odom_in.transport = core.LCMTransport("/state_estimation", Odometry) # type: ignore[attr-defined] - self.connection.odom_pose.transport = core.LCMTransport("/odom", PoseStamped) # type: ignore[attr-defined] + self.connection.cmd_vel.transport = LCMTransport("/cmd_vel", TwistStamped) # type: ignore[attr-defined] + self.connection.mode_cmd.transport = LCMTransport("/b1/mode", Int32) # type: ignore[attr-defined] + self.connection.odom_in.transport = LCMTransport("/state_estimation", Odometry) # type: ignore[attr-defined] + self.connection.odom_pose.transport = LCMTransport("/odom", PoseStamped) # type: ignore[attr-defined] # Configure ROS transports for connection self.connection.ros_cmd_vel.transport = ROSTransport("/cmd_vel", TwistStamped) # type: ignore[attr-defined] @@ -113,8 +112,8 @@ def start(self) -> None: from dimos.robot.unitree.b1.joystick_module import JoystickModule self.joystick = self._dimos.deploy(JoystickModule) # type: ignore[assignment] - self.joystick.twist_out.transport = core.LCMTransport("/cmd_vel", TwistStamped) # type: ignore[attr-defined] - self.joystick.mode_out.transport = core.LCMTransport("/b1/mode", Int32) # type: ignore[attr-defined] + self.joystick.twist_out.transport = LCMTransport("/cmd_vel", TwistStamped) # type: ignore[attr-defined] + self.joystick.mode_out.transport = LCMTransport("/b1/mode", Int32) # type: ignore[attr-defined] logger.info("Joystick module deployed - pygame window will open") self._dimos.start_all_modules() diff --git a/dimos/robot/unitree/g1/blueprints/perceptive/_perception_and_memory.py b/dimos/robot/unitree/g1/blueprints/perceptive/_perception_and_memory.py index 47dc2588b9..241fcb32a8 100644 --- a/dimos/robot/unitree/g1/blueprints/perceptive/_perception_and_memory.py +++ b/dimos/robot/unitree/g1/blueprints/perceptive/_perception_and_memory.py @@ -18,12 +18,10 @@ from dimos.core.blueprints import autoconnect from dimos.perception.object_tracker import object_tracking from dimos.perception.spatial_perception import spatial_memory -from dimos.utils.monitoring import utilization _perception_and_memory = autoconnect( spatial_memory(), object_tracking(frame_id="camera_link"), - utilization(), ) __all__ = ["_perception_and_memory"] diff --git a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1.py b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1.py index 483928ec54..faea2ce0a8 100644 --- a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1.py +++ b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1.py @@ -24,6 +24,6 @@ unitree_g1 = autoconnect( unitree_g1_basic, _perception_and_memory, -).global_config(n_dask_workers=8) +).global_config(n_workers=8) __all__ = ["unitree_g1"] diff --git a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_detection.py b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_detection.py index 6e2da40a2c..25bff97c73 100644 --- a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_detection.py +++ b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_detection.py @@ -15,6 +15,8 @@ """G1 stack with person tracking and 3D detection.""" +from typing import Any + from dimos_lcm.foxglove_msgs import SceneUpdate from dimos_lcm.foxglove_msgs.ImageAnnotations import ImageAnnotations @@ -30,6 +32,11 @@ from dimos.perception.detection.person_tracker import PersonTracker, person_tracker_module from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_basic import unitree_g1_basic + +def _person_only(det: Any) -> bool: + return bool(det.class_id == 0) + + unitree_g1_detection = ( autoconnect( unitree_g1_basic, @@ -40,13 +47,13 @@ ), detection_db_module( camera_info=zed.CameraInfo.SingleWebcam, - filter=lambda det: det.class_id == 0, # Filter for person class only + filter=_person_only, # Filter for person class only ), person_tracker_module( cameraInfo=zed.CameraInfo.SingleWebcam, ), ) - .global_config(n_dask_workers=8) + .global_config(n_workers=8) .remappings( [ # Connect detection modules to camera and lidar diff --git a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_sim.py b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_sim.py index 059102c7a5..d69966455e 100644 --- a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_sim.py +++ b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_sim.py @@ -24,6 +24,6 @@ unitree_g1_sim = autoconnect( unitree_g1_basic_sim, _perception_and_memory, -).global_config(n_dask_workers=8) +).global_config(n_workers=8) __all__ = ["unitree_g1_sim"] diff --git a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py index 36bf569f72..0379abf4da 100644 --- a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py +++ b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py @@ -15,6 +15,8 @@ """Minimal G1 stack without navigation, used as a base for larger blueprints.""" +from typing import Any + from dimos_lcm.sensor_msgs import CameraInfo from dimos.core.blueprints import autoconnect @@ -33,30 +35,47 @@ from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.web.websocket_vis.websocket_vis_module import websocket_vis + +def _convert_camera_info(camera_info: Any) -> Any: + return camera_info.to_rerun( + image_topic="/world/color_image", + optical_frame="camera_optical", + ) + + +def _convert_global_map(grid: Any) -> Any: + return grid.to_rerun(voxel_size=0.1, mode="boxes") + + +def _convert_navigation_costmap(grid: Any) -> Any: + return grid.to_rerun( + colormap="Accent", + z_offset=0.015, + opacity=0.2, + background="#484981", + ) + + +def _static_base_link(rr: Any) -> list[Any]: + return [ + rr.Boxes3D( + half_sizes=[0.2, 0.15, 0.75], + colors=[(0, 255, 127)], + fill_mode="MajorWireframe", + ), + rr.Transform3D(parent_frame="tf#/base_link"), + ] + + rerun_config = { "pubsubs": [LCM(autoconf=True)], "visual_override": { - "world/camera_info": lambda camera_info: camera_info.to_rerun( - image_topic="/world/color_image", - optical_frame="camera_optical", - ), - "world/global_map": lambda grid: grid.to_rerun(voxel_size=0.1, mode="boxes"), - "world/navigation_costmap": lambda grid: grid.to_rerun( - colormap="Accent", - z_offset=0.015, - opacity=0.2, - background="#484981", - ), + "world/camera_info": _convert_camera_info, + "world/global_map": _convert_global_map, + "world/navigation_costmap": _convert_navigation_costmap, }, "static": { - "world/tf/base_link": lambda rr: [ - rr.Boxes3D( - half_sizes=[0.2, 0.15, 0.75], - colors=[(0, 255, 127)], - fill_mode="MajorWireframe", - ), - rr.Transform3D(parent_frame="tf#/base_link"), - ] + "world/tf/base_link": _static_base_link, }, } @@ -76,6 +95,16 @@ case _: _with_vis = autoconnect() + +def _create_webcam() -> Webcam: + return Webcam( + camera_index=0, + fps=15, + stereo_slice="left", + camera_info=zed.CameraInfo.SingleWebcam, + ) + + _camera = ( autoconnect( camera_module( @@ -85,12 +114,7 @@ frame_id="sensor", child_frame_id="camera_link", ), - hardware=lambda: Webcam( - camera_index=0, - fps=15, - stereo_slice="left", - camera_info=zed.CameraInfo.SingleWebcam, - ), + hardware=_create_webcam, ), ) if not global_config.simulation @@ -107,7 +131,7 @@ # Visualization websocket_vis(), ) - .global_config(n_dask_workers=4, robot_model="unitree_g1") + .global_config(n_workers=4, robot_model="unitree_g1") .transports( { # G1 uses Twist for movement commands diff --git a/dimos/robot/unitree/g1/connection.py b/dimos/robot/unitree/g1/connection.py index f12d0ee0e6..17a66945f9 100644 --- a/dimos/robot/unitree/g1/connection.py +++ b/dimos/robot/unitree/g1/connection.py @@ -13,17 +13,23 @@ # limitations under the License. -from typing import Any +from typing import TYPE_CHECKING, Any from reactivex.disposable import Disposable from dimos import spec -from dimos.core import DimosCluster, In, Module, rpc +from dimos.core.core import rpc from dimos.core.global_config import GlobalConfig, global_config +from dimos.core.module import Module +from dimos.core.module_coordinator import ModuleCoordinator +from dimos.core.stream import In from dimos.msgs.geometry_msgs import Twist from dimos.robot.unitree.connection import UnitreeWebRTCConnection from dimos.utils.logging_config import setup_logger +if TYPE_CHECKING: + from dimos.core.rpc_client import ModuleProxy + logger = setup_logger() @@ -92,11 +98,11 @@ def publish_request(self, topic: str, data: dict[str, Any]) -> dict[Any, Any]: g1_connection = G1Connection.blueprint -def deploy(dimos: DimosCluster, ip: str, local_planner: spec.LocalPlanner) -> G1Connection: +def deploy(dimos: ModuleCoordinator, ip: str, local_planner: spec.LocalPlanner) -> "ModuleProxy": connection = dimos.deploy(G1Connection, ip) # type: ignore[attr-defined] connection.cmd_vel.connect(local_planner.cmd_vel) connection.start() - return connection # type: ignore[no-any-return] + return connection __all__ = ["G1Connection", "deploy", "g1_connection"] diff --git a/dimos/robot/unitree/g1/sim.py b/dimos/robot/unitree/g1/sim.py index 6888ae74aa..f969bfbe04 100644 --- a/dimos/robot/unitree/g1/sim.py +++ b/dimos/robot/unitree/g1/sim.py @@ -20,8 +20,10 @@ from reactivex.disposable import Disposable -from dimos.core import In, Module, Out, rpc +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 import ( PoseStamped, Quaternion, 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 3135031738..2537e86632 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -15,6 +15,7 @@ # limitations under the License. import platform +from typing import Any from dimos.constants import DEFAULT_CAPACITY_COLOR_IMAGE from dimos.core.blueprints import autoconnect @@ -39,6 +40,38 @@ autoconnect() if platform.system() == "Linux" else autoconnect().transports(_mac_transports) ) + +def _convert_camera_info(camera_info: Any) -> Any: + return camera_info.to_rerun( + image_topic="/world/color_image", + optical_frame="camera_optical", + ) + + +def _convert_global_map(grid: Any) -> Any: + return grid.to_rerun(voxel_size=0.1, mode="boxes") + + +def _convert_navigation_costmap(grid: Any) -> Any: + return grid.to_rerun( + colormap="Accent", + z_offset=0.015, + opacity=0.2, + background="#484981", + ) + + +def _static_base_link(rr: Any) -> list[Any]: + return [ + rr.Boxes3D( + half_sizes=[0.35, 0.155, 0.2], + colors=[(0, 255, 127)], + fill_mode="wireframe", + ), + rr.Transform3D(parent_frame="tf#/base_link"), + ] + + rerun_config = { # any pubsub that supports subscribe_all and topic that supports str(topic) # is acceptable here @@ -49,28 +82,13 @@ # # This is unsustainable once we move to multi robot etc "visual_override": { - "world/camera_info": lambda camera_info: camera_info.to_rerun( - image_topic="/world/color_image", - optical_frame="camera_optical", - ), - "world/global_map": lambda grid: grid.to_rerun(voxel_size=0.1, mode="boxes"), - "world/navigation_costmap": lambda grid: grid.to_rerun( - colormap="Accent", - z_offset=0.015, - opacity=0.2, - background="#484981", - ), + "world/camera_info": _convert_camera_info, + "world/global_map": _convert_global_map, + "world/navigation_costmap": _convert_navigation_costmap, }, # slapping a go2 shaped box on top of tf/base_link "static": { - "world/tf/base_link": lambda rr: [ - rr.Boxes3D( - half_sizes=[0.35, 0.155, 0.2], - colors=[(0, 255, 127)], - fill_mode="wireframe", - ), - rr.Transform3D(parent_frame="tf#/base_link"), - ] + "world/tf/base_link": _static_base_link, }, } @@ -100,11 +118,10 @@ go2_connection(), websocket_vis(), ) - .global_config(n_dask_workers=4, robot_model="unitree_go2") + .global_config(n_workers=4, robot_model="unitree_go2") .configurators(ClockSyncConfigurator()) ) - __all__ = [ "unitree_go2_basic", ] diff --git a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py index 5d096444d5..c38e0eefa5 100644 --- a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py +++ b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py @@ -26,6 +26,6 @@ cost_mapper(), replanning_a_star_planner(), wavefront_frontier_explorer(), -).global_config(n_dask_workers=6, robot_model="unitree_go2") +).global_config(n_workers=6, robot_model="unitree_go2") __all__ = ["unitree_go2"] diff --git a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_spatial.py b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_spatial.py index e2695f9bfb..879bd828e4 100644 --- a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_spatial.py +++ b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_spatial.py @@ -16,12 +16,10 @@ from dimos.core.blueprints import autoconnect from dimos.perception.spatial_perception import spatial_memory from dimos.robot.unitree.go2.blueprints.smart.unitree_go2 import unitree_go2 -from dimos.utils.monitoring import utilization unitree_go2_spatial = autoconnect( unitree_go2, spatial_memory(), - utilization(), -).global_config(n_dask_workers=8) +).global_config(n_workers=8) __all__ = ["unitree_go2_spatial"] diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 82aa34c97a..4c7d4755fb 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -15,7 +15,7 @@ import logging from threading import Thread import time -from typing import Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol from reactivex.disposable import Disposable from reactivex.observable import Observable @@ -23,8 +23,15 @@ from dimos import spec from dimos.agents.annotation import skill -from dimos.core import DimosCluster, In, LCMTransport, Module, Out, pSHMTransport, rpc +from dimos.core.core import rpc from dimos.core.global_config import GlobalConfig, global_config +from dimos.core.module import Module +from dimos.core.module_coordinator import ModuleCoordinator +from dimos.core.stream import In, Out +from dimos.core.transport import LCMTransport, pSHMTransport + +if TYPE_CHECKING: + from dimos.core.rpc_client import ModuleProxy from dimos.msgs.geometry_msgs import ( PoseStamped, Quaternion, @@ -310,7 +317,7 @@ def observe(self) -> Image | None: go2_connection = GO2Connection.blueprint -def deploy(dimos: DimosCluster, ip: str, prefix: str = "") -> GO2Connection: +def deploy(dimos: ModuleCoordinator, ip: str, prefix: str = "") -> "ModuleProxy": from dimos.constants import DEFAULT_CAPACITY_COLOR_IMAGE connection = dimos.deploy(GO2Connection, ip) # type: ignore[attr-defined] @@ -327,7 +334,7 @@ def deploy(dimos: DimosCluster, ip: str, prefix: str = "") -> GO2Connection: connection.camera_info.transport = LCMTransport(f"{prefix}/camera_info", CameraInfo) connection.start() - return connection # type: ignore[no-any-return] + return connection __all__ = ["GO2Connection", "deploy", "go2_connection"] diff --git a/dimos/robot/unitree/keyboard_teleop.py b/dimos/robot/unitree/keyboard_teleop.py index 3d7d4c263e..14be8432e5 100644 --- a/dimos/robot/unitree/keyboard_teleop.py +++ b/dimos/robot/unitree/keyboard_teleop.py @@ -18,7 +18,9 @@ import pygame -from dimos.core import Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import Out from dimos.msgs.geometry_msgs import Twist, Vector3 # Force X11 driver to avoid OpenGL threading issues diff --git a/dimos/robot/unitree/modular/detect.py b/dimos/robot/unitree/modular/detect.py index e5999e9fd8..99faddc946 100644 --- a/dimos/robot/unitree/modular/detect.py +++ b/dimos/robot/unitree/modular/detect.py @@ -112,7 +112,7 @@ def broadcast( # type: ignore[no-untyped-def] ImageAnnotations, ) - from dimos.core import LCMTransport + from dimos.core.transport import LCMTransport from dimos.msgs.geometry_msgs import PoseStamped lidar_transport = LCMTransport("/lidar", PointCloud2) # type: ignore[var-annotated] diff --git a/dimos/robot/unitree/rosnav.py b/dimos/robot/unitree/rosnav.py index 7a9b98b678..adc97eb4a2 100644 --- a/dimos/robot/unitree/rosnav.py +++ b/dimos/robot/unitree/rosnav.py @@ -16,7 +16,9 @@ import logging import time -from dimos.core import In, Module, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import PoseStamped from dimos.msgs.sensor_msgs import Joy from dimos.msgs.std_msgs.Bool import Bool diff --git a/dimos/robot/unitree/testing/test_actors.py b/dimos/robot/unitree/testing/test_actors.py index 0fee2175fc..ed0b05d664 100644 --- a/dimos/robot/unitree/testing/test_actors.py +++ b/dimos/robot/unitree/testing/test_actors.py @@ -13,24 +13,25 @@ # limitations under the License. import asyncio from collections.abc import Callable -import time import pytest -from dimos import core -from dimos.core import Module, rpc +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.module_coordinator import ModuleCoordinator +from dimos.core.transport import LCMTransport from dimos.msgs.sensor_msgs import PointCloud2 from dimos.robot.unitree.type.map import Map as Mapper @pytest.fixture def dimos(): - return core.start(2) - - -@pytest.fixture -def client(): - return core.start(2) + ret = ModuleCoordinator() + ret.start() + try: + yield ret + finally: + ret.stop() class Consumer: @@ -62,20 +63,6 @@ def addten(self, x: int): return x + 10 -@pytest.mark.tool -def test_wait(client) -> None: - counter = client.submit(Counter, actor=True).result() - - async def addten(n): - return await counter.addten(n) - - consumer = client.submit(Consumer, counter=addten, actor=True).result() - - print("waitcall1", consumer.waitcall(2).result()) - print("waitcall2", consumer.waitcall(2).result()) - time.sleep(1) - - @pytest.mark.tool def test_basic(dimos) -> None: counter = dimos.deploy(Counter) @@ -98,7 +85,7 @@ def test_basic(dimos) -> None: @pytest.mark.tool def test_mapper_start(dimos) -> None: mapper = dimos.deploy(Mapper) - mapper.lidar.transport = core.LCMTransport("/lidar", PointCloud2) + mapper.lidar.transport = LCMTransport("/lidar", PointCloud2) print("start res", mapper.start().result()) diff --git a/dimos/robot/unitree/type/map.py b/dimos/robot/unitree/type/map.py index a771467246..95b2bf6f6b 100644 --- a/dimos/robot/unitree/type/map.py +++ b/dimos/robot/unitree/type/map.py @@ -20,8 +20,12 @@ from reactivex import interval from reactivex.disposable import Disposable -from dimos.core import DimosCluster, In, LCMTransport, Module, Out, rpc +from dimos.core.core import rpc from dimos.core.global_config import GlobalConfig, global_config +from dimos.core.module import Module +from dimos.core.module_coordinator import ModuleCoordinator +from dimos.core.stream import In, Out +from dimos.core.transport import LCMTransport from dimos.mapping.pointclouds.accumulators.general import GeneralPointCloudAccumulator from dimos.mapping.pointclouds.accumulators.protocol import PointCloudAccumulator from dimos.mapping.pointclouds.occupancy import general_occupancy @@ -116,7 +120,7 @@ def _publish(self, _: Any) -> None: mapper = Map.blueprint -def deploy(dimos: DimosCluster, connection: Go2ConnectionProtocol): # type: ignore[no-untyped-def] +def deploy(dimos: ModuleCoordinator, connection: Go2ConnectionProtocol): # type: ignore[no-untyped-def] mapper = dimos.deploy(Map, global_publish_interval=1.0) # type: ignore[attr-defined] mapper.global_map.transport = LCMTransport("/global_map", PointCloud2) mapper.global_costmap.transport = LCMTransport("/global_costmap", OccupancyGrid) diff --git a/dimos/simulation/manipulators/sim_module.py b/dimos/simulation/manipulators/sim_module.py index 4f1bb986d3..831ea6ee34 100644 --- a/dimos/simulation/manipulators/sim_module.py +++ b/dimos/simulation/manipulators/sim_module.py @@ -14,25 +14,22 @@ """Simulator-agnostic manipulator simulation module.""" -from __future__ import annotations - +from collections.abc import Callable from dataclasses import dataclass +from pathlib import Path import threading import time -from typing import TYPE_CHECKING, Any +from typing import Any from reactivex.disposable import Disposable -from dimos.core import In, Module, Out, rpc -from dimos.core.module import ModuleConfig +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out from dimos.msgs.sensor_msgs import JointCommand, JointState, RobotState from dimos.simulation.engines import EngineType, get_engine from dimos.simulation.manipulators.sim_manip_interface import SimManipInterface -if TYPE_CHECKING: - from collections.abc import Callable - from pathlib import Path - @dataclass(kw_only=True) class SimulationModuleConfig(ModuleConfig): diff --git a/dimos/simulation/mujoco/mujoco_process.py b/dimos/simulation/mujoco/mujoco_process.py index 8529de976b..27217afadd 100755 --- a/dimos/simulation/mujoco/mujoco_process.py +++ b/dimos/simulation/mujoco/mujoco_process.py @@ -16,6 +16,7 @@ import base64 import json +import os import pickle import signal import sys @@ -231,7 +232,8 @@ def _run_simulation(config: GlobalConfig, shm: ShmReader) -> None: if __name__ == "__main__": def signal_handler(_signum: int, _frame: Any) -> None: - sys.exit(0) + # os._exit is the documented way of exiting a child process immediatly. + os._exit(0) signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) diff --git a/dimos/simulation/mujoco/shared_memory.py b/dimos/simulation/mujoco/shared_memory.py index 70ba50af2b..bd96ad2025 100644 --- a/dimos/simulation/mujoco/shared_memory.py +++ b/dimos/simulation/mujoco/shared_memory.py @@ -83,12 +83,7 @@ def from_names(cls, shm_names: dict[str, str]) -> "ShmSet": @classmethod def from_sizes(cls) -> "ShmSet": - return cls( - **{ - k: _unregister(SharedMemory(create=True, size=_shm_sizes[k])) - for k in _shm_sizes.keys() - } - ) + return cls(**{k: SharedMemory(create=True, size=_shm_sizes[k]) for k in _shm_sizes.keys()}) def to_names(self) -> dict[str, str]: return {k: getattr(self, k).name for k in _shm_sizes.keys()} diff --git a/dimos/spec/control.py b/dimos/spec/control.py index e2024c5a09..48d58a926a 100644 --- a/dimos/spec/control.py +++ b/dimos/spec/control.py @@ -14,7 +14,7 @@ from typing import Protocol -from dimos.core import Out +from dimos.core.stream import Out from dimos.msgs.geometry_msgs import Twist diff --git a/dimos/spec/mapping.py b/dimos/spec/mapping.py index f8e7e1a04f..0ba88cfaa9 100644 --- a/dimos/spec/mapping.py +++ b/dimos/spec/mapping.py @@ -14,7 +14,7 @@ from typing import Protocol -from dimos.core import Out +from dimos.core.stream import Out from dimos.msgs.nav_msgs import OccupancyGrid from dimos.msgs.sensor_msgs import PointCloud2 diff --git a/dimos/spec/nav.py b/dimos/spec/nav.py index d1f62c0846..08f6f42b35 100644 --- a/dimos/spec/nav.py +++ b/dimos/spec/nav.py @@ -14,7 +14,7 @@ from typing import Protocol -from dimos.core import In, Out +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import PoseStamped, Twist from dimos.msgs.nav_msgs import Path @@ -23,9 +23,4 @@ class Nav(Protocol): goal_req: In[PoseStamped] goal_active: Out[PoseStamped] path_active: Out[Path] - ctrl: Out[Twist] - - # identity quaternion (Quaternion(0,0,0,1)) represents "no rotation requested" - def navigate_to_target(self, target: PoseStamped) -> None: ... - - def stop_navigating(self) -> None: ... + cmd_vel: Out[Twist] diff --git a/dimos/spec/perception.py b/dimos/spec/perception.py index 1cecdb4d2f..1cfe352390 100644 --- a/dimos/spec/perception.py +++ b/dimos/spec/perception.py @@ -14,7 +14,7 @@ from typing import Protocol -from dimos.core import Out +from dimos.core.stream import Out from dimos.msgs.nav_msgs.Odometry import Odometry as OdometryMsg from dimos.msgs.sensor_msgs import CameraInfo, Image as ImageMsg, Imu, PointCloud2 @@ -25,7 +25,6 @@ class Image(Protocol): class Camera(Image): camera_info: Out[CameraInfo] - _camera_info: CameraInfo class DepthCamera(Camera): diff --git a/dimos/teleop/keyboard/keyboard_teleop_module.py b/dimos/teleop/keyboard/keyboard_teleop_module.py index ff42ce9a1a..cc3c301804 100644 --- a/dimos/teleop/keyboard/keyboard_teleop_module.py +++ b/dimos/teleop/keyboard/keyboard_teleop_module.py @@ -42,8 +42,9 @@ pygame = None # type: ignore[assignment] from dimos.control.examples.cartesian_ik_jogger import JogState -from dimos.core import Module, Out, rpc -from dimos.core.module import ModuleConfig +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import Out from dimos.msgs.geometry_msgs import PoseStamped # Force X11 driver to avoid OpenGL threading issues diff --git a/dimos/teleop/phone/phone_extensions.py b/dimos/teleop/phone/phone_extensions.py index f0a8fd4d01..0f52fce2e0 100644 --- a/dimos/teleop/phone/phone_extensions.py +++ b/dimos/teleop/phone/phone_extensions.py @@ -19,7 +19,7 @@ - SimplePhoneTeleop: Filters to ground robot axes and outputs cmd_vel: Out[Twist] """ -from dimos.core import Out +from dimos.core.stream import Out from dimos.msgs.geometry_msgs import Twist, TwistStamped, Vector3 from dimos.teleop.phone.phone_teleop_module import PhoneTeleopModule diff --git a/dimos/teleop/phone/phone_teleop_module.py b/dimos/teleop/phone/phone_teleop_module.py index c0da85c27c..ec33df58eb 100644 --- a/dimos/teleop/phone/phone_teleop_module.py +++ b/dimos/teleop/phone/phone_teleop_module.py @@ -34,8 +34,9 @@ from reactivex.disposable import Disposable -from dimos.core import In, Module, Out, rpc -from dimos.core.module import ModuleConfig +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import Twist, TwistStamped, Vector3 from dimos.msgs.std_msgs.Bool import Bool from dimos.utils.logging_config import setup_logger diff --git a/dimos/teleop/quest/quest_extensions.py b/dimos/teleop/quest/quest_extensions.py index b4e38de546..9139b671ea 100644 --- a/dimos/teleop/quest/quest_extensions.py +++ b/dimos/teleop/quest/quest_extensions.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 # Copyright 2025-2026 Dimensional Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -21,11 +20,10 @@ - VisualizingTeleopModule: Adds Rerun visualization (inherits press-and-hold engage) """ -from __future__ import annotations - from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any +from typing import Any +from dimos.core.stream import Out from dimos.msgs.geometry_msgs import PoseStamped, TwistStamped from dimos.teleop.quest.quest_teleop_module import Hand, QuestTeleopConfig, QuestTeleopModule from dimos.teleop.utils.teleop_visualization import ( @@ -33,9 +31,6 @@ visualize_pose, ) -if TYPE_CHECKING: - from dimos.core import Out - @dataclass class TwistTeleopConfig(QuestTeleopConfig): diff --git a/dimos/teleop/quest/quest_teleop_module.py b/dimos/teleop/quest/quest_teleop_module.py index ea77bb5fc0..adfdb6aaff 100644 --- a/dimos/teleop/quest/quest_teleop_module.py +++ b/dimos/teleop/quest/quest_teleop_module.py @@ -32,8 +32,9 @@ from reactivex.disposable import Disposable -from dimos.core import In, Module, Out, rpc -from dimos.core.module import ModuleConfig +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import PoseStamped from dimos.msgs.sensor_msgs import Joy from dimos.teleop.quest.quest_types import Buttons, QuestControllerState diff --git a/dimos/utils/actor_registry.py b/dimos/utils/actor_registry.py deleted file mode 100644 index 6f6d219594..0000000000 --- a/dimos/utils/actor_registry.py +++ /dev/null @@ -1,84 +0,0 @@ -# Copyright 2025-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. - -"""Shared memory registry for tracking actor deployments across processes.""" - -import json -from multiprocessing import shared_memory - - -class ActorRegistry: - """Shared memory registry of actor deployments.""" - - SHM_NAME = "dimos_actor_registry" - SHM_SIZE = 65536 # 64KB should be enough for most deployments - - @staticmethod - def update(actor_name: str, worker_id: str) -> None: - """Update registry with new actor deployment.""" - try: - shm = shared_memory.SharedMemory(name=ActorRegistry.SHM_NAME) - except FileNotFoundError: - shm = shared_memory.SharedMemory( - name=ActorRegistry.SHM_NAME, create=True, size=ActorRegistry.SHM_SIZE - ) - - # Read existing data - data = ActorRegistry._read_from_shm(shm) - - # Update with new actor - data[actor_name] = worker_id - - # Write back - ActorRegistry._write_to_shm(shm, data) - shm.close() - - @staticmethod - def get_all() -> dict[str, str]: - """Get all actor->worker mappings.""" - try: - shm = shared_memory.SharedMemory(name=ActorRegistry.SHM_NAME) - data = ActorRegistry._read_from_shm(shm) - shm.close() - return data - except FileNotFoundError: - return {} - - @staticmethod - def clear() -> None: - """Clear the registry and free shared memory.""" - try: - shm = shared_memory.SharedMemory(name=ActorRegistry.SHM_NAME) - ActorRegistry._write_to_shm(shm, {}) - shm.close() - shm.unlink() - except FileNotFoundError: - pass - - @staticmethod - def _read_from_shm(shm) -> dict[str, str]: # type: ignore[no-untyped-def] - """Read JSON data from shared memory.""" - raw = bytes(shm.buf[:]).rstrip(b"\x00") - if not raw: - return {} - return json.loads(raw.decode("utf-8")) # type: ignore[no-any-return] - - @staticmethod - def _write_to_shm(shm, data: dict[str, str]): # type: ignore[no-untyped-def] - """Write JSON data to shared memory.""" - json_bytes = json.dumps(data).encode("utf-8") - if len(json_bytes) > ActorRegistry.SHM_SIZE: - raise ValueError("Registry data too large for shared memory") - shm.buf[: len(json_bytes)] = json_bytes - shm.buf[len(json_bytes) :] = b"\x00" * (ActorRegistry.SHM_SIZE - len(json_bytes)) diff --git a/dimos/utils/cli/human/humancli.py b/dimos/utils/cli/human/humancli.py index cf7fd8a258..ee3ba5aa37 100644 --- a/dimos/utils/cli/human/humancli.py +++ b/dimos/utils/cli/human/humancli.py @@ -29,7 +29,7 @@ from textual.geometry import Size from textual.widgets import Input, RichLog -from dimos.core import pLCMTransport +from dimos.core.transport import pLCMTransport from dimos.utils.cli import theme from dimos.utils.generic import truncate_display_string diff --git a/dimos/core/colors.py b/dimos/utils/colors.py similarity index 100% rename from dimos/core/colors.py rename to dimos/utils/colors.py diff --git a/dimos/utils/logging_config.py b/dimos/utils/logging_config.py index a9bfc5031d..1ddac0a23c 100644 --- a/dimos/utils/logging_config.py +++ b/dimos/utils/logging_config.py @@ -14,7 +14,6 @@ from collections.abc import Mapping from datetime import datetime -import inspect import logging import logging.handlers import os @@ -202,8 +201,7 @@ def setup_logger(*, level: int | None = None) -> Any: A configured structlog logger instance. """ - caller_frame = inspect.stack()[1] - name = caller_frame.filename + name = sys._getframe(1).f_code.co_filename # Convert absolute path to relative path try: diff --git a/dimos/utils/metrics.py b/dimos/utils/metrics.py index bf7bf45cdc..3292d7220f 100644 --- a/dimos/utils/metrics.py +++ b/dimos/utils/metrics.py @@ -20,7 +20,8 @@ from dimos_lcm.std_msgs import Float32 import rerun as rr -from dimos.core import LCMTransport, Transport +from dimos.core.stream import Transport +from dimos.core.transport import LCMTransport F = TypeVar("F", bound=Callable[..., Any]) diff --git a/dimos/utils/monitoring.py b/dimos/utils/monitoring.py deleted file mode 100644 index ca3e03c55e..0000000000 --- a/dimos/utils/monitoring.py +++ /dev/null @@ -1,307 +0,0 @@ -# Copyright 2025-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. - -""" -Note, to enable ps-spy to run without sudo you need: - - echo 0 | sudo tee /proc/sys/kernel/yama/ptrace_scope -""" - -from functools import cache -import os -import re -import shutil -import subprocess -import threading - -from distributed import get_client -from distributed.client import Client - -from dimos.core import Module, rpc -from dimos.utils.actor_registry import ActorRegistry -from dimos.utils.logging_config import setup_logger - -logger = setup_logger() - - -def print_data_table(data) -> None: # type: ignore[no-untyped-def] - headers = [ - "cpu_percent", - "active_percent", - "gil_percent", - "n_threads", - "pid", - "worker_id", - "modules", - ] - numeric_headers = {"cpu_percent", "active_percent", "gil_percent", "n_threads", "pid"} - - # Add registered modules. - modules = ActorRegistry.get_all() - for worker in data: - worker["modules"] = ", ".join( - module_name.split("-", 1)[0] - for module_name, worker_id_str in modules.items() - if worker_id_str == str(worker["worker_id"]) - ) - - # Determine column widths - col_widths = [] - for h in headers: - max_len = max(len(str(d[h])) for d in data) - col_widths.append(max(len(h), max_len)) - - # Print header with DOS box characters - header_row = " │ ".join(h.ljust(col_widths[i]) for i, h in enumerate(headers)) - border_parts = ["─" * w for w in col_widths] - border_line = "─┼─".join(border_parts) - print(border_line) - print(header_row) - print(border_line) - - # Print rows - for row in data: - formatted_cells = [] - for i, h in enumerate(headers): - value = str(row[h]) - if h in numeric_headers: - formatted_cells.append(value.rjust(col_widths[i])) - else: - formatted_cells.append(value.ljust(col_widths[i])) - print(" │ ".join(formatted_cells)) - - -class UtilizationThread(threading.Thread): - _module: "UtilizationModule" - _stop_event: threading.Event - _monitors: dict # type: ignore[type-arg] - - def __init__(self, module) -> None: # type: ignore[no-untyped-def] - super().__init__(daemon=True) - self._module = module - self._stop_event = threading.Event() - self._monitors = {} - - def run(self) -> None: - while not self._stop_event.is_set(): - workers = self._module.client.scheduler_info()["workers"] # type: ignore[union-attr] - pids = {pid: None for pid in get_worker_pids()} # type: ignore[no-untyped-call] - for worker, info in workers.items(): - pid = get_pid_by_port(worker.rsplit(":", 1)[-1]) - if pid is None: - continue - pids[pid] = info["id"] - data = [] - for pid, worker_id in pids.items(): - if pid not in self._monitors: - self._monitors[pid] = GilMonitorThread(pid) - self._monitors[pid].start() - cpu, gil, active, n_threads = self._monitors[pid].get_values() - data.append( - { - "cpu_percent": cpu, - "worker_id": worker_id, - "pid": pid, - "gil_percent": gil, - "active_percent": active, - "n_threads": n_threads, - } - ) - data.sort(key=lambda x: x["pid"]) - self._fix_missing_ids(data) - print_data_table(data) - self._stop_event.wait(1) - - def stop(self) -> None: - self._stop_event.set() - for monitor in self._monitors.values(): - monitor.stop() - monitor.join(timeout=2) - - def _fix_missing_ids(self, data) -> None: # type: ignore[no-untyped-def] - """ - Some worker IDs are None. But if we order the workers by PID and all - non-None ids are in order, then we can deduce that the None ones are the - missing indices. - """ - if all(x["worker_id"] in (i, None) for i, x in enumerate(data)): - for i, worker in enumerate(data): - worker["worker_id"] = i - - -class UtilizationModule(Module): - client: Client | None - _utilization_thread: UtilizationThread | None - - def __init__(self) -> None: - super().__init__() - self.client = None - self._utilization_thread = None - - if not os.getenv("MEASURE_GIL_UTILIZATION"): - logger.info("Set `MEASURE_GIL_UTILIZATION=true` to print GIL utilization.") - return - - if not _can_use_py_spy(): # type: ignore[no-untyped-call] - logger.warning( - "Cannot start UtilizationModule because in order to run py-spy without " - "being root you need to enable this:\n" - "\n" - " echo 0 | sudo tee /proc/sys/kernel/yama/ptrace_scope" - ) - return - - if not shutil.which("py-spy"): - logger.warning("Cannot start UtilizationModule because `py-spy` is not installed.") - return - - self.client = get_client() - self._utilization_thread = UtilizationThread(self) - - @rpc - def start(self) -> None: - super().start() - - if self._utilization_thread: - self._utilization_thread.start() - - @rpc - def stop(self) -> None: - if self._utilization_thread: - self._utilization_thread.stop() - self._utilization_thread.join(timeout=2) - super().stop() - - -utilization = UtilizationModule.blueprint - - -__all__ = ["UtilizationModule", "utilization"] - - -def _can_use_py_spy(): # type: ignore[no-untyped-def] - try: - with open("/proc/sys/kernel/yama/ptrace_scope") as f: - value = f.read().strip() - return value == "0" - except Exception: - pass - return False - - -@cache -def get_pid_by_port(port: int) -> int | None: - try: - result = subprocess.run( - ["lsof", "-ti", f":{port}"], capture_output=True, text=True, check=True - ) - pid_str = result.stdout.strip() - return int(pid_str) if pid_str else None - except subprocess.CalledProcessError: - return None - - -def get_worker_pids(): # type: ignore[no-untyped-def] - pids = [] - for pid in os.listdir("/proc"): - if not pid.isdigit(): - continue - try: - with open(f"/proc/{pid}/cmdline") as f: - cmdline = f.read().replace("\x00", " ") - if "spawn_main" in cmdline: - pids.append(int(pid)) - except (FileNotFoundError, PermissionError): - continue - return pids - - -class GilMonitorThread(threading.Thread): - pid: int - _latest_values: tuple[float, float, float, int] - _stop_event: threading.Event - _lock: threading.Lock - - def __init__(self, pid: int) -> None: - super().__init__(daemon=True) - self.pid = pid - self._latest_values = (-1.0, -1.0, -1.0, -1) - self._stop_event = threading.Event() - self._lock = threading.Lock() - - def run(self): # type: ignore[no-untyped-def] - command = ["py-spy", "top", "--pid", str(self.pid), "--rate", "100"] - process = None - try: - process = subprocess.Popen( - command, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - bufsize=1, # Line-buffered output - ) - - for line in iter(process.stdout.readline, ""): # type: ignore[union-attr] - if self._stop_event.is_set(): - break - - if "GIL:" not in line: - continue - - match = re.search( - r"GIL:\s*([\d.]+?)%,\s*Active:\s*([\d.]+?)%,\s*Threads:\s*(\d+)", line - ) - if not match: - continue - - try: - cpu_percent = _get_cpu_percent(self.pid) - gil_percent = float(match.group(1)) - active_percent = float(match.group(2)) - num_threads = int(match.group(3)) - - with self._lock: - self._latest_values = ( - cpu_percent, - gil_percent, - active_percent, - num_threads, - ) - except (ValueError, IndexError): - pass - except Exception as e: - logger.error(f"An error occurred in GilMonitorThread for PID {self.pid}: {e}") - raise - finally: - if process: - process.terminate() - process.wait(timeout=1) - self._stop_event.set() - - def get_values(self): # type: ignore[no-untyped-def] - with self._lock: - return self._latest_values - - def stop(self) -> None: - self._stop_event.set() - - -def _get_cpu_percent(pid: int) -> float: - try: - result = subprocess.run( - ["ps", "-p", str(pid), "-o", "%cpu="], capture_output=True, text=True, check=True - ) - return float(result.stdout.strip()) - except Exception: - return -1.0 diff --git a/dimos/utils/testing/moment.py b/dimos/utils/testing/moment.py index e92d771687..a0aa4219cb 100644 --- a/dimos/utils/testing/moment.py +++ b/dimos/utils/testing/moment.py @@ -21,7 +21,7 @@ from dimos.utils.testing.replay import TimedSensorReplay if TYPE_CHECKING: - from dimos.core import Transport + from dimos.core.stream import Transport T = TypeVar("T", bound=Timestamped) diff --git a/dimos/utils/testing/test_moment.py b/dimos/utils/testing/test_moment.py index 6764610d0e..75f11d2657 100644 --- a/dimos/utils/testing/test_moment.py +++ b/dimos/utils/testing/test_moment.py @@ -13,7 +13,7 @@ # limitations under the License. import time -from dimos.core import LCMTransport +from dimos.core.transport import LCMTransport from dimos.msgs.geometry_msgs import PoseStamped, Transform from dimos.msgs.sensor_msgs import CameraInfo, Image, PointCloud2 from dimos.protocol.tf import TF diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 1dc104f1b4..af91f1b8b8 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -33,8 +33,8 @@ from toolz import pipe # type: ignore[import-untyped] import typer -from dimos.core import Module, rpc -from dimos.core.module import ModuleConfig +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.protocol.pubsub.patterns import Glob, pattern_matches from dimos.utils.logging_config import setup_logger diff --git a/dimos/web/dimos_interface/api/server.py b/dimos/web/dimos_interface/api/server.py index 6692e90f46..7f7bdc23b8 100644 --- a/dimos/web/dimos_interface/api/server.py +++ b/dimos/web/dimos_interface/api/server.py @@ -65,7 +65,6 @@ def __init__( # type: ignore[no-untyped-def] audio_subject=None, **streams, ) -> None: - print("Starting FastAPIServer initialization...") # Debug print super().__init__(dev_name, edge_type) self.app = FastAPI() self._server: uvicorn.Server | None = None @@ -119,9 +118,7 @@ def __init__( # type: ignore[no-untyped-def] self.text_disposables[key] = disposable self.disposables.add(disposable) - print("Setting up routes...") # Debug print self.setup_routes() - print("FastAPIServer initialization complete") # Debug print def process_frame_fastapi(self, frame): # type: ignore[no-untyped-def] """Convert frame to JPEG format for streaming.""" diff --git a/dimos/web/websocket_vis/README.md b/dimos/web/websocket_vis/README.md index c04235958e..3f0776bed5 100644 --- a/dimos/web/websocket_vis/README.md +++ b/dimos/web/websocket_vis/README.md @@ -39,17 +39,17 @@ Control: ```python from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule -from dimos import core +from dimos.core.transport import LCMTransport, pLCMTransport # Deploy the WebSocket visualization module websocket_vis = dimos.deploy(WebsocketVisModule, port=7779) # Receive control from the Foxglove plugin. -websocket_vis.click_goal.transport = core.LCMTransport("/goal_request", PoseStamped) -websocket_vis.explore_cmd.transport = core.LCMTransport("/explore_cmd", Bool) -websocket_vis.stop_explore_cmd.transport = core.LCMTransport("/stop_explore_cmd", Bool) -websocket_vis.movecmd.transport = core.LCMTransport("/cmd_vel", Twist) -websocket_vis.gps_goal.transport = core.pLCMTransport("/gps_goal") +websocket_vis.click_goal.transport = LCMTransport("/goal_request", PoseStamped) +websocket_vis.explore_cmd.transport = LCMTransport("/explore_cmd", Bool) +websocket_vis.stop_explore_cmd.transport = LCMTransport("/stop_explore_cmd", Bool) +websocket_vis.movecmd.transport = LCMTransport("/cmd_vel", Twist) +websocket_vis.gps_goal.transport = pLCMTransport("/gps_goal") # Send visualization data to the Foxglove plugin. websocket_vis.robot_pose.connect(connection.odom) diff --git a/dimos/web/websocket_vis/websocket_vis_module.py b/dimos/web/websocket_vis/websocket_vis_module.py index ad93af5c96..2c3ad3009b 100644 --- a/dimos/web/websocket_vis/websocket_vis_module.py +++ b/dimos/web/websocket_vis/websocket_vis_module.py @@ -45,8 +45,10 @@ FilePath(__file__).parent.parent / "command-center-extension" / "dist-standalone" ) -from dimos.core import In, Module, Out, rpc +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.mapping.occupancy.gradient import gradient from dimos.mapping.occupancy.inflation import simple_inflate from dimos.mapping.types import LatLon @@ -194,6 +196,10 @@ def start(self) -> None: @rpc def stop(self) -> None: + if getattr(self, "_ws_stopped", False): + return + self._ws_stopped = True + if self._uvicorn_server: self._uvicorn_server.should_exit = True diff --git a/docs/capabilities/navigation/native/index.md b/docs/capabilities/navigation/native/index.md index 115c6f0ee2..a750d3bfba 100644 --- a/docs/capabilities/navigation/native/index.md +++ b/docs/capabilities/navigation/native/index.md @@ -135,7 +135,7 @@ unitree_go2 = autoconnect( cost_mapper(), # 2D costmap generation replanning_a_star_planner(), # path planning wavefront_frontier_explorer(), # exploration -).global_config(n_dask_workers=6, robot_model="unitree_go2") +).global_config(n_workers=6, robot_model="unitree_go2") to_svg(unitree_go2, "assets/go2_blueprint.svg") ``` diff --git a/docs/development/dimos_run.md b/docs/development/dimos_run.md index 604724a68d..b12f5bc3fb 100644 --- a/docs/development/dimos_run.md +++ b/docs/development/dimos_run.md @@ -59,7 +59,7 @@ class GlobalConfig(BaseSettings): robot_ip: str | None = None simulation: bool = False replay: bool = False - n_dask_workers: int = 2 + n_workers: int = 2 ``` Configuration values can be set from multiple places in order of precedence (later entries override earlier ones): diff --git a/docs/development/profiling_dimos.md b/docs/development/profiling_dimos.md new file mode 100644 index 0000000000..2ff3082299 --- /dev/null +++ b/docs/development/profiling_dimos.md @@ -0,0 +1,9 @@ +# Profiling dimos + +You can use py-spy to profile a particular blueprint: + +```bash +uv run py-spy record --format speedscope --subprocesses -o profile.speedscope.json -- python -m dimos.robot.cli.dimos run unitree-go2-agentic +``` + +Hit `Ctrl+C` when you're done. It will write a `profile.speedscope.json` file which you can upload to [speedscope.app](https://www.speedscope.app/) to visualize it. diff --git a/docs/usage/blueprints.md b/docs/usage/blueprints.md index 54b52ba3c0..ed48670cb4 100644 --- a/docs/usage/blueprints.md +++ b/docs/usage/blueprints.md @@ -8,7 +8,8 @@ You create a `Blueprint` from a single module (say `ConnectionModule`) with: ```python session=blueprint-ex1 from dimos.core.blueprints import Blueprint -from dimos.core import Module, rpc +from dimos.core.core import rpc +from dimos.core.module import Module class ConnectionModule(Module): def __init__(self, arg1, arg2, kwarg='value') -> None: @@ -100,7 +101,9 @@ Imagine you have this code: from functools import partial from dimos.core.blueprints import Blueprint, autoconnect -from dimos.core import Module, rpc, Out, In +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import Out, In from dimos.msgs.sensor_msgs import Image class ModuleA(Module): @@ -161,7 +164,9 @@ Sometimes you need to rename a connection to match what other modules expect. Yo ```python session=blueprint-ex2 from dimos.core.blueprints import autoconnect -from dimos.core import Module, rpc, Out, In +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import Out, In from dimos.msgs.sensor_msgs import Image class ConnectionModule(Module): @@ -204,7 +209,8 @@ blueprint.remappings([ Each module can optionally take global config as a `cfg` option in `__init__`. E.g.: ```python session=blueprint-ex3 -from dimos.core import Module, rpc +from dimos.core.core import rpc +from dimos.core.module import Module from dimos.core.global_config import GlobalConfig class ModuleA(Module): @@ -217,7 +223,7 @@ class ModuleA(Module): The config is normally taken from .env or from environment variables. But you can specifically override the values for a specific blueprint: ```python session=blueprint-ex3 -blueprint = ModuleA.blueprint().global_config(n_dask_workers=8) +blueprint = ModuleA.blueprint().global_config(n_workers=8) ``` ## Calling the methods of other modules @@ -225,7 +231,8 @@ blueprint = ModuleA.blueprint().global_config(n_dask_workers=8) Imagine you have this code: ```python session=blueprint-ex3 -from dimos.core import Module, rpc +from dimos.core.core import rpc +from dimos.core.module import Module class Drone(Module): @@ -243,7 +250,8 @@ And you want to call `ModuleA.get_time` in `ModuleB.request_the_time`. To do this, you can request a module reference. ```python session=blueprint-ex3 -from dimos.core import Module, rpc +from dimos.core.core import rpc +from dimos.core.module import Module class HelperModule(Module): drone_module: Drone @@ -283,7 +291,8 @@ class ModuleB(Module): Skills are methods on a `Module` decorated with `@skill`. The agent automatically discovers all skills from launched modules at startup. ```python session=blueprint-ex4 -from dimos.core import Module, rpc +from dimos.core.core import rpc +from dimos.core.module import Module from dimos.agents.annotation import skill from dimos.core.global_config import GlobalConfig diff --git a/docs/usage/configuration.md b/docs/usage/configuration.md index eaad4a9271..fe6e0029f0 100644 --- a/docs/usage/configuration.md +++ b/docs/usage/configuration.md @@ -49,7 +49,9 @@ Error: Config.__init__() got an unexpected keyword argument 'something' ```python from dataclasses import dataclass -from dimos.core import In, Module, Out, rpc, ModuleConfig +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out from rich import print @dataclass diff --git a/docs/usage/data_streams/advanced_streams.md b/docs/usage/data_streams/advanced_streams.md index e9d9f1d12d..68c1ed5bfa 100644 --- a/docs/usage/data_streams/advanced_streams.md +++ b/docs/usage/data_streams/advanced_streams.md @@ -106,7 +106,8 @@ The `LATEST` strategy means: when the slow subscriber finishes processing, it ge Most module streams offer backpressured observables. ```python session=bp -from dimos.core import Module, In +from dimos.core.module import Module +from dimos.core.stream import In from dimos.msgs.sensor_msgs import Image class MLModel(Module): diff --git a/docs/usage/data_streams/reactivex.md b/docs/usage/data_streams/reactivex.md index 45873b471b..026eb292c4 100644 --- a/docs/usage/data_streams/reactivex.md +++ b/docs/usage/data_streams/reactivex.md @@ -285,7 +285,7 @@ disposed ```python session=rx import time -from dimos.core import Module +from dimos.core.module import Module class MyModule(Module): def start(self): diff --git a/docs/usage/native_modules.md b/docs/usage/native_modules.md index 5a9362839f..929ac18424 100644 --- a/docs/usage/native_modules.md +++ b/docs/usage/native_modules.md @@ -18,7 +18,8 @@ Both the config dataclass and pubsub topics get converted to CLI args passed dow ```python no-result session=nativemodule from dataclasses import dataclass -from dimos.core import Out, LCMTransport +from dimos.core.stream import Out +from dimos.core.transport import LCMTransport from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.msgs.sensor_msgs.Imu import Imu @@ -229,7 +230,7 @@ For language interop examples (subscribing to DimOS topics from C++, TypeScript, The Livox Mid-360 LiDAR driver is a complete example at [`dimos/hardware/sensors/lidar/livox/module.py`](/dimos/hardware/sensors/lidar/livox/module.py): ```python skip -from dimos.core import Out +from dimos.core.stream import Out from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.msgs.sensor_msgs.Imu import Imu diff --git a/docs/usage/sensor_streams/advanced_streams.md b/docs/usage/sensor_streams/advanced_streams.md index c2cd0dbfca..588a7928ac 100644 --- a/docs/usage/sensor_streams/advanced_streams.md +++ b/docs/usage/sensor_streams/advanced_streams.md @@ -106,7 +106,8 @@ The `LATEST` strategy means: when the slow subscriber finishes processing, it ge Most module streams offer backpressured observables. ```python session=bp -from dimos.core import Module, In +from dimos.core.module import Module +from dimos.core.stream import In from dimos.msgs.sensor_msgs import Image class MLModel(Module): diff --git a/docs/usage/sensor_streams/reactivex.md b/docs/usage/sensor_streams/reactivex.md index 45873b471b..026eb292c4 100644 --- a/docs/usage/sensor_streams/reactivex.md +++ b/docs/usage/sensor_streams/reactivex.md @@ -285,7 +285,7 @@ disposed ```python session=rx import time -from dimos.core import Module +from dimos.core.module import Module class MyModule(Module): def start(self): diff --git a/docs/usage/transforms.md b/docs/usage/transforms.md index 2435839feb..8b98e4e81d 100644 --- a/docs/usage/transforms.md +++ b/docs/usage/transforms.md @@ -172,7 +172,7 @@ Modules in DimOS automatically get a `frame_id` property. This is controlled by - `frame_id_prefix` - Optional prefix for namespacing ```python -from dimos.core import Module, ModuleConfig +from dimos.core.module import Module, ModuleConfig from dataclasses import dataclass @dataclass @@ -221,8 +221,10 @@ This example demonstrates how multiple modules publish and receive transforms. T import time import reactivex as rx from reactivex import operators as ops -from dimos.core import Module, rpc, start +from dimos.core.core import rpc +from dimos.core.module import Module from dimos.msgs.geometry_msgs import Quaternion, Transform, Vector3 +from dimos.core.module_coordinator import ModuleCoordinator class RobotBaseModule(Module): """Publishes the robot's position in the world frame at 10Hz.""" @@ -306,16 +308,14 @@ class PerceptionModule(Module): if __name__ == "__main__": - dimos = start(3) + dimos = ModuleCoordinator() + dimos.start() - # Deploy and start modules robot = dimos.deploy(RobotBaseModule) camera = dimos.deploy(CameraModule) perception = dimos.deploy(PerceptionModule) - robot.start() - camera.start() - perception.start() + dimos.start_all_modules() time.sleep(1.0) diff --git a/docs/usage/transports/index.md b/docs/usage/transports/index.md index 1c8745d117..5cfe9caaa8 100644 --- a/docs/usage/transports/index.md +++ b/docs/usage/transports/index.md @@ -92,7 +92,7 @@ nav = autoconnect( cost_mapper(), replanning_a_star_planner(), wavefront_frontier_explorer(), -).global_config(n_dask_workers=6, robot_model="unitree_go2") +).global_config(n_workers=6, robot_model="unitree_go2") ros = nav.transports( { @@ -113,10 +113,12 @@ Each **stream** on a module can use a different transport. Set `.transport` on t ```python ansi=false import time -from dimos.core import In, Module, start +from dimos.core.module import Module +from dimos.core.stream import In from dimos.core.transport import LCMTransport from dimos.hardware.sensors.camera.module import CameraModule from dimos.msgs.sensor_msgs import Image +from dimos.core.module_coordinator import ModuleCoordinator class ImageListener(Module): @@ -129,7 +131,8 @@ class ImageListener(Module): if __name__ == "__main__": # Start local cluster and deploy modules to separate processes - dimos = start(2) + dimos = ModuleCoordinator() + dimos.start() camera = dimos.deploy(CameraModule, frequency=2.0) listener = dimos.deploy(ImageListener) @@ -140,8 +143,7 @@ if __name__ == "__main__": # Connect listener input to camera output listener.image.connect(camera.color_image) - camera.start() - listener.start() + dimos.start_all_modules() time.sleep(2) dimos.stop() diff --git a/examples/camera_grayscale.py b/examples/camera_grayscale.py index 0d21a0449f..52fb7166b6 100644 --- a/examples/camera_grayscale.py +++ b/examples/camera_grayscale.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.core import In, Module, Out, autoconnect, rpc +from dimos.core.blueprints import autoconnect +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out from dimos.hardware.sensors.camera.module import CameraModule from dimos.msgs.sensor_msgs.Image import Image from dimos.visualization.rerun.bridge import RerunBridgeModule diff --git a/examples/rpc_calls.py b/examples/rpc_calls.py index c39ff8d4e8..ea3df37217 100644 --- a/examples/rpc_calls.py +++ b/examples/rpc_calls.py @@ -14,8 +14,9 @@ from typing import Protocol -from dimos.core import Module, rpc from dimos.core.blueprints import autoconnect +from dimos.core.core import rpc +from dimos.core.module import Module from dimos.spec.utils import Spec diff --git a/examples/simplerobot/README.md b/examples/simplerobot/README.md index 8c7b6dfbc7..8d2f9432b3 100644 --- a/examples/simplerobot/README.md +++ b/examples/simplerobot/README.md @@ -33,7 +33,7 @@ Use `lcmspy` in another terminal to inspect messages. Press `q` or `Esc` to quit From any language with LCM bindings, publish `Twist` messages to `/cmd_vel`: ```python -from dimos.core import LCMTransport +from dimos.core.transport import LCMTransport from dimos.msgs.geometry_msgs import Twist transport = LCMTransport("/cmd_vel", Twist) diff --git a/examples/simplerobot/simplerobot.py b/examples/simplerobot/simplerobot.py index b959fa7d6f..010b3bf2eb 100644 --- a/examples/simplerobot/simplerobot.py +++ b/examples/simplerobot/simplerobot.py @@ -29,7 +29,9 @@ import reactivex as rx -from dimos.core import In, Module, ModuleConfig, Out, rpc +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs import Pose, PoseStamped, Quaternion, Twist, Vector3 @@ -102,22 +104,13 @@ def _update(self) -> None: if __name__ == "__main__": import argparse - from dimos.core import LCMTransport + from dimos.core.transport import LCMTransport parser = argparse.ArgumentParser(description="Simple virtual robot") parser.add_argument("--headless", action="store_true") parser.add_argument("--selftest", action="store_true", help="Run demo movements") args = parser.parse_args() - # If running in a dimos cluster we'd call - # - # from dimos.core import start - # dimos = start() - # robot = dimos.deploy(SimpleRobot) - # - # but this is a standalone example - # and we don't mind running in the main thread - robot = SimpleRobot() robot.pose.transport = LCMTransport("/odom", PoseStamped) robot.cmd_vel.transport = LCMTransport("/cmd_vel", Twist) diff --git a/pyproject.toml b/pyproject.toml index b34bee992c..84ba5ea25c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,10 +42,8 @@ dependencies = [ "python-dotenv", "annotation-protocol>=1.4.0", "lazy_loader", - - # Multiprocess - "dask[complete]==2025.5.1", "plum-dispatch==2.5.7", + # Logging "structlog>=25.5.0,<26", "colorlog==6.9.0", @@ -66,6 +64,7 @@ dependencies = [ # remove this once rerun is optional in core "rerun-sdk>=0.20.0", "toolz>=1.1.0", + "protobuf>=6.33.5,<7", ] @@ -283,7 +282,6 @@ docker = [ "numpy>=1.26.4", "scipy>=1.15.1", "reactivex", - "dask[distributed]==2025.5.1", "plum-dispatch==2.5.7", "structlog>=25.5.0,<26", "pydantic", diff --git a/uv.lock b/uv.lock index aa32378bd0..bd25dd1d10 100644 --- a/uv.lock +++ b/uv.lock @@ -454,30 +454,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/cb/f2ad4230dc2eb1a74edf38f1a38b9b52277f75bef262d8908e60d957e13c/blinker-1.9.0-py3-none-any.whl", hash = "sha256:ba0efaa9080b619ff2f3459d1d500c57bddea4a6b424b60a91141db6fd2f08bc", size = 8458, upload-time = "2024-11-08T17:25:46.184Z" }, ] -[[package]] -name = "bokeh" -version = "3.8.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "contourpy", version = "1.3.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "contourpy", version = "1.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "jinja2" }, - { name = "narwhals" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "packaging" }, - { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "pandas", version = "3.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "pillow" }, - { name = "pyyaml" }, - { name = "tornado", marker = "sys_platform != 'emscripten'" }, - { name = "xyzservices" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e4/31/7ee0c4dfd0255631b0624ce01be178704f91f763f02a1879368eb109befd/bokeh-3.8.2.tar.gz", hash = "sha256:8e7dcacc21d53905581b54328ad2705954f72f2997f99fc332c1de8da53aa3cc", size = 6529251, upload-time = "2026-01-06T00:20:06.568Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f6/a8/877f306720bc114c612579c5af36bcb359026b83d051226945499b306b1a/bokeh-3.8.2-py3-none-any.whl", hash = "sha256:5e2c0d84f75acb25d60efb9e4d2f434a791c4639b47d685534194c4e07bd0111", size = 7207131, upload-time = "2026-01-06T00:20:04.917Z" }, -] - [[package]] name = "brax" version = "0.14.1" @@ -1639,41 +1615,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0b/8c/dd63d210b28a7589f4bc1e84880525368147425c717d12834ab562f52d14/dash-4.0.0-py3-none-any.whl", hash = "sha256:e36b4b4eae9e1fa4136bf4f1450ed14ef76063bc5da0b10f8ab07bd57a7cb1ab", size = 7247521, upload-time = "2026-02-03T19:42:25.01Z" }, ] -[[package]] -name = "dask" -version = "2025.5.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "cloudpickle" }, - { name = "fsspec" }, - { name = "importlib-metadata", marker = "python_full_version < '3.12'" }, - { name = "packaging" }, - { name = "partd" }, - { name = "pyyaml" }, - { name = "toolz" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/3d/29/05feb8e2531c46d763547c66b7f5deb39b53d99b3be1b4ddddbd1cec6567/dask-2025.5.1.tar.gz", hash = "sha256:979d9536549de0e463f4cab8a8c66c3a2ef55791cd740d07d9bf58fab1d1076a", size = 10969324, upload-time = "2025-05-20T19:54:30.688Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/38/30/53b0844a7a4c6b041b111b24ca15cc9b8661a86fe1f6aaeb2d0d7f0fb1f2/dask-2025.5.1-py3-none-any.whl", hash = "sha256:3b85fdaa5f6f989dde49da6008415b1ae996985ebdfb1e40de2c997d9010371d", size = 1474226, upload-time = "2025-05-20T19:54:20.309Z" }, -] - -[package.optional-dependencies] -complete = [ - { name = "bokeh" }, - { name = "distributed" }, - { name = "jinja2" }, - { name = "lz4" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "pandas", version = "3.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "pyarrow" }, -] -distributed = [ - { name = "distributed" }, -] - [[package]] name = "dataclasses-json" version = "0.6.7" @@ -1741,7 +1682,6 @@ source = { editable = "." } dependencies = [ { name = "annotation-protocol" }, { name = "colorlog" }, - { name = "dask", extra = ["complete"] }, { name = "dimos-lcm" }, { name = "lazy-loader" }, { name = "llvmlite" }, @@ -1754,6 +1694,7 @@ dependencies = [ { name = "pin" }, { name = "plotext" }, { name = "plum-dispatch" }, + { name = "protobuf" }, { name = "pydantic" }, { name = "pydantic-settings" }, { name = "python-dotenv" }, @@ -1905,7 +1846,6 @@ dev = [ { name = "watchdog" }, ] docker = [ - { name = "dask", extra = ["distributed"] }, { name = "dimos-lcm" }, { name = "lcm" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -2041,8 +1981,6 @@ requires-dist = [ { name = "ctransformers", extras = ["cuda"], marker = "extra == 'cuda'", specifier = "==0.2.27" }, { name = "cupy-cuda12x", marker = "platform_machine == 'x86_64' and extra == 'cuda'", specifier = "==13.6.0" }, { name = "cyclonedds", marker = "extra == 'dds'", specifier = ">=0.10.5" }, - { name = "dask", extras = ["complete"], specifier = "==2025.5.1" }, - { name = "dask", extras = ["distributed"], marker = "extra == 'docker'", specifier = "==2025.5.1" }, { name = "dimos", extras = ["agents", "web", "perception", "visualization", "sim"], marker = "extra == 'base'" }, { name = "dimos", extras = ["base"], marker = "extra == 'unitree'" }, { name = "dimos", extras = ["dev"], marker = "extra == 'dds'" }, @@ -2109,6 +2047,7 @@ requires-dist = [ { name = "plum-dispatch", marker = "extra == 'docker'", specifier = "==2.5.7" }, { name = "portal", marker = "extra == 'misc'" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = "==4.2.0" }, + { name = "protobuf", specifier = ">=6.33.5,<7" }, { name = "psycopg2-binary", marker = "extra == 'psql'", specifier = ">=2.9.11" }, { name = "py-spy", marker = "extra == 'dev'" }, { name = "pydantic" }, @@ -2213,32 +2152,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl", hash = "sha256:9659f7d87e46584a30b5780e43ac7a2143098441670ff0a49d5f9034c54a6c16", size = 469047, upload-time = "2025-07-17T16:51:58.613Z" }, ] -[[package]] -name = "distributed" -version = "2025.5.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "cloudpickle" }, - { name = "dask" }, - { name = "jinja2" }, - { name = "locket" }, - { name = "msgpack" }, - { name = "packaging" }, - { name = "psutil" }, - { name = "pyyaml" }, - { name = "sortedcontainers" }, - { name = "tblib" }, - { name = "toolz" }, - { name = "tornado" }, - { name = "urllib3" }, - { name = "zict" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/29/ba/45950f405d023a520a4d10753ef40209a465b86c8fdc131236ec29bcb15c/distributed-2025.5.1.tar.gz", hash = "sha256:cf1d62a2c17a0a9fc1544bd10bb7afd39f22f24aaa9e3df3209c44d2cfb16703", size = 1107874, upload-time = "2025-05-20T19:54:26.005Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/25/65/89601dcc7383f0e5109e59eab90677daa9abb260d821570cd6089c8894bf/distributed-2025.5.1-py3-none-any.whl", hash = "sha256:74782b965ddb24ce59c6441fa777e944b5962d82325cc41f228537b59bb7fbbe", size = 1014789, upload-time = "2025-05-20T19:54:21.935Z" }, -] - [[package]] name = "distro" version = "1.9.0" @@ -4423,15 +4336,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/30/a8/e61a8c2b3cc7a597073d9cde1fcbb567e9d827f1db30c93cf80422eac70d/llvmlite-0.46.0-cp314-cp314-win_amd64.whl", hash = "sha256:7821eda3ec1f18050f981819756631d60b6d7ab1a6cf806d9efefbe3f4082d61", size = 39153056, upload-time = "2025-12-08T18:15:33.938Z" }, ] -[[package]] -name = "locket" -version = "1.0.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/2f/83/97b29fe05cb6ae28d2dbd30b81e2e402a3eed5f460c26e9eaa5895ceacf5/locket-1.0.0.tar.gz", hash = "sha256:5c0d4c052a8bbbf750e056a8e65ccd309086f4f0f18a2eac306a8dfa4112a632", size = 4350, upload-time = "2022-04-20T22:04:44.312Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/db/bc/83e112abc66cd466c6b83f99118035867cecd41802f8d044638aa78a106e/locket-1.0.0-py2.py3-none-any.whl", hash = "sha256:b6c819a722f7b6bd955b80781788e4a66a55628b858d347536b7e81325a3a5e3", size = 4398, upload-time = "2022-04-20T22:04:42.23Z" }, -] - [[package]] name = "logistro" version = "2.0.1" @@ -6487,15 +6391,14 @@ version = "2.3.3" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version < '3.11' and sys_platform == 'darwin'", - "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", "python_full_version < '3.11' and sys_platform == 'win32'", "(python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "python-dateutil", marker = "python_full_version < '3.11'" }, - { name = "pytz", marker = "python_full_version < '3.11'" }, - { name = "tzdata", marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'aarch64') or (python_full_version < '3.11' and sys_platform != 'linux')" }, + { name = "python-dateutil", marker = "(python_full_version < '3.11' and platform_machine != 'aarch64') or (python_full_version < '3.11' and sys_platform != 'linux')" }, + { name = "pytz", marker = "(python_full_version < '3.11' and platform_machine != 'aarch64') or (python_full_version < '3.11' and sys_platform != 'linux')" }, + { name = "tzdata", marker = "(python_full_version < '3.11' and platform_machine != 'aarch64') or (python_full_version < '3.11' and sys_platform != 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/33/01/d40b85317f86cf08d853a4f495195c73815fdf205eef3993821720274518/pandas-2.3.3.tar.gz", hash = "sha256:e05e1af93b977f7eafa636d043f9f94c7ee3ac81af99c13508215942e64c993b", size = 4495223, upload-time = "2025-09-29T23:34:51.853Z" } wheels = [ @@ -6556,9 +6459,6 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'darwin'", "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version >= '3.14' and platform_machine == 'aarch64' and sys_platform == 'linux'", - "python_full_version == '3.13.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", - "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'win32'", "(python_full_version >= '3.14' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')", @@ -6566,13 +6466,12 @@ resolution-markers = [ "python_full_version == '3.12.*' and sys_platform == 'win32'", "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')", "python_full_version == '3.11.*' and sys_platform == 'darwin'", - "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", "python_full_version == '3.11.*' and sys_platform == 'win32'", "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')", ] dependencies = [ - { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "python-dateutil", marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'aarch64') or (python_full_version >= '3.11' and sys_platform != 'linux')" }, + { name = "python-dateutil", marker = "(python_full_version >= '3.11' and platform_machine != 'aarch64') or (python_full_version >= '3.11' and sys_platform != 'linux')" }, { name = "tzdata", marker = "(python_full_version >= '3.11' and sys_platform == 'emscripten') or (python_full_version >= '3.11' and sys_platform == 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/de/da/b1dc0481ab8d55d0f46e343cfe67d4551a0e14fcee52bd38ca1bd73258d8/pandas-3.0.0.tar.gz", hash = "sha256:0facf7e87d38f721f0af46fe70d97373a37701b1c09f7ed7aeeb292ade5c050f", size = 4633005, upload-time = "2026-01-21T15:52:04.726Z" } @@ -6649,19 +6548,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b6/61/fae042894f4296ec49e3f193aff5d7c18440da9e48102c3315e1bc4519a7/parso-0.8.6-py2.py3-none-any.whl", hash = "sha256:2c549f800b70a5c4952197248825584cb00f033b29c692671d3bf08bf380baff", size = 106894, upload-time = "2026-02-09T15:45:21.391Z" }, ] -[[package]] -name = "partd" -version = "1.4.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "locket" }, - { name = "toolz" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/b2/3a/3f06f34820a31257ddcabdfafc2672c5816be79c7e353b02c1f318daa7d4/partd-1.4.2.tar.gz", hash = "sha256:d022c33afbdc8405c226621b015e8067888173d85f7f5ecebb3cafed9a20f02c", size = 21029, upload-time = "2024-05-06T19:51:41.945Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/71/e7/40fb618334dcdf7c5a316c0e7343c5cd82d3d866edc100d98e29bc945ecd/partd-1.4.2-py3-none-any.whl", hash = "sha256:978e4ac767ec4ba5b86c6eaa52e5a2a3bc748a2ca839e8cc798f1cc6ce6efb0f", size = 18905, upload-time = "2024-05-06T19:51:39.271Z" }, -] - [[package]] name = "pathspec" version = "1.0.4" @@ -9239,15 +9125,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, ] -[[package]] -name = "tblib" -version = "3.2.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f4/8a/14c15ae154895cc131174f858c707790d416c444fc69f93918adfd8c4c0b/tblib-3.2.2.tar.gz", hash = "sha256:e9a652692d91bf4f743d4a15bc174c0b76afc750fe8c7b6d195cc1c1d6d2ccec", size = 35046, upload-time = "2025-11-12T12:21:16.572Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/02/be/5d2d47b1fb58943194fb59dcf222f7c4e35122ec0ffe8c36e18b5d728f0b/tblib-3.2.2-py3-none-any.whl", hash = "sha256:26bdccf339bcce6a88b2b5432c988b266ebbe63a4e593f6b578b1d2e723d2b76", size = 12893, upload-time = "2025-11-12T12:21:14.407Z" }, -] - [[package]] name = "tenacity" version = "9.1.4" @@ -10892,15 +10769,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7b/d9/8d95e906764a386a3d3b596f3c68bb63687dfca806373509f51ce8eea81f/xxhash-3.6.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:15e0dac10eb9309508bfc41f7f9deaa7755c69e35af835db9cb10751adebc35d", size = 31565, upload-time = "2025-10-02T14:37:06.966Z" }, ] -[[package]] -name = "xyzservices" -version = "2025.11.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ee/0f/022795fc1201e7c29e742a509913badb53ce0b38f64b6db859e2f6339da9/xyzservices-2025.11.0.tar.gz", hash = "sha256:2fc72b49502b25023fd71e8f532fb4beddbbf0aa124d90ea25dba44f545e17ce", size = 1135703, upload-time = "2025-11-22T11:31:51.82Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ef/5c/2c189d18d495dd0fa3f27ccc60762bbc787eed95b9b0147266e72bb76585/xyzservices-2025.11.0-py3-none-any.whl", hash = "sha256:de66a7599a8d6dad63980b77defd1d8f5a5a9cb5fc8774ea1c6e89ca7c2a3d2f", size = 93916, upload-time = "2025-11-22T11:31:50.525Z" }, -] - [[package]] name = "yapf" version = "0.40.2" @@ -10915,15 +10783,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/66/c9/d4b03b2490107f13ebd68fe9496d41ae41a7de6275ead56d0d4621b11ffd/yapf-0.40.2-py3-none-any.whl", hash = "sha256:adc8b5dd02c0143108878c499284205adb258aad6db6634e5b869e7ee2bd548b", size = 254707, upload-time = "2023-09-22T18:40:43.297Z" }, ] -[[package]] -name = "zict" -version = "3.0.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d1/ac/3c494dd7ec5122cff8252c1a209b282c0867af029f805ae9befd73ae37eb/zict-3.0.0.tar.gz", hash = "sha256:e321e263b6a97aafc0790c3cfb3c04656b7066e6738c37fffcca95d803c9fba5", size = 33238, upload-time = "2023-04-17T21:41:16.041Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/80/ab/11a76c1e2126084fde2639514f24e6111b789b0bfa4fc6264a8975c7e1f1/zict-3.0.0-py2.py3-none-any.whl", hash = "sha256:5796e36bd0e0cc8cf0fbc1ace6a68912611c1dbd74750a3f3026b9b9d6a327ae", size = 43332, upload-time = "2023-04-17T21:41:13.444Z" }, -] - [[package]] name = "zipp" version = "3.23.0"