refactor(runtime): extract ModuleReloader/ModuleWatcher into AutoreloadManager (#9590)

Author: mscolnick

marimo/_runtime/reload/manager.py

@@ -0,0 +1,101 @@
+# Copyright 2026 Marimo. All rights reserved.
+"""Autoreload manager: owns `ModuleReloader` and `ModuleWatcher` on behalf of the kernel."""
+
+from __future__ import annotations
+
+import contextlib
+import sys
+from typing import TYPE_CHECKING, Literal
+
+from marimo._runtime.reload.autoreload import ModuleReloader
+from marimo._runtime.reload.module_watcher import ModuleWatcher
+from marimo._utils.platform import is_pyodide
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from marimo._ast.cell import CellImpl
+    from marimo._runtime.runner.hook_context import OnFinishHookContext
+    from marimo._runtime.runtime import Kernel
+
+AutoReloadMode = Literal["off", "lazy", "autorun"]
+
+
+class AutoreloadManager:
+    """Owns ModuleReloader + ModuleWatcher and reacts to config changes."""
+
+    def __init__(self, kernel: Kernel) -> None:
+        self._kernel = kernel
+        self._reloader: ModuleReloader | None = None
+        self._watcher: ModuleWatcher | None = None
+
+        # Re-arm the watcher after every kernel run, regardless of trigger.
+        kernel._hooks.add_on_finish(self._on_finish_hook)
+
+    @property
+    def reloader(self) -> ModuleReloader | None:
+        return self._reloader
+
+    @property
+    def watcher(self) -> ModuleWatcher | None:
+        return self._watcher
+
+    def update_from_config(self, mode: AutoReloadMode) -> None:
+        """Start, stop, or swap the watcher/reloader to match `runtime.auto_reload`."""
+        # Pyodide doesn't support hot module reloading.
+        if (mode == "lazy" or mode == "autorun") and not is_pyodide():
+            if self._reloader is None:
+                self._reloader = ModuleReloader()
+
+            if self._watcher is not None and self._watcher.mode != mode:
+                self._watcher.stop()
+                self._watcher = None
+
+            if self._watcher is None:
+                self._watcher = ModuleWatcher(
+                    self._kernel.graph,
+                    reloader=self._reloader,
+                    enqueue_run_stale_cells=self._kernel._execute_stale_cells_callback,
+                    mode=mode,
+                    stream=self._kernel.stream,
+                )
+        else:
+            self._reloader = None
+            if self._watcher is not None:
+                self._watcher.stop()
+                self._watcher = None
+
+    def teardown(self) -> None:
+        if self._watcher is not None:
+            self._watcher.stop()
+            self._watcher = None
+        self._reloader = None
+
+    def flag_if_imports_stale(self, cell: CellImpl) -> None:
+        reloader = self._reloader
+        if reloader is None:
+            return
+        if reloader.cell_uses_stale_modules(cell):
+            self._kernel.graph.set_stale({cell.cell_id}, prune_imports=True)
+
+    @contextlib.contextmanager
+    def cell_scope(self) -> Iterator[None]:
+        """Reload modified modules on entry; record mtimes for newly-imported modules on exit."""
+        if self._reloader is None:
+            yield
+            return
+        snapshot = set(sys.modules)
+        self._reloader.check(modules=sys.modules, reload=True)
+        try:
+            yield
+        finally:
+            new_modules = set(sys.modules) - snapshot
+            self._reloader.check(
+                modules={m: sys.modules[m] for m in new_modules},
+                reload=False,
+            )
+
+    def _on_finish_hook(self, ctx: OnFinishHookContext) -> None:
+        del ctx
+        if self._watcher is not None:
+            self._watcher.run_is_processed.set()

marimo/_runtime/runtime.py

@@ -129,8 +129,7 @@
     start_parent_poller,
 )
 from marimo._runtime.redirect_streams import redirect_streams
