feat(dask): remove dask by paul-nechifor · Pull Request #1365 · dimensionalOS/dimos

paul-nechifor · 2026-02-25T06:33:22Z

Problem

Dask is slow to start.
Causes issues with shutdown (you get ugly errors when hitting Ctrl+C).
Required a lot of ugly patches in core/__init__.py
Causes confusion because things are imported once per worker. (E.g.: see Jeff's rerun issue, or my actor confusion)

Closes DIM-609

Solution

Removed the dask package entirely.
Moved imports out of dimos.core. Now you have to import from the original place.
Removed everything in dimos/core/__init__.py.
Made dimos shut down cleanly, without any exceptions, and right away.
Inlined some of the imports which were causing slowdowns.
Added a short doc about py-spy so others can investigate speed if they want to.
Fixed mujoco errors on shutdown.

Breaking Changes

UtilizationModule is gone.
Blueprint configurations cannot use lambdas anymore because multiprocessing requires params to be pickleable, and lambdas are not because they don't have names. This worked with Dask because Dask serialized the bytecode of the lambda and reconstructed it on the other side.

How to Test

Test any blueprint. Example:

uv run dimos run unitree-go2

Contributor License Agreement

I have read and approved the CLA.

greptile-apps · 2026-02-25T06:33:40Z

Too many files changed for review. (135 files found, 100 file limit)

paul-nechifor · 2026-02-25T07:49:51Z

dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py

-            opacity=0.2,
-            background="#484981",
-        ),
+        "world/camera_info": _convert_camera_info,


The reason this change is needed is because multiprocessing can't pickle lambdas. Lambdas don't have names so they can't be pickled because there's nothing to reference them from the other side.

Dask cheats by serializing the bytecode of the lambda.

paul-nechifor · 2026-02-27T20:08:29Z

dimos/core/__init__.py

@@ -1,278 +0,0 @@
-from __future__ import annotations


deleted dask stuff here

paul-nechifor · 2026-02-27T20:10:02Z

dimos/core/module.py

        self._close_module()

    def _close_module(self) -> None:
+        with self._module_closed_lock:


don't double close

paul-nechifor · 2026-02-27T20:10:33Z

dimos/core/module.py

            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:


no longer need these dask functions

paul-nechifor · 2026-02-27T20:11:37Z

dimos/core/module_coordinator.py

+        for module_class, module in reversed(self._deployed_modules.items()):
+            logger.info("Stopping module...", module=module_class.__name__)
+            try:
+                module.stop()


continue stopping other modules even if one errors

paul-nechifor · 2026-02-27T20:15:50Z

dimos/core/test_worker.py

+
+
+@pytest.mark.slow
+def test_worker_pool_modules_share_workers(create_worker_manager):


add test that verifies that two modules can deploy to the same worker

paul-nechifor · 2026-02-27T20:17:59Z

dimos/core/worker.py

-        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."""


Changed Worker to support deploying multiple modules to it.

paul-nechifor · 2026-02-27T20:21:54Z

dimos/core/worker_manager.py

        self._closed = False
+        self._started = False
+
+    def start(self) -> None:


Start the workers beforehand.

paul-nechifor · 2026-02-27T20:32:13Z

dimos/hardware/sensors/camera/test_webcam.py

@@ -1,60 +0,0 @@
-# Copyright 2025-2026 Dimensional Inc.


I don't think this was used and we have a blueprint to test the cam

paul-nechifor · 2026-02-27T20:33:06Z

dimos/models/vl/base.py

    Returns:
        Detection2DBBox instance or None if invalid
    """
+    from dimos.perception.detection.type import Detection2DBBox


Inline to avoid torch import.

paul-nechifor · 2026-02-27T20:36:36Z

dimos/robot/drone/drone.py

-        else:
-            logger.warning("Tracking module failed to start")
-
+        self.connection.start()


start no longer returns a boolean

paul-nechifor · 2026-02-27T20:37:07Z

dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_detection.py

 from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_basic import unitree_g1_basic

+
+def _person_only(det: Any) -> bool:


cannot serialize lambdas, so using functions

leshy

large PR but actually acan't find any objections to core/ stuff

spomichter · 2026-02-28T09:50:07Z

dimos/agents/agent.py

            model = MockModel(json_path=self.config.model_fixture)

        with self._lock:
+            # Here to prevent unwanted imports in the file.


@paul-nechifor I assume LLM generated comments its on every import shift. not useful and looks ugly

They're not LLM generated. I specifically added them there to indicate that those imports should not be moved to the top because they cause imports (like that of torch) which slow down the start of dimos run.

Without those comments, people might move them to the top. I often move unnecessary local imports because they're added by LLMs. LLMs often prefer to add local imports because they have the habit of producing the smallest diff possible, even when it doesn't make sense.

spomichter · 2026-02-28T09:50:13Z

dimos/agents/skills/person_follow.py


        with self._lock:
            if self._tracker is None:
+                # Here to prevent unwanted imports in the file.


spomichter · 2026-02-28T09:50:20Z

dimos/agents/skills/navigation.py

        super().__init__()
        self._skill_started = False
+
+        # Here to prevent unwanted imports in the file.


paul-nechifor marked this pull request as draft February 25, 2026 06:33

paul-nechifor commented Feb 25, 2026

View reviewed changes

paul-nechifor marked this pull request as ready for review February 27, 2026 00:34

paul-nechifor force-pushed the paul/feat/remove-dask branch from 86f65d9 to 889662a Compare February 27, 2026 00:36

paul-nechifor commented Feb 27, 2026

View reviewed changes

dimos/core/__init__.py

@@ -1,278 +0,0 @@

from __future__ import annotations

Copy link

Contributor Author

paul-nechifor Feb 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

deleted dask stuff here

paul-nechifor commented Feb 27, 2026

View reviewed changes

feat(dask): remove dask

31f8829

paul-nechifor force-pushed the paul/feat/remove-dask branch from 04f07b2 to 31f8829 Compare February 27, 2026 22:57

leshy approved these changes Feb 28, 2026

View reviewed changes

leshy merged commit ed35bd7 into dev Feb 28, 2026
12 checks passed

spomichter reviewed Feb 28, 2026

View reviewed changes

spomichter mentioned this pull request Feb 28, 2026

chore: remove unnecessary "prevent unwanted imports" comments #1379

Closed

3 tasks



		@pytest.mark.slow
		def test_worker_pool_modules_share_workers(create_worker_manager):

		from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_basic import unitree_g1_basic


		def _person_only(det: Any) -> bool:

Conversation

paul-nechifor commented Feb 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Solution

Breaking Changes

How to Test

Contributor License Agreement

Uh oh!

greptile-apps bot commented Feb 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

leshy left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

paul-nechifor Feb 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

paul-nechifor commented Feb 25, 2026 •

edited

Loading

greptile-apps bot commented Feb 25, 2026 •

edited

Loading

paul-nechifor Feb 28, 2026 •

edited

Loading