diff --git a/dimos/core/coordination/test_worker.py b/dimos/core/coordination/test_worker.py index c606dcda73..a9708fc845 100644 --- a/dimos/core/coordination/test_worker.py +++ b/dimos/core/coordination/test_worker.py @@ -14,14 +14,16 @@ from typing import TYPE_CHECKING +from pydantic import Field import pytest -from dimos.core.coordination.worker_manager_python import WorkerManagerPython +from dimos.core.coordination.worker_manager_python import WorkerManagerPython, _merge_config_kwargs 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 import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.protocol.service.spec import BaseConfig if TYPE_CHECKING: from dimos.core.resource_monitor.stats import WorkerStats @@ -81,6 +83,15 @@ def get_multiplier(self) -> int: return self.multiplier +class NestedConfig(BaseConfig): + a: int = 0 + b: int = 2 + + +class ModuleWithNestedConfig(ModuleConfig): + sub: NestedConfig = Field(default_factory=lambda: NestedConfig(a=10, b=20)) + + class HeavyModule(Module): dedicated_worker = True @@ -189,6 +200,17 @@ def test_worker_manager_parallel_deployment(create_worker_manager): module3.stop() +def test_nested_blueprint_config_defaults_survive_cli_override(): + blueprint_kwargs = {"sub": NestedConfig(a=10, b=20)} + cli_kwargs = {"sub": {"a": "1"}} + + merged = _merge_config_kwargs(blueprint_kwargs, cli_kwargs) + config = ModuleWithNestedConfig(**merged) + + assert config.sub.a == 1 + assert config.sub.b == 20 + + @pytest.mark.skipif_macos_bug def test_collect_stats(create_worker_manager): from dimos.core.resource_monitor.monitor import StatsMonitor diff --git a/dimos/core/coordination/worker_manager_python.py b/dimos/core/coordination/worker_manager_python.py index 4963db2dac..2956dff71a 100644 --- a/dimos/core/coordination/worker_manager_python.py +++ b/dimos/core/coordination/worker_manager_python.py @@ -17,6 +17,8 @@ from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any +from pydantic import BaseModel + from dimos.core.coordination.python_worker import PythonWorker from dimos.core.global_config import GlobalConfig from dimos.core.module import ModuleBase, ModuleSpec @@ -30,6 +32,26 @@ logger = setup_logger() +def _merge_config_kwargs(base: Mapping[str, Any], overrides: Mapping[str, Any]) -> dict[str, Any]: + merged = dict(base) + for key, override_value in overrides.items(): + merged[key] = _merge_config_value(merged.get(key), override_value) + return merged + + +def _merge_config_value(base_value: Any, override_value: Any) -> Any: + if not isinstance(override_value, Mapping): + return override_value + + if isinstance(base_value, BaseModel): + base_value = base_value.model_dump(mode="python") + + if not isinstance(base_value, Mapping): + return dict(override_value) + + return _merge_config_kwargs(base_value, override_value) + + class WorkerManagerPython: deployment_identifier: str = "python" @@ -161,7 +183,11 @@ def deploy_parallel( module_class, _, kwargs = specs[i] worker = self._select_worker(dedicated=module_class.dedicated_worker) worker.reserve_slot() - kwargs.update(blueprint_args.get(module_class.name, {})) + specs[i] = ( + specs[i][0], + specs[i][1], + _merge_config_kwargs(kwargs, blueprint_args.get(module_class.name, {})), + ) workers_by_index[i] = worker assignments = [(workers_by_index[i], specs[i]) for i in range(len(specs))]