Conversation
- Add mod.stop() to test_process_crash_triggers_stop so watchdog, LCM, and event-loop threads are properly joined from the test thread - Filter third-party daemon threads with generic names (Thread-\d+) in conftest monitor_threads to ignore torch/HF background threads that have no cleanup API
Convert test_process_crash_triggers_stop to use a fixture that calls mod.stop() in teardown. The watchdog thread calls self.stop() but can't join itself, so an explicit stop() from the test thread is needed to properly clean up all threads. Drop the broad conftest regex filter for generic daemon thread names per review feedback.
mod.stop() is a no-op when the watchdog already called it, so capture thread IDs before the test and join new ones in teardown.
…dimos into jeff/fix/native_threading
Greptile SummaryThis PR introduces a suite of smart, auto-cleaning thread primitives ( Key changes and their rationale:
Confidence Score: 4/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant Main as Main Thread
participant MT as ModuleThread
participant AT as AsyncModuleThread
participant MP as ModuleProcess
participant WD as Watchdog Thread
participant Disp as CompositeDisposable
Main->>AT: __init__(module)
AT->>AT: new_event_loop()
AT->>AT: thread.start() [runs loop.run_forever]
AT->>Disp: add(Disposable(self.stop))
Main->>MP: __init__(module)
MP->>Disp: add(Disposable(self.stop))
MP->>MP: subprocess.Popen(...)
MP->>MT: ModuleThread(module, target=_watch)
MT->>Disp: add(Disposable(self.stop))
MT->>WD: thread.start()
Note over WD: process exits naturally
WD->>WD: proc.wait() returns
WD->>WD: _stopped? No → call on_exit()
WD->>Main: on_exit() → module.stop()
Main->>Disp: dispose()
Disp->>MP: stop() → _stopped=True, process=None
Disp->>MT: stop() → _stop_event.set()
MT->>MT: current_thread == watchdog? → skip join()
Note over MT: No deadlock ✓
Note over Main: Normal teardown path
Main->>Disp: dispose()
Disp->>AT: stop() → loop.call_soon_threadsafe(loop.stop)
Disp->>MP: stop() → SIGTERM → wait → SIGKILL if needed
Disp->>MT: stop() → join(timeout)
Reviews (1): Last reviewed commit: "misc improve" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| assert done.wait(timeout=10), "Deadlock with slow ModuleThread.stop()" | ||
|
|
||
|
|
||
| from dimos.utils.typing_utils import ExceptionGroup |
There was a problem hiding this comment.
ExceptionGroup imported at bottom of file, used earlier
ExceptionGroup is imported on line 888 but first used on line 750 inside TestSafeThreadMap methods. This works at runtime because the full module is loaded before any test runs, but it's confusing to readers: the symbol appears to be undefined at its use sites, and any linter or static analysis tool will flag these as NameErrors. The import should be moved to the top-level imports block alongside the other third-party imports.
| from dimos.utils.typing_utils import ExceptionGroup | |
| from dimos.utils.typing_utils import ExceptionGroup |
(Move this to the top of the file alongside the other dimos.utils imports, and remove line 888.)
| """ | ||
|
|
||
| @staticmethod | ||
| def _make_fake_stop(mod: FakeModule, done: threading.Event) -> Callable: |
There was a problem hiding this comment.
Callable import used in return-type annotation
Callable is referenced as a return-type annotation in _make_fake_stop but is never imported in this file. With from __future__ import annotations in effect, the annotation is stored as a string at definition time and won't raise a NameError at runtime. However, any call to typing.get_type_hints(_make_fake_stop) — including some test introspection tools — will fail with NameError: name 'Callable' is not defined.
Add to the imports at the top of the file:
from collections.abc import Callable| self._watchdog = ModuleThread( | ||
| module=self._module, | ||
| target=self._watch, | ||
| name=f"proc-{self._process.pid}-watchdog", | ||
| ) |
There was a problem hiding this comment.
ModuleProcess.start() call adds a new ModuleThread disposable
Every time start() is called (line 388), a new ModuleThread is constructed for the watchdog. ModuleThread.__init__ immediately registers a Disposable(self.stop) in module._disposables (line 155). CompositeDisposable simply appends, so restarting the process accumulates stale disposables for watchdog threads that have already exited.
For the single-use lifecycle this is fine. But if start() is ever called more than once (e.g. after a failed first attempt, or the deferred-start path), the module's disposable list grows unboundedly, and on teardown each old watchdog's stop() is called even though it already finished, which — while idempotent — is surprising and hard to debug.
Consider either:
- Explicitly removing the old watchdog disposable before creating a new one, or
- Documenting clearly that
start()is a one-shot operation and raising an error on re-entry.
There was a problem hiding this comment.
super().start() will throw if its called more than once. We can/should assume start isn't being called multiple times AFAIK.
setstate getstate are different though, start could be called after setstate I believe
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
…dimos into jeff/fix/native_threading
| with self.mod_state as state: | ||
| if state == "stopped": | ||
| raise RuntimeError(f"{type(self).__name__} cannot be restarted after stop") | ||
| self.mod_state.set("started") |
There was a problem hiding this comment.
I know lots of modules don't call super().start() but they also wouldn't be using mod_state cause its a new thing.
Different/off-topic discussion, but I think core2 should have ModuleBase as class decorator instead of an inherited class (we can basically wrap methods instead of saying "please remember to call super").
| loop = getattr(self, "_loop", None) | ||
| # dispose of things BEFORE making aspects like rpc and _tf invalid | ||
| if hasattr(self, "_disposables"): | ||
| self._disposables.dispose() # stops _async_thread via disposable |
There was a problem hiding this comment.
I think its important to move disposables up before the rpc stop and the tf stop
| if self._uvicorn_server: | ||
| self._uvicorn_server.should_exit = True | ||
| loop = self._loop | ||
| if loop is not None and self._serve_future is not None: |
There was a problem hiding this comment.
the loop is always there until super().stop() is called
| server = uvicorn.Server(config) | ||
| self._uvicorn_server = server | ||
| loop = self._loop | ||
| assert loop is not None |
There was a problem hiding this comment.
loop always there until stop is called
| return s.getsockname()[1] | ||
|
|
||
|
|
||
| def test_mcp_server_lifecycle() -> None: |
dimos/core/test_core.py
Outdated
| assert hasattr(class_rpcs["start"], "__rpc__"), "start should have __rpc__ attribute" | ||
|
|
||
| nav._close_module() | ||
| nav._stop() |
There was a problem hiding this comment.
I'm trying to consolidate our naming to be "stop" instead of half "stop" half "close"
| # ThreadSafeVal: a lock-protected value with context-manager support | ||
|
|
||
|
|
||
| class ThreadSafeVal(Generic[T]): |
There was a problem hiding this comment.
this is my favorite util. I hate having _thing and _thing_lock and _thing2 and _thing2_lock, but I also hate seeing _thing being used in a method and thinking "hmm ... does _thing have a lock thats not being used?". This prevents ambiguity about what vals need locks and what vals don't
| self._thread.start() | ||
|
|
||
| def stop(self) -> None: | ||
| """Signal the thread to stop and join it. |
There was a problem hiding this comment.
this is probably the part that needs the most review
| # safe_thread_map: parallel map that collects all results before raising | ||
|
|
||
|
|
||
| def safe_thread_map( |
There was a problem hiding this comment.
Not used in this PR, but is used by the docker branch so getting it in here a bit early cause this is the util file it belongs in
|
|
||
| if sys.version_info < (3, 11): | ||
|
|
||
| class ExceptionGroup(Exception): # type: ignore[no-redef] # noqa: N818 |
There was a problem hiding this comment.
I didn't want to repeat all this cludge so I put it here. Let me know if there's a better spot
| if self._thread.is_alive() and self._thread is not threading.current_thread(): | ||
| self._thread.join(timeout=self._close_timeout) | ||
|
|
||
| def join(self, timeout: float | None = None) -> None: |
There was a problem hiding this comment.
I don't think you need join since you're already join()-ing in stop.
dimos/utils/thread_utils.py
Outdated
| self._stopped = False | ||
| self._stop_lock = threading.Lock() |
There was a problem hiding this comment.
Why do you need _stopped and _stop_lock? You have _stop_event.
dimos/utils/thread_utils.py
Outdated
|
|
||
| def start(self) -> None: | ||
| """Start the underlying thread.""" | ||
| self._stop_event.clear() |
There was a problem hiding this comment.
You don't need this. It's already off. If you want ModuleThread to be restartable, then you need to use another thread since threads aren't restartable.
dimos/utils/thread_utils.py
Outdated
| if start: | ||
| self.start() |
There was a problem hiding this comment.
Noooooo, don't autostart in the constructor. 😭
There was a problem hiding this comment.
😈 no boilerplate
But fr, how do you feel about ModuleThread().start()
| self._worker = ModuleThread( | ||
| module=self, | ||
| target=self._run_loop, | ||
| name="my-worker", |
There was a problem hiding this comment.
It would be nice if ModuleThread used self.module.__class__.__name__ as the prefix so we can just leave name blank most of the time and it still produces a useful name for debugging.
dimos/utils/thread_utils.py
Outdated
| return f"ThreadSafeVal({self._value!r})" | ||
|
|
||
|
|
||
| # ModuleThread: a thread that auto-registers with a module's disposables |
There was a problem hiding this comment.
Why add this if there's a docstring below?
There was a problem hiding this comment.
cause AI loves redundancy
(I'll remove it, thanks for bringing attention)
dimos/core/module.py
Outdated
| def _close_module(self) -> None: | ||
| with self._module_closed_lock: | ||
| if self._module_closed: | ||
| def _stop(self) -> None: |
There was a problem hiding this comment.
_close_module is a remnant from the the Module class hierarchy was more complicated. Some classes were skipping Module.__init__ and didn't initialize self._disposables for example. That's why I'm using hasattr(self, "_disposables") or hasattr(self, "_tf"). We didn't even have stop then.
I think it's not needed at all anymore. This could be deleted if you want and moved into def stop.
There was a problem hiding this comment.
happily! I though it was a rpc vs non-rpc thing
| self._worker = ModuleThread( | ||
| module=self, | ||
| target=self._run_loop, | ||
| name="my-worker", |
There was a problem hiding this comment.
| name="my-worker", | |
| name=self.module.__class__.__name__+"_my_worker", |
- Merge _stop() into stop() in ModuleBase (removes unnecessary indirection) - Update all callers of _stop() to use stop() directly - Add thread_start() convenience function that creates + starts a ModuleThread
AsyncModuleThread no longer spawns the event loop thread in __init__. The loop is created on the first call to start(), which ModuleBase.start() now calls. This means module construction no longer has side effects — no threads are spawned until the module is explicitly started.
3 participants
Problem
Flakey test, and a history of Flakey tests surrounding threads in modules.
Solution
Smart thread tooling with auto-cleanup to reduce the risk of bad cleanup and reduce bloat.
Breaking Changes
None, did more testing than usual because this touches core.
How to Test
Contributor License Agreement