Repository: marimo-team/marimo
Pull request: Persistent cache
## 📝 Summary

closes #1306

Introduces `persistent_cache` which is a precursor to native saving in marimo.

TODO:
 - [x] Finish tests
 - [x] General cleanup (hash.py)
 - [x] Add / finish documentation

Following iteration:
 - `max_size` (cleanup)
 - version pinning
 - Handling state and UI

Even further
 - remote cache ?

## 🔍 Description of Changes

Moved over save notebook + fixed a few more bugs that I came across, lint passes locally, but I think there's a version difference with CI- so good to get it up here.

Have some tests locally, but needs refinement (hence draft status)

## 📋 Checklist

- [x] I have read the [contributor guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [x] For large changes, or changes that affect the public API: this change was discussed or approved through an issue, on [Discord](https://discord.gg/JE7nhX6mD8), or the community [discussions](https://github.com/marimo-team/marimo/discussions) (Please provide a link if applicable).
- [x] I have added tests for the changes made.
- [x] I have run the code and verified that it works as expected.

## 📜 Reviewers

@akshayka OR @mscolnick


Problem statement:
Native Cached Results
### Description

I find myself using this pattern a lot

```python
if not os.path.isfile("sample_selection_test_all.npy"):
    np.save(
        "sample_selection_test_all", data_sample_export(test_data, test_sample)
    )
    np.save(
        "sample_selection_train_all",
        data_sample_export(train_data, train_sample),
    )

exported_train = np.load("sample_selection_train_all")
exported_test = np.load("sample_selection_test_all")
```

I wonder if saving computations + saving outputs is on the roadmap.
(docs recommend functools memo, but I like to restart notebooks sometime)

### Suggested solution


result = mo.saved(x)

execution unwraps result, saves it to file, or loads it.
or after cell execution could unwrap result

for stateless, could hash ref dag as file name. (hash concated cell contents)

---

mo.save(plt.gca()) # Last line

allows cell to have a default output. State ensured by the ref dag hash

### Alternative

Tons of ways doing this, just provided a suggestion

### Additional context

_No response_

Hints from issue discussion:
It's something Myles and I have talked about, yes! We even considered using the same API name. So yes, we'll likely support something like this.
@akshayka @mscolnick Do you all have a road map for this? Here is a very rough outline

```python
def save(callback : Callable[[], T],
         sign_and_verify : Union[None, str, tuple[Callable[[str], str], Callable[[str], bool]] = None,
         store_and_load : Union[None, str, tuple[Callable[[str], str], Callable[[str], bool]] = None,
         pin_marimo_lib=False,
         pin_sys_libs=False,
         pin_local_libs=True,
         save_path="outputs",
         *,
         name) -> T:
    """
    Save functionality.
    
    Save a function output under a hash of cell and ancestors, as such relevant changes to the
    compute graph will invalidate a saved value. For an external example of a similar procedure,
    see the NAR spec on how inputs can be used to hash and verify file compute results:
    https://edolstra.github.io/pubs/phd-thesis.pdf (5.2.2. Store paths, p.g 93) , and the relevant
    section motivating store paths.

    Args:
        callback (Callable[[], T]): The function to be executed and its result saved.
        sign_and_verify (Union[None, str, tuple[Callable[[str], str], Callable[[str], bool]]): Optional. Function or functions to sign and verify the output. If a string is provided, it is assumed to be the name of a known signing and verifying method.
        store_and_load (Union[None, str, tuple[Callable[[str], str], Callable[[str], bool]]): Optional. Function or functions to store and load the output. If a string is provided, it is assumed to be the name of a known storing and loading method.
        pin_marimo_lib (bool): If True, the marimo library version is used for cache invalidation. Default is False.
        pin_sys_libs (bool): If True, the system libraries versions are used for cache invalidation. Default is False.
        pin_local_libs (bool): If True, the local libraries, using the __version__ variable, are used for cache invalidation. Default is True.
        save_path (str): The path where the output will be saved. Default is "outputs".
        name (str): The name of the output file.

    Returns:
        T: The result of the callback function.

    Raises:
        Exception: If the output cannot be signed, stored, or loaded.
    """

    Args:
    """
    ...
    # Run traceback, get content for code up to function call, and compile
    # Hash ancestors
    # Create digest of combined ancestors and current code content
    # if {save_path}/{hash}/{name} doesn't exist
    #  - Run function
    #  - Run sign on outputs
    #  - Save in {save_path}/{hash}/{name}
    # else:
    #  - load data!
```

I have a stub I could flesh out, but I wanted to discuss it in more detail first. Will definitely be smaller than my last PR lol, but all the same
@dmadisetti — thanks for the design sketch. It's really cool!

We don't yet have a roadmap or design spec — yours is a great place to start.

**Usage.**

To check my understanding: Is this usage correct?

Cell 1:

```python
x_value = 1
y_value = 2
```

Cell 2:

```python
def add(x, y):
  return x + y

func = lambda: add(x_value, y_value)
x_plus_y = mo.save(func)
```

**Decorator**

Would it be possible to have an API that could be used as a decorator?

```python
@mo.save
def add(x, y):
  return x + y

x_plus_y = mo.save(add)
```

or even

```python

@mo.save
@functools.cache
def add(x, y):
  return x + y

x_plus_y = mo.save(add)
```

**Signature.**

The function has many arguments. Is `sign_and_verify` really necessary? What use case do you foresee for that?


```python
@mo.save("name")
def add(x, y):
  return x + y

x_plus_y = add()
```

?

I think that's doable. The implementation I was thinking of uses the call frame to build call context. 

If a function is "pure" then the result is dependent directly on the arguments and not necessarily the execution path, so maybe that needs to be decoupled from this spec. `functools.cache` is set up to capture these cases.

There's also a nix analog for this (idk if it's in the thesis) with content addressable storage, which checks the contents of the inputs and not the input paths (what's in my spec here).

> Is sign_and_verify really necessary?

The NAR scheme gives you guarantees of the values that go into your computation, but there's no method of checking the "correctness" of the result.

Nix gets around this by making the "store" strictly read-only + caching results externally. That means that the result is essentially tamper proof. marimo doesn't have these guarantees baked in. I think an external cache is a bit much, and an RO file system is not easily enforceable (nix is pretty invasive in this regard). But we can use something like a pgp signature to verify a computation result. Probably a bit more academic than practical, but was still trying to think through the "correct" way of doing things. It could just be a set it up and forget it like GitHub signatures.

As for the other arguments, it's more to allow overriding of needed, with the expectation it'll mostly be default.
How about

```python
@mo.save("name")
def add(x, y):
  return x + y

x_plus_y = add(1, 2)
```

as well as

```python
@mo.save("name")
def add(x, y):
  return x + y + some_ref

x_plus_y = add(1, 2)
```

with `mo.save()` returning a callable whose signature matches the wrapped function? I understand that this doesn't match your spec (statically determined input paths). But otherwise (IIUC) users would need to rewrite this as

```python
def add(x, y):
  return x + y

func = lambda: add(x_value, y_value)
x_plus_y = mo.save(func, "name")
```

which feels non-obvious.

> If a function is "pure" then the result is dependent directly on the arguments and not necessarily the execution path,  so maybe that needs to be decoupled from this spec ... There's also a nix analog for this (idk if it's in the thesis) with content addressable storage, which checks the contents of the inputs and not the input paths (what's in my spec here).

It would be really nice to be able to use the same function — `mo.save` — regardless of whether the function was "pure" ("refless") or not. Especially since`functools.cache` doesn't save to disk.

It's okay if this isn't possible ... you've thought about this much more than I have. Thoughts? 

> Probably a bit more academic than practical, but was still trying to think through the "correct" way of doing things. It could just be a set it up and forget it like GitHub signatures.

Okay, sounds good, makes sense.

> As for the other arguments, it's more to allow overriding of needed, with the expectation it'll mostly be default.

Sounds good.
`save` could hide the details of hashing by content vs execution. A little naive, but the "primitive check" like in strict execution could be a way to determine "pure" functions. The pure case could kick in for the your first case example, and the execution path for the second, resolved by required_refs.

Keeping save results in memory like functools seems reasonable, execution saves could also invalidate the cache more meaningfully than functools too. I'm sure there's an edge case where it will eat up a users memory, but I think it's a premature optimization until someone complains.

For API

```python
saved = mo.savable(foo, params)(a,b,c)
```

Or

```python
with mo.saved_context(params):
    x = add(1, 2, 3)
```

The latter is tricky, but maybe possible with a bit of static analysis and swapping out the relevant functions, or hackily skipping the execution block and running it through an exec (which I think is less error prone for marimo, but might be harder for a user to debug)

Another option is saving on a cell level. Then you don't have to deal with passing in wrapped values, you can get a saved result through a flag.


> Another option is saving on a cell level. Then you don't have to deal with passing in wrapped values, you can get a saved result through a flag.

Yea I thought about that -- that could be pretty nice. If enabled, notebooks could populate with outputs instantly on startup, which would be pretty great.

> saved = mo.savable(foo, params)(a,b,c)

So this would be a different API than `mo.save`? What are `params` and how are they different from `a`, `b`, and `c`?

I guess fundamentally there are two different types of caching, which to the user may feel the same but whose implementations are different (content addressed vs path addressed) And I'm wondering if we can have the same API for both. It's okay if we can't, but we should just be clear on that.

By the way. For all of this -- `mo.save` acts like a cache, in that it can and will invalidate old entries and delete them from disk. Right? Which means that you shouldn't use `mo.save` if you want to permanently save an artifact to disk or elsewhere (eg a plot, a numpy array, etc) . If I got that right, maybe the API should be called `mo.cache`.
>      So this would be a different API than mo.save? What are params and how are they different from a, b, and c?

`params` being the keywords like `name` and `save_path`. It's a little different. The `with` block is definitely possible though

I think it's doable, but it might be worth exposing the individual functions (`execution_cache`, `function_cache`) as well, because the behavior is going to be subtly different.

Right, so we haven't mentioned deleting at all- if runs are going to be invalidated then that's definitely worth documenting and renaming from `save`.

I think the execution path case has less of a reason to invalidate cache, but a function called many times may should likely be a "cache" over "save".
> params being the keywords like name and save_path

Got it, makes sense.

> I think it's doable, but it might be worth exposing the individual functions (execution_cache, function_cache) as well, because the behavior is going to be subtly different.

Ok, good point about the behavior being different — could be `save` for the execution path case, `cache` for the one based on runtime values, if we end up adopting the saving/caching behavior you suggest.

> Right, so we haven't mentioned deleting at all- if runs are going to be invalidated then that's definitely worth documenting and renaming from save.

Ah — in your spec I assumed "invalidate" meant "delete", but you're right saved values could be invalidated without actually deleting them from disk.

To bring it all together, we've proposed ...

1. An API that saves to disk (and loads from disk) values based on ancestors + code, ie based on the execution path. I like your spec. When a saved value is "invalidated", we need to decide whether the saved artifact remains on disk. 
2. Maybe an API for caching function values to disk.
3. Maybe overloading the API from 1 to handle 2 as well, but I see your point now why the two might be better separate.
4. An option to bake "saving" into the runtime on a cell level — could be controlled via an app config. In this case saved values would likely be deleted from disk upon invalidation.

I'm all for some version of these features in marimo. How do you think we should proceed?
Here's a proposed rollout:

1.) Execution-path first with the `with` statement, to remove any ambiguity w.r.t. the types of caching (the "function" cache doesn't really make sense in this context)
```python
with mo.experimental.cached_context(name="expensive_calls", ...):
    x = expensive_call()
    y = another_expensive_call(x)
# x and y are "cached", there's a python workaround that can just conditionally skip this block, and populate x y from load
...
def cached_context(
                  pin_marimo_lib : bool = False, # could also go by "major" | "minor" | "patch" | None; default minor here on the safe side
                  pin_sys_libs : bool = False, # major
                  pin_local_libs : bool = True, # patch
                  max_size=None,
                  save_path : str = ".marimo_outputs", # or just outputs
                  *,
                  name : str):
    """
    same as above, just reduced. I think the load and sign fns were premature optimization
    additionally "max_size=None" which will specify the max number of saves before LRU cleanup.
    """
```

2.a) This will unlock cell level save / cache (basically the same as wrapping the whole block in a `with`, but can have tighter integration + would need some UI work / a CellConfig attribute)

2.b Can be done concurrently or before 2.a) Function level calls, `mo.persistent_cache` as a drop-in replacement for `functools.cache`, uses the save to disk infrastructure that 1. lays down, and more sophisticated argument checking

