Cache interface options

[andrewhill157]

I was looking more closely into the cache options added recently and saw that you'd added mo.cache and mo.persistent_cache! Sorry if this overlaps with past discussion that I've missed somewhere!

My understanding is that the former is used like

@mo.cache
def my_function():
    sleep(5)
    return 0

whereas the latter can be used like:

with mo.persistent_cache("my_cache"):
    x = expensive_thing()

but can't be used to decorate functions and the docs point you to use of @mo.cache for that.

In looking around at other options I came across joblib.Memory which can be used to decorate functions like @mo.cache (or the various functools cache options) like:

from joblib import Memory
mem = Memory("my_cache")

@mem.cache
def my_function():
    sleep(5)
    return 0

but with a persistent cache.

I don't think you can directly use it as a context manager but you could obvs make something to do that pretty easily.

One thing I liked about this was that it keeps the interface when using a cache vs. a persistent cache very similar. Not sure if you would (or already have) consider anything like this as an alternative or extension of the ways marimo currently supports caching.

joblib.Memory works fine for my purposes if decorating a function defined in a module outside of a notebook, but it doesn't seem to persistently cache if the function is defined inside the notebook. Seems like that is probably a known issue:
joblib/joblib#1498

[akshayka]

@andrewhill157 , thanks for opening this discussion and so sorry for the late response.

You're correct that we have an API asymmetry between mo.cache and mo.persistent_cache, one is a decorator and the other a context manager.

One thing I liked about this was that it keeps the interface when using a cache vs. a persistent cache very similar. Not sure if you would (or already have) consider anything like this as an alternative or extension of the ways marimo currently supports caching.

I see the appeal of having the same interface for both. cc @dmadisetti — is that something that could be possible in the future?

Separately — Andrew, I'm curious if you've gotten a chance to use our caching features, and if you had any feedback

[dmadisetti]