-from marimo._runtime.reload.autoreload import ModuleReloader
-from marimo._runtime.reload.module_watcher import ModuleWatcher
+from marimo._runtime.reload.manager import AutoreloadManager
 from marimo._runtime.request_router import RequestRouter
 from marimo._runtime.runner import cell_runner, hook_context
 from marimo._runtime.runner.hooks import (
@@ -556,8 +555,7 @@ def __init__(
         self.module_registry = ModuleRegistry(
             self.graph, excluded_modules=set()
         )
-        self.module_reloader: ModuleReloader | None = None
-        self.module_watcher: ModuleWatcher | None = None
+        self.autoreload_manager = AutoreloadManager(self)
 
         # Load runtime settings from user config
         self.user_config = user_config
@@ -627,8 +625,7 @@ def teardown(self) -> None:
             self.stdin._stop()
         self.stream.stop()
 
-        if self.module_watcher is not None:
-            self.module_watcher.stop()
+        self.autoreload_manager.teardown()
 
         # TODO(akshayka): There's a memory leak in run mode, with memory
         # usage increasing with each session creation. Somehow the kernel
@@ -653,35 +650,7 @@ def _update_runtime_from_user_config(self, config: MarimoConfig) -> None:
         self.user_config = config
 
         self.packages_callbacks.update_package_manager(package_manager)
-
-        if (
-            (autoreload_mode == "lazy" or autoreload_mode == "autorun")
-            # Pyodide doesn't support hot module reloading
-            and not is_pyodide()
-        ):
-            if self.module_reloader is None:
-                self.module_reloader = ModuleReloader()
-
-            if (
-                self.module_watcher is not None
-                and self.module_watcher.mode != autoreload_mode
-            ):
-                self.module_watcher.stop()
-                self.module_watcher = None
-
-            if self.module_watcher is None:
-                self.module_watcher = ModuleWatcher(
-                    self.graph,
-                    reloader=self.module_reloader,
-                    enqueue_run_stale_cells=self._execute_stale_cells_callback,
-                    mode=autoreload_mode,
-                    stream=self.stream,
-                )
-        else:
-            self.module_reloader = None
-            if self.module_watcher is not None:
-                self.module_watcher.stop()
-                self.module_watcher = None
+        self.autoreload_manager.update_from_config(autoreload_mode)
 
     @property
     def globals(self) -> dict[Any, Any]:
@@ -756,25 +725,12 @@ def _install_execution_context(
                 stderr=self.stderr,
                 stdin=self.stdin,
             ),
+            self.autoreload_manager.cell_scope(),
         ):
-            modules = None
             try:
-                if self.module_reloader is not None:
-                    # Reload modules if they have changed
-                    modules = set(sys.modules)
-                    self.module_reloader.check(
-                        modules=sys.modules, reload=True
-                    )
                 yield exec_ctx
             finally:
                 ctx.execution_context = None
-                if self.module_reloader is not None and modules is not None:
-                    # Note timestamps for newly loaded modules
-                    new_modules = set(sys.modules) - modules
-                    self.module_reloader.check(
-                        modules={m: sys.modules[m] for m in new_modules},
-                        reload=False,
-                    )
 
     def _register_cell(
         self,
@@ -796,12 +752,7 @@ def _register_cell(
             self.graph.cells[cell_id].set_stale(stale=True, broadcast=False)
         # leaky abstraction: the graph doesn't know about stale modules, so
         # we have to check for them here.
-        module_reloader = self.module_reloader
-        if (
-            module_reloader is not None
-            and module_reloader.cell_uses_stale_modules(cell)
-        ):
-            self.graph.set_stale({cell.cell_id}, prune_imports=True)
+        self.autoreload_manager.flag_if_imports_stale(cell)
         LOGGER.debug("registered cell %s", cell_id)
         LOGGER.debug("parents: %s", self.graph.parents[cell_id])
         LOGGER.debug("children: %s", self.graph.children[cell_id])
@@ -1805,9 +1756,6 @@ async def run_stale_cells(self) -> None:
             )
         )
 
-        if self.module_watcher is not None:
-            self.module_watcher.run_is_processed.set()
-
     @kernel_tracer.start_as_current_span("set_cell_config")
     async def set_cell_config(self, request: UpdateCellConfigCommand) -> None:
         """Update cell configs.

tests/_runtime/reload/test_module_watcher.py

@@ -944,13 +944,14 @@ async def test_module_watcher_stop(
         await asyncio.sleep(INTERVAL)
 
         # Stop the watcher
-        assert k.module_watcher is not None
-        assert not k.module_watcher.should_exit.is_set()
+        watcher = k.autoreload_manager.watcher
+        assert watcher is not None
+        assert not watcher.should_exit.is_set()
 
-        k.module_watcher.stop()
+        watcher.stop()
 
         # should_exit should be set
-        assert k.module_watcher.should_exit.is_set()
+        assert watcher.should_exit.is_set()
 
     async def test_module_watcher_processes_flag(
         self, execution_kernel: Kernel, exec_req: ExecReqProvider
@@ -965,9 +966,10 @@ async def test_module_watcher_processes_flag(
         # Give watcher time to start
         await asyncio.sleep(INTERVAL)
 
-        assert k.module_watcher is not None
+        watcher = k.autoreload_manager.watcher
+        assert watcher is not None
         # Initially should be set (no run in flight)
-        assert k.module_watcher.run_is_processed.is_set()
+        assert watcher.run_is_processed.is_set()
 
 
 class TestModuleWatcherEdgeCases:
@@ -1075,3 +1077,83 @@ async def test_module_watcher_cache_invalidation(
 
         # The cache should handle the modified imports correctly
         assert not k.graph.cells[er_1.cell_id].stale
+
+
+class TestAutoreloadManagerLifecycle:
+    """Tests for AutoreloadManager responding to runtime.auto_reload changes."""
+
+    async def test_mode_swap_replaces_watcher(
+        self, execution_kernel: Kernel, exec_req: ExecReqProvider
+    ):
+        del exec_req
+        k = execution_kernel
+        config = copy.deepcopy(DEFAULT_CONFIG)
+
+        config["runtime"]["auto_reload"] = "lazy"
+        k.set_user_config(UpdateUserConfigCommand(config=config))
+        lazy_watcher = k.autoreload_manager.watcher
+        assert lazy_watcher is not None
+        assert lazy_watcher.mode == "lazy"
+
+        config["runtime"]["auto_reload"] = "autorun"
+        k.set_user_config(UpdateUserConfigCommand(config=config))
+        autorun_watcher = k.autoreload_manager.watcher
+        assert autorun_watcher is not None
+        assert autorun_watcher.mode == "autorun"
+        assert autorun_watcher is not lazy_watcher
+        assert lazy_watcher.should_exit.is_set()
+
+    async def test_reloader_persists_across_mode_swap(
+        self, execution_kernel: Kernel, exec_req: ExecReqProvider
+    ):
+        """Switching between lazy and autorun must not throw away the
+        reloader's mtime state — otherwise every swap would force a reload
+        of all already-tracked modules."""
+        del exec_req
+        k = execution_kernel
+        config = copy.deepcopy(DEFAULT_CONFIG)
+
+        config["runtime"]["auto_reload"] = "lazy"
+        k.set_user_config(UpdateUserConfigCommand(config=config))
+        reloader = k.autoreload_manager.reloader
+        assert reloader is not None
+
+        config["runtime"]["auto_reload"] = "autorun"
+        k.set_user_config(UpdateUserConfigCommand(config=config))
+        assert k.autoreload_manager.reloader is reloader
+
+    async def test_disable_clears_manager_state(
+        self, execution_kernel: Kernel, exec_req: ExecReqProvider
+    ):
+        del exec_req
+        k = execution_kernel
+        config = copy.deepcopy(DEFAULT_CONFIG)
+
+        config["runtime"]["auto_reload"] = "lazy"
+        k.set_user_config(UpdateUserConfigCommand(config=config))
+        assert k.autoreload_manager.watcher is not None
+        assert k.autoreload_manager.reloader is not None
+
+        config["runtime"]["auto_reload"] = "off"
+        k.set_user_config(UpdateUserConfigCommand(config=config))
+        assert k.autoreload_manager.watcher is None
+        assert k.autoreload_manager.reloader is None
+
+    async def test_run_rearms_watcher_run_is_processed(
+        self, execution_kernel: Kernel, exec_req: ExecReqProvider
+    ):
+        """The on_finish hook fires for every kernel run, not just
+        `_run_stale_cells`. This is a deliberate semantic change from the
+        pre-extraction code, which only set the flag at the end of stale-cell
+        runs."""
+        k = execution_kernel
+        config = copy.deepcopy(DEFAULT_CONFIG)
+        config["runtime"]["auto_reload"] = "lazy"
+        k.set_user_config(UpdateUserConfigCommand(config=config))
+
+        watcher = k.autoreload_manager.watcher
+        assert watcher is not None
+
+        watcher.run_is_processed.clear()
+        await k.run([exec_req.get("x = 1")])
+        assert watcher.run_is_processed.is_set()