perf(autoreload): skip stdlib/site-packages on per-cell check

marimo/_runtime/reload/autoreload.py

@@ -14,6 +14,7 @@
 import modulefinder
 import os
 import sys
+import sysconfig
 import threading
 import traceback
 import types
@@ -73,6 +74,19 @@ def safe_hasattr(obj: M, attr: str) -> bool:
         return False
 
 
+def _non_user_module_roots() -> tuple[str, ...]:
+    """Filesystem prefixes that hold stdlib + site-packages modules."""
+    roots: set[str] = set()
+    for key in ("stdlib", "platstdlib", "purelib", "platlib"):
+        p = sysconfig.get_path(key)
+        if p:
+            roots.add(p)
+    # Covers stdlib for non-venv interpreters, where sysconfig's stdlib path
+    # may not match the runtime location.
+    roots.add(sys.base_prefix)
+    return tuple(roots)
+
+
 def modules_imported_by_cell(
     cell: CellImpl, sys_modules: dict[str, types.ModuleType]
 ) -> set[str]:
@@ -160,10 +174,26 @@ def __init__(self) -> None:
         # for thread-safety
         self.lock = threading.Lock()
         self._module_dependency_finder = ModuleDependencyFinder()
+        # Names known to live in stdlib/site-packages; populated lazily on the
+        # hot per-cell path. Entries are never evicted: a module whose
+        # __file__ moves between roots at runtime would not be re-evaluated.
+        self._skip: set[str] = set()
+        self._non_user_roots = _non_user_module_roots()
 
         # Timestamp existing modules
         self.check(modules=sys.modules, reload=False)
 
+    def _is_user_module(self, module: types.ModuleType) -> bool:
+        """True for modules whose source lives outside stdlib/site-packages.
+
+        Editable installs (e.g. ``pip install -e .``) point ``__file__`` at the
+        source tree, so they are correctly classified as user code.
+        """
+        f = safe_getattr(module, "__file__", None)
+        if not f:
+            return False
+        return not f.startswith(self._non_user_roots)
+
     def filename_and_mtime(
         self, module: types.ModuleType
     ) -> ModuleMTime | None:
@@ -206,12 +236,21 @@ def cell_uses_stale_modules(self, cell: CellImpl) -> bool:
             )
 
     def check(
-        self, modules: dict[str, types.ModuleType], reload: bool
+        self,
+        modules: dict[str, types.ModuleType],
+        reload: bool,
+        *,
+        skip_non_user_modules: bool = False,
     ) -> set[types.ModuleType]:
         """Check timestamps of modules, optionally reload them.
 
         Also patches existing objects with hot-reloaded ones.
 
+        When ``skip_non_user_modules`` is True, stdlib and site-packages
+        modules are skipped via a persistent cache — intended for the per-cell
+        hot path. The background ``ModuleWatcher`` keeps the default so that
+        edits to installed packages remain detectable at watcher latency.
+
         Returns a set of modules that were found to have been modified.
         """
 