This would be very easy to implement, they share the same underlying API (we'd just swap out MemoryLoader for PickleLoader); the hardest part is just thinking how to expose the functionality.

I think mixing the namespace to be a decorator and a context manager will make this a little difficult from a maintenance perspective and might be a little bit confusing in documentation. But maybe there's a certain elegance to it? Others have talked about wanting other save formats than pickle as well, so I think an API extension is due. With a shared namespace we might have:

from marimo import memory_cache, pickle_cache, numpy_cache

@memory_cache
def fn():
    ...

@remote_cache(endpoint="example.api")
def cached():
    ...

with numpy_cache("expensive_block", compression=True):
    ...

For implementation, I see this working by switching on the constructor arguments types- which is messy but maybe fine.
Conversely,

from marimo.save import MemoryLoader, RemoteLoader, NumpyLoader

# `persistent_cache` nomenclature may be misleading if the Loader doesn't actually write to disk
with MemoryLoader.context("my_cache") as cache:
    x = expensive_thing()


@RemoteLoader.cache(endpoint="example.api")
def fn():
    ...

with NumpyLoader.context("expensive_block", compression=True):
    ...

open to more thoughts too

[akshayka]

Others have talked about wanting other save formats than pickle as well, so I think an API extension is due.

Do we need to expose other save formats, or can we just handle this for users behind-the-scenes? The latter would be my preference. From the code sample it's not clear to me what numpy_cache would do (only cache NumPy arrays, and nothing else? Not sure we need that).

[dmadisetti]

npy and npz can save basic structures, just given as an example here, because npz has better compression than pickle. But there are other formats like dill, which are more expansive than pickle.

Choosing a file format behinds the scenes could lead to unexpected cache misses on dependency additions or "behind the scenes changes". I think it's OK to be explicit and let the user choose the save format, knowing the limitations of the chosen format.

For instance, a parquet saver would be great for saving dataframes, but won't work for general python objects

marimo could try to Frankenstein multiple save types together, but I think this makes the system more opaque and brittle. Allowing the user to choose the format type based on its strength (or weakness), and optimize their code for the given format, seems like best option

But we should still keep the general cache or persistent_cache for general cases

But open for discussion- how would a behind-the-scenes format best work?

[akshayka]

But we should still keep the general cache or persistent_cache for general cases

Yea that makes sense.

I think I'm just concerned about API scope creep. But I do see of course why more specific serialization types would be helpful.

I'm also a bit confused by the semantics. For

with NumpyLoader.context("expensive_block", compression=True):
    ...

if the code in ... created variables of several types, some of which were not serializable as .npy, would a RuntimeError be raised?

[dmadisetti]

Yes, which is the same as pickle right now

[andrewhill157]

@akshayka sorry for the late reply -- I've only used them in some simple test examples so not sure I have much feedback at this point other than the interface detail and some downstream limitations (or user error).

For example, in my case the function I want to cache is called within another function, which doesn't seem to play nice with marimo's implementation currently -- for example:

import marimo as mo
import time

def test(i):
    with mo.persistent_cache("my_cache"):
        time.sleep(5)
        result = 5
    return result * i

test(1)

yields a BlockException marimo._save.ast.BlockException: persistent_cache cannot be invoked within a block (try moving the block within the persistent_cache scope).

I also tried to make my own decorator but run into the same error when doing so.

I don't currently have particular requirements re: save formats (pickle is fine for current use case), but like the idea of maintaining a generic (or auto) default personally.

[dmadisetti]

So this is an intentional check, since we can't get the same guarantees of the execution path if the block is called from somewhere like in a loop.

My suggestion for something like this would be:

with persistent_cache("precomputed"):
    data = [expensive_call(i) for _ in range(i)]

Or just

with persistent_cache(f"compute_{i}"):
    data = expensive_call(i)

which will call on a change to i.

I get the use case, that you are calling expensive_call many times, in a variety of situations, so it doesn't fit into this format.
I'm still open to mixing the namespace, but I do want to leave room for other caching types

persistent_cache will likely not work with any user level decorator by design- if this is really pressing, you can swap out MemoryLoader for PickleLoader in marimo._save._cache_base- but I think it will be an easy feature to add once a api is agreed.

---

There is also this spec:

with mo.persistent_cache("block", loader=MemoryLoader) as cache:
    ...

@mo.cache(loader=PickleLoader):
def fn():
    ...

with mo.persistent_cache("block", loader="memory") as cache:
    ...

@mo.cache(loader="pickle") # This is what you would want
def fn():
    ...

Which I'm fine with, except that persistent_cache is potentially a misnomer (with particular loaders).
So maybe, keep persistent_cache behavior the same, and have a different function for explicit loaders:

with context_cache("block", loader="parquet"):
    ...

[akshayka]

@mo.cache(loader="pickle") # This is what you would want

Would this have the same semantics as persistent_cache?

So maybe, keep persistent_cache behavior the same, and have a different function for explicit loaders:

That makes sense, though what happens if some of the objects saved/loaded are not compatible with the selected loader? That's on the user I guess?

[dmadisetti]

Hash names should be the same, but I think this should be on the user

[akshayka]

Makes sense

[leventov]

As a general meta-note, I think would be nice to have a clearer conceptual picture about what we want from the notebook source files. Not just for developers, but more generally for users of marimo, and documented.

One simple distinction from Jupyter is already obvious from the system premise, that is, in virtue of being .py files, marimo notebooks don't store cell results. But there is still a fork on the road after that: do we want it to mostly "specification of behavior" and deployment concerns shoveled elsewhere, or we want it to bundle information about its own deployment (and caching, especially remote, is kinda a deployment concern), too. And a third option is perhaps to make it work both ways.

I don't think this is an easy question. There are good arguments that could be made for each option.

This question also interplays with another important API design dimension, namely the question of timeliness of cells (what does it mean for cells to have a time parameter, e.g. defaulting to now() but could be pegged for illustration?), and which is also an important deployment concern, if cells to be deployed in a kind of data processing DAG or executed at certain intervals.

[leventov]

Few thoughts:

1. I think there is probably little reason to configure storage formats per cell within a notebook. That is, if I want to use npy or Arrow or whatever compressed format for columns for storage, this expands to all cells within a notebook outputting columnar data. It's because the choice is motivated by space/network bandwidth usage -- load/store speed tradeoff, and this extends to the entire notebook as it's the unit of deployment.

So, maybe create App can be configured with Cacher (which in turn can be configured with different formats), and then there is @app.persist_results decorator. Indirection through App to save a cacher = Cacher() line if the user doesn't change the defaults.

2. Regarding specific formats, maybe use the approach of ORMs and other similar serialisers that permit configuring serialiser per type. That is, "numpy" as a kind of top level Cache doesn't make sense - it should be configured just for specific object types, whenever they appear. Others may default to pickle or whatever. So, it's more like Cacher((np.array, "npy"), (pd.dataframe, "arrow")), than NumpyCacher().

[leventov]

Actually, if Cacher is an App's field, @app.cell(persist_results=true) may be simpler, no need for a separate decorator.