---

Future
- Get cache to work on wasm (load over HTTP vs from disk)
- Consider `sign` and `load`

Commit message: tidy: typing and comments

Patch to review:
diff --git a/marimo/_save/__init__.py b/marimo/_save/__init__.py
new file mode 100644
index 00000000000..be52407b506
--- /dev/null
+++ b/marimo/_save/__init__.py
@@ -0,0 +1 @@
+# Copyright 2024 Marimo. All rights reserved.
diff --git a/marimo/_save/ast.py b/marimo/_save/ast.py
new file mode 100644
index 00000000000..f3dd3f0e66b
--- /dev/null
+++ b/marimo/_save/ast.py
@@ -0,0 +1,123 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import ast
+from typing import Any, Sequence, Union
+
+
+class BlockException(Exception):
+    pass
+
+
+def compiled_ast(block: Sequence[Union[ast.AST | ast.stmt]]) -> ast.Module:
+    return compile(
+        ast.Module(block, type_ignores=[]),
+        # <ast> is non-standard as a filename, but easier to debug than
+        # <module> everywhere.
+        "<ast>",
+        mode="exec",
+        flags=ast.PyCF_ONLY_AST | ast.PyCF_ALLOW_TOP_LEVEL_AWAIT,
+        optimize=0,
+        dont_inherit=True,
+    )
+
+
+def clean_to_modules(
+    pre_block: list[ast.AST], block: ast.With
+) -> tuple[ast.Module, ast.Module]:
+    """Standardizes a `with` block to modules.
+
+    Consider
+
+    >>> <pre_block>
+    >>> ...
+    >>> </pre_block>
+    >>> with fn(val:=other) as x:
+    >>>   <block>
+    >>>   ...
+    >>>   </block>
+
+    We want to compile <pre_block>, and <block> into two separate modules,
+    however, line 4 removes the context of "x" and "val", so this is adjusted
+    to:
+
+    >>> <pre_block>
+    >>> ...
+    >>> x = fn(val := other)
+    >>> </pre_block>
+    >>> <block>
+    >>> ...
+    >>> </block>
+    """
+    assert len(block.items) == 1, "Unexpected with block structure."
+    (with_block,) = block.items
+    initializer: ast.AST = with_block.context_expr
+    if with_block.optional_vars:
+        initializer = ast.Assign(
+            targets=[with_block.optional_vars],
+            value=initializer,
+        )
+    else:
+        # Edgecase with no "as" clause.
+        initializer = ast.Expr(value=initializer)
+    initializer.lineno = len(pre_block) + 1
+    initializer.col_offset = 0
+    pre_block.append(initializer)
+    return (compiled_ast(pre_block), compiled_ast(block.body))
+
+
+class ExtractWithBlock(ast.NodeTransformer):
+    def __init__(self, line: int, *arg: Any, **kwargs: Any) -> None:
+        super().__init__(*arg, **kwargs)
+        self.target_line = line
+
+    def generic_visit(self, node: ast.AST) -> tuple[ast.Module, ast.Module]:
+        pre_block = []
+
+        # There are a few strange edges cases like:
+        # >>> with open("file.txt"): pass
+        #
+        # It's difficult to properly delineate the block, so if multiple
+        # things map to the same line number, it's best to throw an error.
+        on_line = []
+        previous = None
+
+        assert isinstance(node, list), "Unexpected block structure."
+        for n in node:
+            if n.lineno < self.target_line:
+                pre_block.append(n)
+                previous = n
+            elif n.lineno == self.target_line:
+                on_line.append(n)
+            # The target line can easily be skipped if there are comments or
+            # white space or if the block is contained within another block.
+            else:
+                break
+
+        # Capture the edge case when the block is contained within another.
+        # These cases are explicitly restricted to If and other With blocks,
+        # excluding by omission try, for, classes and functions.
+        if len(on_line) == 0:
+            if isinstance(previous, (ast.With, ast.If)):
+                try:
+                    # Recursion by referring the the containing block also
+                    # captures the case where the target line number was not
+                    # exactly hit.
+                    return ExtractWithBlock(self.target_line).generic_visit(
+                        previous.body
+                    )
+                except BlockException:
+                    on_line.append(previous)
+            else:
+                raise BlockException(
+                    "persistent_cache cannot be invoked within a block "
+                    "(try moving the block within the persistent_cache scope)."
+                )
+        elif len(on_line) == 1:
+            assert isinstance(on_line[0], ast.With), "Unexpected block."
+            return clean_to_modules(pre_block, on_line[0])
+        # It should be possible to relate the lines with the AST,
+        # but reduce potential bugs by just throwing an error.
+        raise BlockException(
+            "Saving on a shared line may lead to unexpected behavior."
+        )
diff --git a/marimo/_save/cache.py b/marimo/_save/cache.py
new file mode 100644
index 00000000000..c21e9cea1bf
--- /dev/null
+++ b/marimo/_save/cache.py
@@ -0,0 +1,41 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import re
+from collections import namedtuple
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal
+
+from marimo._runtime.context import get_context
+
+if TYPE_CHECKING:
+    from marimo._ast.visitor import Name
+
+CacheType = Literal["ContentAddressed", "ExecutionPath", "Unknown"]
+# Easy visual identification of cache type.
+CACHE_PREFIX: dict[CacheType, str] = {
+    "ContentAddressed": "C_",
+    "ExecutionPath": "E_",
+    "Unknown": "U_",
+}
+
+ValidCacheSha = namedtuple("ValidCacheSha", ("sha", "cache_type"))
+
+
+@dataclass
+class Cache:
+    defs: dict[Name, Any]
+    hash: str
+    cache_type: CacheType
+    hit: bool
+
+
+def contextual_defs(cache: Cache) -> dict[tuple[Name, Name], Any]:
+    """Uses context to resolve private variable names."""
+    context = get_context().execution_context
+    assert context is not None, "Context could not be resolved"
+    private_prefix = f"_cell_{context.cell_id}"
+    return {
+        (var, re.sub(r"^_", private_prefix, var)): value
+        for var, value in cache.defs.items()
+    }
diff --git a/marimo/_save/hash.py b/marimo/_save/hash.py
new file mode 100644
index 00000000000..8f705206c3b
--- /dev/null
+++ b/marimo/_save/hash.py
@@ -0,0 +1,134 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import ast
+import base64
+import hashlib
+import numbers
+import types
+from typing import TYPE_CHECKING, Optional
+
+from marimo._ast.visitor import ScopedVisitor
+from marimo._save.cache import Cache, ValidCacheSha
+
+if TYPE_CHECKING:
+    from types import CodeType
+
+    from marimo._ast.cell import CellId_t, CellImpl
+    from marimo._runtime.dataflow import DirectedGraph
+    from marimo._save.loaders import Loader
+
+BASE_PRIMITIVES = (str, numbers.Number, type(None))
+
+# marimo: __marimo__.__version__
+# Module refs
+# UI refs
+# get_context().ui_element_registry._bindings
+# State refs
+# Need the same for state
+
+# depth == 1
+
+# "pure" Function refs
+# primitive refs
+# other
+
+
+def hash_module(code: Optional[CodeType]) -> bytes:
+    if not code:
+        return b"0" * 32
+
+    sha = hashlib.sha256()
+
+    def process(code_obj: CodeType) -> None:
+        # Recursively hash the constants that are also code objects
+        for const in code_obj.co_consts:
+            if isinstance(const, types.CodeType):
+                process(const)
+            else:
+                sha.update(str(const).encode("utf8"))
+        # Concatenate the names and bytecode of the current code object
+        # Will cause invalidation of variable naming at the top level
+        sha.update(bytes("|".join(code_obj.co_names), "utf8"))
+        sha.update(code_obj.co_code)
+
+    process(code)
+    return sha.digest()
+
+
+def hash_raw_module(module: ast.Module) -> bytes:
+    return hash_module(
+        compile(
+            module,
+            "<hash>",
+            mode="exec",
+            flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT,
+        )
+    )
+
+
+def hash_cell_impl(cell: CellImpl) -> bytes:
+    return hash_module(cell.body) + hash_module(cell.last_expr)
+
+
+def build_execution_hash(
+    graph: DirectedGraph, cell_id: CellId_t
+) -> Optional[ValidCacheSha]:
+    sha = hashlib.sha256()
+    ancestors = graph.ancestors(cell_id)
+    references = sorted(
+        [hash_cell_impl(graph.cells[cell_id]) for cell_id in ancestors]
+    )
+    for code in references:
+        sha.update(code)
+    return ValidCacheSha(sha, "ExecutionPath")
+
+
+def build_content_hash(
+    graph: DirectedGraph, visitor: ScopedVisitor
+) -> Optional[ValidCacheSha]:
+    sha = hashlib.sha256()
+    for ref in sorted(
+        graph.get_transitive_references(visitor.defs, inclusive=False)
+    ):
+        if ref not in globals():
+            if ref in __builtins__:
+                continue
+            return None
+        else:
+            value = globals()[ref]
+        if isinstance(value, BASE_PRIMITIVES):
+            sha.update(str(value).encode("utf8"))
+            continue
+        return None
+    return ValidCacheSha(sha, "ContentAddressed")
+
+
+def hash_context(
+    module: ast.Module,
+    graph: DirectedGraph,
+    cell_id: CellId_t,
+    *,
+    context: Optional[ast.Module] = None,
+    loader: Loader,
+) -> Cache:
+    # Empty name, so we can match and fill in cell context on load.
+    visitor = ScopedVisitor("")
+    visitor.visit(module)
+
+    # Attempt content hash
+    valid_cache_sha = build_content_hash(graph, visitor)
+    if not valid_cache_sha:
+        # Execution path hash
+        valid_cache_sha = build_execution_hash(graph, cell_id)
+        assert valid_cache_sha
+        if context:
+            valid_cache_sha.sha.update(hash_raw_module(context))
+
+    sha, cache_type = valid_cache_sha
+    sha.update(hash_raw_module(module))
+    hashed_context = (
+        base64.urlsafe_b64encode(sha.digest()).decode("utf-8").strip("=")
+    )
+
+    return loader.cache_attempt(visitor.defs, hashed_context, cache_type)
diff --git a/marimo/_save/loaders.py b/marimo/_save/loaders.py
new file mode 100644
index 00000000000..8af1c198d5e
--- /dev/null
+++ b/marimo/_save/loaders.py
@@ -0,0 +1,88 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import os
+import pickle
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from marimo._save.cache import CACHE_PREFIX, Cache, CacheType
+
+if TYPE_CHECKING:
+    from marimo._ast.visitor import Name
+
+
+class Loader(ABC):
+    """Loaders are responsible for saving and loading persistent caches.
+
+    Loaders are provided a name, a save path and a cache key or "hash", which
+    should be deterministically determined given the notebook context.
+
+    In the future, they may be specialized for different types of data (such as
+    numpy or pandas dataframes), or remote storage (such as S3 or marimo
+    cloud).
+    """
+
+    def __init__(self, name: str, save_path: str) -> None:
+        self.name = name
+        self.save_path = Path(save_path) / name
+        self.save_path.mkdir(parents=True, exist_ok=True)
+
+    def build_path(self, hashed_context: str, cache_type: CacheType) -> Path:
+        prefix = CACHE_PREFIX.get(cache_type, "U_")
+        return self.save_path / f"{prefix}{hashed_context}"
+
+    def cache_attempt(
+        self, defs: set[Name], hashed_context: str, cache_type: CacheType
+    ) -> Cache:
+        if not self.cache_hit(hashed_context, cache_type):
+            return Cache(
+                {d: None for d in defs}, hashed_context, cache_type, False
+            )
+        loaded = self.load_cache(hashed_context, cache_type)
+        # TODO: Consider more robust verification
+        assert loaded.hash == hashed_context
+        assert set(defs) == set(loaded.defs)
+        return Cache(loaded.defs, hashed_context, cache_type, True)
+
+    @abstractmethod
+    def cache_hit(self, hashed_context: str, cache_type: CacheType) -> bool:
+        """Check if cache has been hit given a result hash.
+
+        Args:
+            hashed_context: The hash of the result context
+            cache_type: The type of cache to check for
+
+        Returns:
+            bool: Whether the cache has been hit
+        """
+
+    @abstractmethod
+    def load_cache(self, hashed_context: str, cache_type: CacheType) -> Cache:
+        """Load Cache"""
+
+    @abstractmethod
+    def save_cache(self, cache: Cache) -> None:
+        """Save Cache"""
+
+
+class PickleLoader(Loader):
+    """General loader for serialable objects."""
+
+    def build_path(self, hashed_context: str, cache_type: CacheType) -> Path:
+        return Path(f"{super().build_path(hashed_context, cache_type)}.pickle")
+
+    def cache_hit(self, hashed_context: str, cache_type: CacheType) -> bool:
+        return os.path.exists(self.build_path(hashed_context, cache_type))
+
+    def load_cache(self, hashed_context: str, cache_type: CacheType) -> Cache:
+        assert self.cache_hit(hashed_context, cache_type)
+        with open(self.build_path(hashed_context, cache_type), "rb") as handle:
+            cache = pickle.load(handle)
+            assert isinstance(cache, Cache), "Excepted cache object"
+            return cache
+
+    def save_cache(self, cache: Cache) -> None:
+        with open(self.build_path(cache.hash, cache.cache_type), "wb") as f:
+            pickle.dump(cache, f, protocol=pickle.HIGHEST_PROTOCOL)
diff --git a/marimo/_save/save.py b/marimo/_save/save.py
new file mode 100644
index 00000000000..e100182e8a4
--- /dev/null
+++ b/marimo/_save/save.py
@@ -0,0 +1,161 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import ast
+import sys
+import traceback
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Optional,
+    Self,
+    Type,
+    Union,
+)
+
+from marimo._runtime.context import get_context
+from marimo._save.ast import ExtractWithBlock
+from marimo._save.cache import Cache, contextual_defs
+from marimo._save.hash import hash_context
+from marimo._save.loaders import PickleLoader
+
+if TYPE_CHECKING:
+    from types import FrameType, TracebackType
+
+    from _typeshed import TraceFunction
+
+
+class SkipWithBlock(Exception):
+    """Special exception to get around executing the with block body."""
+
+
+class CacheException(BaseException):
+    pass
+
+
+class persistent_cache(object):
+    """Context block for cache lookup of a block of code.
+
+    Example usage:
+
+    >>> with persistent_cache(name="my_cache"):
+    >>>     variable = expensive_function() # This will be cached.
+
+    For an implementation sibling regarding the block skipping, see `withhacks`
+    in pypi.
+
+    NB: Since context abuses sys frame trace, this may conflict with debugging
+    tools that also use sys.settrace.
+    """
+
+    def __init__(
+        self,
+        *,
+        save_path: str = "outputs",
+        name: str,
+    ) -> None:
+        # TODO: consider construction injection
+        self._loader = PickleLoader(name, save_path)
+        self.name = name
+
+        self._skipped = True
+        self._cache: Optional[Cache] = None
+        self._entered_trace = False
+        self._old_trace: Optional[TraceFunction] = None
+
+    def __enter__(self) -> Self:
+        sys.settrace(lambda *_args, **_keys: None)
+        frame = sys._getframe(1)
+        # Attempt to hold on to the previous trace.
+        self._old_trace = frame.f_trace
+        # Setting the frametrace, will cause the function to be run on _every_
+        # single context call until the trace is cleared.
+        frame.f_trace = self.trace
+        return self
+
+    def trace(
+        self, _frame: FrameType, _event: str, _arg: Any
+    ) -> Union[TraceFunction | None]:
+        # General flow is as follows:
+        #   1) Follow the stack trace backwards to the first instance of a
+        # "<module>" function call, which corresponds to a cell level block.
+        #   2) Run static analysis to determine whether the call meets our
+        # criteria. The procedure is a little brittle as such, certain contexts
+        # are not allow (e.g. called within a function or a loop).
+        #  3) Hash the execution and lookup the cache, and return!
+        #  otherwise) Set _skipped such that the block continues to execute.
+
+        self._entered_trace = True
+        stack = traceback.extract_stack()
+        if not self._skipped:
+            return self._old_trace
+
+        # This only executes on the first line of code in the block. If the
+        # cache is hit, the block terminates early with a SkipWithBlock
+        # exception, if the block is not hit, self._skipped is set to False,
+        # causing this function to terminate before reaching this block.
+        for i, frame in enumerate(stack[::-1]):
+            _filename, lineno, function_name, _code = frame
+            if function_name == "<module>":
+                ctx = get_context()
+                assert (
+                    ctx.execution_context is not None
+                ), "Could not resolve context for cache."
+                graph = ctx.graph
+                cell_id = ctx.execution_context.cell_id
+                pre_module, save_module = ExtractWithBlock(lineno - 1).visit(
+                    ast.parse(ctx.graph.cells[cell_id].code).body
+                )
+
+                self._cache = hash_context(
+                    save_module,
+                    graph,
+                    cell_id,
+                    loader=self._loader,
+                    context=pre_module,
+                )
+                if self._cache.hit:
+                    raise SkipWithBlock()
+
+                self.cache_type = self._cache
+                self._skipped = False
+                return self._old_trace
+            elif i > 1:
+                raise CacheException(
+                    "persistent_cache must be invoked from cell level "
+                    "(cannot be in a function or class)"
+                )
+        raise CacheException("persistent_cache could not resolve block")
+
+    def __exit__(
+        self,
+        exception: Optional[Type[BaseException]],
+        instance: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> bool:
+        sys.settrace(self._old_trace)  # Clear to previous set trace.
+        if not self._entered_trace:
+            raise CacheException("Unexpected block format.")
+
+        # Backfill the loaded values into global scope.
+        if self._cache and self._cache.hit:
+            for lookup, var in contextual_defs(self._cache):
+                globals()[var] = self._cache.defs[lookup]
+            # Return true to suppress the SkipWithBlock exception.
+            return True
+
+        # NB: exception is a type.
+        if exception:
+            assert not isinstance(
+                instance, SkipWithBlock
+            ), "Cache was not correctly set."
+            if isinstance(instance, BaseException):
+                raise instance
+            raise exception
+
+        # Fill the cache object and save.
+        assert self._cache is not None
+        for lookup, var in contextual_defs(self._cache):
+            self._cache.defs[lookup] = globals()[var]
+        self._loader.save_cache(self._cache)
+        return False


Code review comments:
diff_hunk: @@ -0,0 +1,161 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import ast
+import sys
+import traceback
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Optional,
+    Self,
+    Type,
+    Union,
+)
+
+from marimo._runtime.context import get_context
+from marimo._save.ast import ExtractWithBlock
+from marimo._save.cache import Cache, contextual_defs
+from marimo._save.hash import hash_context
+from marimo._save.loaders import PickleLoader
+
+if TYPE_CHECKING:
+    from types import FrameType, TracebackType
+
+    from _typeshed import TraceFunction
+
+
+class SkipWithBlock(Exception):
+    """Special exception to get around executing the with block body."""
+
+
+class CacheException(BaseException):
+    pass
+
+
+class persistent_cache(object):
+    """Context block for cache lookup of a block of code.
+
+    Example usage:
+
+    >>> with persistent_cache(name="my_cache"):
+    >>>     variable = expensive_function() # This will be cached.
+
+    For an implementation sibling regarding the block skipping, see `withhacks`
+    in pypi.
+
+    NB: Since context abuses sys frame trace, this may conflict with debugging
+    tools that also use sys.settrace.
+    """
+
+    def __init__(
+        self,
+        *,
+        save_path: str = "outputs",
+        name: str,
+    ) -> None:
+        # TODO: consider construction injection
+        self._loader = PickleLoader(name, save_path)
+        self.name = name
+
+        self._skipped = True
+        self._cache: Optional[Cache] = None
+        self._entered_trace = False
+        self._old_trace: Optional[TraceFunction] = None
+
+    def __enter__(self) -> Self:
+        sys.settrace(lambda *_args, **_keys: None)
+        frame = sys._getframe(1)
+        # Attempt to hold on to the previous trace.
+        self._old_trace = frame.f_trace
+        # Setting the frametrace, will cause the function to be run on _every_
+        # single context call until the trace is cleared.
+        frame.f_trace = self.trace
+        return self
+
+    def trace(
+        self, _frame: FrameType, _event: str, _arg: Any
+    ) -> Union[TraceFunction | None]:
+        # General flow is as follows:
+        #   1) Follow the stack trace backwards to the first instance of a
+        # "<module>" function call, which corresponds to a cell level block.
+        #   2) Run static analysis to determine whether the call meets our
+        # criteria. The procedure is a little brittle as such, certain contexts
+        # are not allow (e.g. called within a function or a loop).
+        #  3) Hash the execution and lookup the cache, and return!
+        #  otherwise) Set _skipped such that the block continues to execute.
+
+        self._entered_trace = True
+        stack = traceback.extract_stack()
+        if not self._skipped:
+            return self._old_trace
+
+        # This only executes on the first line of code in the block. If the
+        # cache is hit, the block terminates early with a SkipWithBlock
+        # exception, if the block is not hit, self._skipped is set to False,
+        # causing this function to terminate before reaching this block.
+        for i, frame in enumerate(stack[::-1]):
+            _filename, lineno, function_name, _code = frame
+            if function_name == "<module>":
+                ctx = get_context()
+                assert (
+                    ctx.execution_context is not None
+                ), "Could not resolve context for cache."
+                graph = ctx.graph
+                cell_id = ctx.execution_context.cell_id
+                pre_module, save_module = ExtractWithBlock(lineno - 1).visit(
+                    ast.parse(ctx.graph.cells[cell_id].code).body
+                )
+
+                self._cache = hash_context(
+                    save_module,
+                    graph,
+                    cell_id,
+                    loader=self._loader,
+                    context=pre_module,
+                )
+                if self._cache.hit:
+                    raise SkipWithBlock()
+
+                self.cache_type = self._cache
+                self._skipped = False
+                return self._old_trace
+            elif i > 1:
+                raise CacheException(
+                    "persistent_cache must be invoked from cell level "
+                    "(cannot be in a function or class)"
+                )
+        raise CacheException("persistent_cache could not resolve block")
path: marimo/_save/save.py
text: @user1:
Are there valid cases in which this would be hit, or would it be considered a bug in marimo if this exception were raised?

