feat(save): stub serialization toolkit (class/lazy/module stubs) (#9896)

Expands the current stubbing mechanism for larger coverage for various
python instances (classes, lambdas, modules, and module imports)

Introduces:
 - UnhashableStub
 - ClassStub

Additionally:
- Handles module imports explicitly to prevent their serialization to
disk
 - and replaces pytorch tensor Pickled objects with a `.pt` loader

Author: dmadisetti

marimo/_runtime/exceptions.py

@@ -5,6 +5,7 @@
 
 if TYPE_CHECKING:
     from marimo._runtime.dataflow import DirectedGraph
+    from marimo._types.ids import CellId_t
 
 
 class MarimoRuntimeException(BaseException):
@@ -26,6 +27,39 @@ def __init__(self, ref: str, name_error: NameError | None = None) -> None:
         self.name_error = name_error
 
 
+class MarimoCancelCellError(BaseException):
+    """Soft-cancel signal raised by a lifecycle.
+
+    Subclasses BaseException (not Exception) so user-code `except Exception`
+    blocks don't swallow the control-flow signal.
+    """
+
+    cells_to_rerun: set[CellId_t]
+
+    def __init__(
+        self,
+        *args: object,
+        cells_to_rerun: set[CellId_t] | None = None,
+    ) -> None:
+        super().__init__(*args)
+        self.cells_to_rerun = cells_to_rerun or set()
+
+
+class MarimoUnhashableCacheError(MarimoCancelCellError):
+    """Raised when cell-level caching encounters a value that cannot be
+    hashed/serialized for cache restoration."""
+
+    def __init__(
+        self,
+        cells_to_rerun: set[CellId_t],
+        variables: list[str],
+        error_details: str,
+    ) -> None:
+        super().__init__(error_details, cells_to_rerun=cells_to_rerun)
+        self.variables = variables
+        self.error_details = error_details
+
+
 def unwrap_user_exception(
     exc: MarimoRuntimeException,
     graph: DirectedGraph | None = None,

marimo/_save/cache.py

@@ -4,6 +4,7 @@
 import abc
 import inspect
 import re
+import textwrap
 from collections import namedtuple
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Literal, get_args
@@ -14,13 +15,15 @@
 from marimo._runtime.state import SetFunctor
 from marimo._save.stubs import (
     CUSTOM_STUBS,
+    ClassStub,
     CustomStub,
     FunctionStub,
     ModuleStub,
     ReferenceStub,
     UIElementStub,
     maybe_register_stub,
 )
+from marimo._save.stubs.lazy_stub import UnhashableStub
 
 # Many assertions are for typing and should always pass. This message is a
 # catch all to motive users to report if something does fail.
@@ -114,6 +117,24 @@ def dispose(self, context: RuntimeContext, deletion: bool) -> bool:  # noqa: ARG
 )
 
 
+def _source_refs(code: str) -> set[Name]:
+    """Free references of a captured class/function source block.
+
+    Reuses the cell `ScopedVisitor` so the notion of "what this def needs"
+    matches marimo's dataflow rather than a bespoke parser.
+    """
+    from marimo._ast.parse import ast_parse
+    from marimo._ast.visitor import ScopedVisitor
+
+    try:
+        tree = ast_parse(textwrap.dedent(code))
+    except SyntaxError:
+        return set()
+    visitor = ScopedVisitor("cache_restore")
+    visitor.visit(tree)
+    return set(visitor.refs)
+
+
 # BaseException because "raise _ as e" is utilized.
 class CacheException(BaseException):
     pass
@@ -136,7 +157,7 @@ class Cache:
     def restore(self, scope: dict[str, Any]) -> None:
         """Restores values from cache, into scope."""
         memo: dict[int, Any] = {}  # Track processed objects to handle cycles
-        for var, lookup in self.contextual_defs():
+        for var, lookup in self._restore_order(self.contextual_defs()):
             value = self.defs.get(var, None)
             scope[lookup] = self._restore_from_stub_if_needed(
                 value, scope, memo
@@ -167,6 +188,44 @@ def restore(self, scope: dict[str, Any]) -> None:
                     f"({type(ref)}:{ref})."
                 )
 
+    def _restore_deps(self, value: Any) -> set[Name]:
+        """Cross-def names *value* needs before it can be restored.
+
+        - A re-exec'd class/function needs the defs it references at
+          definition time (bases, decorators, class-body calls).
+        - A pickled instance of a cell-defined class needs that class
+          materialized first (tagged via `requires`).
+        """
+        if isinstance(value, (ClassStub, FunctionStub)):
+            return _source_refs(value.code)
+        requires = getattr(value, "requires", "")
+        return {requires} if requires else set()
+
+    def _restore_order(
+        self, contextual_defs: dict[tuple[Name, Name], Any]
+    ) -> list[tuple[Name, Name]]:
+        """Order `(var, lookup)` pairs so each def's cross-def dependencies
+        restore first (depth-first; tolerant of cycles via `seen`)."""
+        lookups = {var: lookup for var, lookup in contextual_defs}
+        deps = {
+            var: self._restore_deps(self.defs.get(var)) & lookups.keys()
+            for var in lookups
+        }
+        order: list[tuple[Name, Name]] = []
+        seen: set[Name] = set()
+
+        def visit(var: Name) -> None:
+            if var in seen:
+                return
+            seen.add(var)
+            for dep in deps[var]:
+                visit(dep)
+            order.append((var, lookups[var]))
+
+        for var in lookups:
+            visit(var)
+        return order
+
     def _restore_from_stub_if_needed(
         self,
         value: Any,
@@ -223,6 +282,16 @@ def _restore_from_stub_if_needed(
             result = value
         elif isinstance(value, ReferenceStub):
             result = value.load(scope)
+        elif isinstance(value, ClassStub):
+            # Re-exec the captured source into the cell namespace so the
+            # name rebinds to a live class (not the stub). Must run before
+            # any pickle blob referencing the class deserializes, so the
+            # class is resolvable as `__main__.<name>` in `scope`.
+            result = value.load(scope)
+        elif isinstance(value, UnhashableStub):
+            # Marker for a def whose value couldn't be serialized. Place it
+            # in scope as-is.
+            result = value
         elif isinstance(value, CustomStub):
             # CustomStub is a placeholder for a custom type, which cannot be
             # restored directly.
@@ -315,8 +384,28 @@ def _convert_to_stub_if_needed(
 
         if inspect.ismodule(value):
             result = ModuleStub(value)
-        elif inspect.isfunction(value):
+        elif (
+            inspect.isfunction(value)
+            and value.__name__ != "<lambda>"
+            and getattr(value, "__module__", "__main__") == "__main__"
+        ):
+            # NB. Lambdas can't round-trip via FunctionStub: inspect.getsource
+            # returns the line *containing* the lambda (e.g. "return model,
+            # lambda inp: model(inp)"), which fails to compile as a module.
             result = FunctionStub(value)
+        elif (
+            inspect.isclass(value)
+            and getattr(value, "__module__", None) == "__main__"
+        ):
+            # Attempt to capture classes by source so the loader can rebuild
+            # them in the cell namespace before pickle blobs that reference them
+            # deserialize. Pass the executing cell's source filename as a hint
+            # so attribute-only classes (no method code object to read a
+            # filename from) still resolve against `linecache`.
+            try:
+                result = ClassStub(value, filename=self._cell_filename())
+            except (TypeError, OSError):
+                result = value
         elif isinstance(value, UIElement):
             result = UIElementStub(value)
         elif isinstance(value, tuple):
@@ -397,6 +486,20 @@ def _convert_to_stub_if_needed(
 
         return result
 
+    @staticmethod
+    def _cell_filename() -> str | None:
+        """Source filename of the executing cell, for sourcing classes that
+        lack a method code object. `None` when there is no runtime context
+        (e.g. script-mode / direct API use)."""
+        try:
+            from marimo._ast.compiler import get_filename
+
+            context = get_context().execution_context
+            assert context is not None
+            return get_filename(context.cell_id)
+        except (ContextNotInitializedError, AssertionError):
+            return None
+
     def contextual_defs(self) -> dict[tuple[Name, Name], Any]:
         """Uses context to resolve private variable names."""
         try:

marimo/_save/hash.py

@@ -705,11 +705,13 @@ def serialize_and_dequeue_content_refs(
                 version = ""
                 module = None
                 if self.pin_modules:
-                    module = sys.modules[imports[ref].module]
-                    version = getattr(module, "__version__", "")
+                    # `.get`: a cached module def restored as a missing-
+                    # module placeholder is in scope but not sys.modules.
+                    module = sys.modules.get(imports[ref].module)
+                    version = getattr(module, "__version__", "") or ""
                     if not version:
-                        module = sys.modules[imports[ref].namespace]
-                        version = getattr(module, "__version__", "")
+                        module = sys.modules.get(imports[ref].namespace)
+                        version = getattr(module, "__version__", "") or ""
 
                 content_serialization[ref] = type_sign(
                     bytes(f"module:{ref}:{version}", "utf-8"), "module"
@@ -1078,6 +1080,7 @@ def cache_attempt_from_hash(
         hasher.defs,
         hasher.key,
         hasher.stateful_refs,
+        glbls=scope,
     )
 
 
@@ -1166,4 +1169,5 @@ def content_cache_attempt_from_base(
         hasher.defs,
         hasher.key,
         stateful_refs,
+        glbls=scope,
     )