@@ -222,12 +261,18 @@ def check(
         # and also iterates over it
         with self.lock:
             modified_modules: set[types.ModuleType] = set()
+            skip = self._skip if skip_non_user_modules else None
             # materialize the module keys, since we'll be reloading while
             # iterating
             for modname in list(modules.keys()):
+                if skip is not None and modname in skip:
+                    continue
                 m = modules.get(modname, None)
                 if m is None:
                     continue
+                if skip is not None and not self._is_user_module(m):
+                    skip.add(modname)
+                    continue
 
                 module_mtime = self.filename_and_mtime(m)
                 if module_mtime is None:

marimo/_runtime/reload/manager.py

@@ -85,14 +85,20 @@ def cell_scope(self) -> Iterator[None]:
             yield
             return
         snapshot = set(sys.modules)
-        self._reloader.check(modules=sys.modules, reload=True)
+        # Hot path: skip stdlib/site-packages. Edits inside an installed
+        # package are still picked up by the background ModuleWatcher's
+        # full scan, just at watcher latency rather than per-cell.
+        self._reloader.check(
+            modules=sys.modules, reload=True, skip_non_user_modules=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,
+                skip_non_user_modules=True,
             )
 
     def _on_finish_hook(self, ctx: OnFinishHookContext) -> None:

tests/_runtime/reload/test_autoreload.py

@@ -425,6 +425,133 @@ def test_check_reload_clears_stale_modules(
         assert len(reloader.stale_modules) == 0
 
 
+class TestSkipCache:
+    def test_is_user_module_stdlib(self):
+        reloader = ModuleReloader()
+        assert reloader._is_user_module(sys.modules["os"]) is False
+        assert reloader._is_user_module(sys.modules["pathlib"]) is False
+
+    def test_is_user_module_builtin_has_no_file(self):
+        reloader = ModuleReloader()
+        assert reloader._is_user_module(sys.modules["sys"]) is False
+        assert reloader._is_user_module(sys.modules["builtins"]) is False
+
+    def test_is_user_module_user_code(
+        self, tmp_path: pathlib.Path, py_modname: str
+    ):
+        sys.path.append(str(tmp_path))
+        py_file = tmp_path / pathlib.Path(py_modname + ".py")
+        py_file.write_text("x = 1")
+        mod = importlib.import_module(py_modname)
+        reloader = ModuleReloader()
+        assert reloader._is_user_module(mod) is True
+
+    def test_default_check_does_not_populate_skip(self):
+        # The watcher path must keep scanning everything so that edits inside
+        # installed packages remain detectable.
+        reloader = ModuleReloader()
+        assert reloader._skip == set()
+
+    def test_hot_path_populates_skip_with_stdlib(self):
+        reloader = ModuleReloader()
+        reloader.check(sys.modules, reload=False, skip_non_user_modules=True)
+        assert "os" in reloader._skip
+        assert "pathlib" in reloader._skip
+
+    def test_user_module_not_skipped_on_hot_path(
+        self, tmp_path: pathlib.Path, py_modname: str
+    ):
+        sys.path.append(str(tmp_path))
+        py_file = tmp_path / pathlib.Path(py_modname + ".py")
+        py_file.write_text("x = 1")
+        importlib.import_module(py_modname)
+        reloader = ModuleReloader()
+        reloader.check(sys.modules, reload=False, skip_non_user_modules=True)
+        assert py_modname not in reloader._skip
+
+    def test_skipped_modules_are_not_restated(self, monkeypatch):
+        reloader = ModuleReloader()
+        reloader.check(sys.modules, reload=False, skip_non_user_modules=True)
+        assert "os" in reloader._skip
+
+        calls: list[str] = []
+        orig = reloader.filename_and_mtime
+
+        def spy(module):
+            calls.append(getattr(module, "__name__", "?"))
+            return orig(module)
+
+        monkeypatch.setattr(reloader, "filename_and_mtime", spy)
+        reloader.check(sys.modules, reload=False, skip_non_user_modules=True)
+        assert "os" not in calls
+        assert "pathlib" not in calls
+
+    def test_watcher_path_still_sees_installed_packages(
+        self, tmp_path: pathlib.Path, py_modname: str, monkeypatch
+    ):
+        # Regression guard for the most surprising user-visible behavior:
+        # editing a file inside site-packages must still be detected, just
+        # via the background watcher's full scan instead of the hot path.
+        # We simulate "this module is in site-packages" by treating tmp_path
+        # as a non-user root for the duration of the test.
+        sys.path.append(str(tmp_path))
+        py_file = tmp_path / pathlib.Path(py_modname + ".py")
+        py_file.write_text("x = 1")
+        mod = importlib.import_module(py_modname)
+
+        reloader = ModuleReloader()
+        monkeypatch.setattr(
+            reloader,
+            "_non_user_roots",
+            (str(tmp_path),) + reloader._non_user_roots,
+        )
+        assert reloader._is_user_module(mod) is False
+
+        # Watcher path (default) sees it.
+        reloader.check(sys.modules, reload=False)
+        assert py_modname in reloader.modules_mtimes
+        update_file(py_file, "x = 2")
+        assert any(m is mod for m in reloader.check(sys.modules, reload=False))
+
+        # Hot path skips it.
+        reloader2 = ModuleReloader()
+        monkeypatch.setattr(
+            reloader2,
+            "_non_user_roots",
+            (str(tmp_path),) + reloader2._non_user_roots,
+        )
+        reloader2.check(sys.modules, reload=False, skip_non_user_modules=True)
+        assert py_modname in reloader2._skip
+
+    def test_user_module_reload_still_works(
+        self, tmp_path: pathlib.Path, py_modname: str
+    ):
+        sys.path.append(str(tmp_path))
+        py_file = tmp_path / pathlib.Path(py_modname + ".py")
+        py_file.write_text(
+            textwrap.dedent(
+                """
+                def foo():
+                    return 1
+                """
+            )
+        )
+        mod = importlib.import_module(py_modname)
+        reloader = ModuleReloader()
+        reloader.check(sys.modules, reload=False)
+        assert mod.foo() == 1
+
+        update_file(
+            py_file,
+            """
+            def foo():
+                return 2
+            """,
+        )
+        reloader.check(sys.modules, reload=True)
+        assert mod.foo() == 2
+
+
 class TestUpdateFunctions:
     """Tests for update_* functions"""