If latter, the error message could be "persistent_cache could not resolve block — this is a bug in marimo. Please file an issue at https://github.com/marimo-team/marimo/issues"

diff_hunk: @@ -0,0 +1,161 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import ast
+import sys
+import traceback
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Optional,
+    Self,
+    Type,
+    Union,
+)
+
+from marimo._runtime.context import get_context
+from marimo._save.ast import ExtractWithBlock
+from marimo._save.cache import Cache, contextual_defs
+from marimo._save.hash import hash_context
+from marimo._save.loaders import PickleLoader
+
+if TYPE_CHECKING:
+    from types import FrameType, TracebackType
+
+    from _typeshed import TraceFunction
+
+
+class SkipWithBlock(Exception):
+    """Special exception to get around executing the with block body."""
+
+
+class CacheException(BaseException):
+    pass
+
+
+class persistent_cache(object):
+    """Context block for cache lookup of a block of code.
+
+    Example usage:
+
+    >>> with persistent_cache(name="my_cache"):
+    >>>     variable = expensive_function() # This will be cached.
+
+    For an implementation sibling regarding the block skipping, see `withhacks`
+    in pypi.
+
+    NB: Since context abuses sys frame trace, this may conflict with debugging
+    tools that also use sys.settrace.
+    """
+
+    def __init__(
+        self,
+        *,
+        save_path: str = "outputs",
+        name: str,
+    ) -> None:
+        # TODO: consider construction injection
+        self._loader = PickleLoader(name, save_path)
+        self.name = name
+
+        self._skipped = True
+        self._cache: Optional[Cache] = None
+        self._entered_trace = False
+        self._old_trace: Optional[TraceFunction] = None
+
+    def __enter__(self) -> Self:
+        sys.settrace(lambda *_args, **_keys: None)
+        frame = sys._getframe(1)
+        # Attempt to hold on to the previous trace.
+        self._old_trace = frame.f_trace
+        # Setting the frametrace, will cause the function to be run on _every_
+        # single context call until the trace is cleared.
+        frame.f_trace = self.trace
+        return self
+
+    def trace(
+        self, _frame: FrameType, _event: str, _arg: Any
+    ) -> Union[TraceFunction | None]:
+        # General flow is as follows:
+        #   1) Follow the stack trace backwards to the first instance of a
+        # "<module>" function call, which corresponds to a cell level block.
+        #   2) Run static analysis to determine whether the call meets our
+        # criteria. The procedure is a little brittle as such, certain contexts
+        # are not allow (e.g. called within a function or a loop).
+        #  3) Hash the execution and lookup the cache, and return!
+        #  otherwise) Set _skipped such that the block continues to execute.
+
+        self._entered_trace = True
+        stack = traceback.extract_stack()
+        if not self._skipped:
+            return self._old_trace
+
+        # This only executes on the first line of code in the block. If the
+        # cache is hit, the block terminates early with a SkipWithBlock
+        # exception, if the block is not hit, self._skipped is set to False,
+        # causing this function to terminate before reaching this block.
+        for i, frame in enumerate(stack[::-1]):
+            _filename, lineno, function_name, _code = frame
+            if function_name == "<module>":
+                ctx = get_context()
+                assert (
+                    ctx.execution_context is not None
+                ), "Could not resolve context for cache."
+                graph = ctx.graph
+                cell_id = ctx.execution_context.cell_id
+                pre_module, save_module = ExtractWithBlock(lineno - 1).visit(
+                    ast.parse(ctx.graph.cells[cell_id].code).body
+                )
+
+                self._cache = hash_context(
+                    save_module,
+                    graph,
+                    cell_id,
+                    loader=self._loader,
+                    context=pre_module,
+                )
+                if self._cache.hit:
+                    raise SkipWithBlock()
+
+                self.cache_type = self._cache
+                self._skipped = False
+                return self._old_trace
+            elif i > 1:
+                raise CacheException(
+                    "persistent_cache must be invoked from cell level "
+                    "(cannot be in a function or class)"
+                )
+        raise CacheException("persistent_cache could not resolve block")
+
+    def __exit__(
+        self,
+        exception: Optional[Type[BaseException]],
+        instance: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> bool:
+        sys.settrace(self._old_trace)  # Clear to previous set trace.
+        if not self._entered_trace:
+            raise CacheException("Unexpected block format.")
+
+        # Backfill the loaded values into global scope.
+        if self._cache and self._cache.hit:
+            for lookup, var in contextual_defs(self._cache):
+                globals()[var] = self._cache.defs[lookup]
+            # Return true to suppress the SkipWithBlock exception.
+            return True
+
+        # NB: exception is a type.
+        if exception:
+            assert not isinstance(
+                instance, SkipWithBlock
+            ), "Cache was not correctly set."
+            if isinstance(instance, BaseException):
+                raise instance
+            raise exception
+
+        # Fill the cache object and save.
+        assert self._cache is not None
+        for lookup, var in contextual_defs(self._cache):
+            self._cache.defs[lookup] = globals()[var]
+        self._loader.save_cache(self._cache)
path: marimo/_save/save.py
text: @user1:
What's the behavior if the loader fails to serialize the object? For example not all objects are pickle-able.

@author:
That's a good question. marimo doesn't really have a "warning", This could print to stderr that cache save did not occur. I don't think failure is completely correct.

Thoughts?

@user1:
Print to stderr sounds good to me.

diff_hunk: @@ -0,0 +1,88 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import os
+import pickle
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from marimo._save.cache import CACHE_PREFIX, Cache, CacheType
+
+if TYPE_CHECKING:
+    from marimo._ast.visitor import Name
+
+
+class Loader(ABC):
+    """Loaders are responsible for saving and loading persistent caches.
+
+    Loaders are provided a name, a save path and a cache key or "hash", which
+    should be deterministically determined given the notebook context.
+
+    In the future, they may be specialized for different types of data (such as
+    numpy or pandas dataframes), or remote storage (such as S3 or marimo
+    cloud).
+    """
+
+    def __init__(self, name: str, save_path: str) -> None:
+        self.name = name
+        self.save_path = Path(save_path) / name
+        self.save_path.mkdir(parents=True, exist_ok=True)
+
+    def build_path(self, hashed_context: str, cache_type: CacheType) -> Path:
+        prefix = CACHE_PREFIX.get(cache_type, "U_")
+        return self.save_path / f"{prefix}{hashed_context}"
+
+    def cache_attempt(
+        self, defs: set[Name], hashed_context: str, cache_type: CacheType
+    ) -> Cache:
+        if not self.cache_hit(hashed_context, cache_type):
+            return Cache(
+                {d: None for d in defs}, hashed_context, cache_type, False
+            )
+        loaded = self.load_cache(hashed_context, cache_type)
+        # TODO: Consider more robust verification
+        assert loaded.hash == hashed_context
+        assert set(defs) == set(loaded.defs)
+        return Cache(loaded.defs, hashed_context, cache_type, True)
+
+    @abstractmethod
+    def cache_hit(self, hashed_context: str, cache_type: CacheType) -> bool:
+        """Check if cache has been hit given a result hash.
+
+        Args:
+            hashed_context: The hash of the result context
+            cache_type: The type of cache to check for
+
+        Returns:
+            bool: Whether the cache has been hit
+        """
+
+    @abstractmethod
+    def load_cache(self, hashed_context: str, cache_type: CacheType) -> Cache:
+        """Load Cache"""
+
+    @abstractmethod
+    def save_cache(self, cache: Cache) -> None:
+        """Save Cache"""
+
+
+class PickleLoader(Loader):
+    """General loader for serialable objects."""
path: marimo/_save/loaders.py
text: @user1:
nit
```suggestion
    """General loader for serializable objects."""
```