From 45fbfdef2598bafdff237abc58c46e19166bca57 Mon Sep 17 00:00:00 2001 From: "dmadisetti@coreweave.com" Date: Wed, 24 Jun 2026 15:49:46 -0700 Subject: [PATCH 1/6] feat(save): LazyStore dual-mode backend + WASM export store layer Flesh out the LazyStore placeholder into a dual-mode store: - LazyStore (native): wraps an inner FileStore, tracks written/touched keys for export. - WasmLazyStore (Pyodide): writes to a shared in-session DictStore; reads fall through to concurrent HTTP fetch from notebook_location()/public/cache/, with path-traversal-safe keys and poisoned-key eviction on corrupt restore. - The single native/WASM decision is made once via a DualLoader registry entry (resolve_loader), so nothing downstream re-checks the platform. Adds WasmExportableStore (export_manifest tracking) + DictStore, the _ACTIVE_LAZY_LOADERS registry, and LazyLoader.flush_all() for export. Also folds in the unserializable-def robustness mechanism: rather than writing a placeholder blob, the loader marks the manifest Item with unserializable_type and reconstructs the UnhashableStub tripwire on load (from_item). --- marimo/_save/loaders/__init__.py | 33 +- marimo/_save/loaders/lazy.py | 526 +++++++++++++++++++--- marimo/_save/save.py | 6 +- marimo/_save/stores/__init__.py | 3 +- marimo/_save/stores/dict_store.py | 32 ++ marimo/_save/stores/store.py | 28 ++ marimo/_save/stubs/lazy_stub.py | 29 +- tests/_save/loaders/test_lazy_wasm.py | 294 ++++++++++++ tests/_save/loaders/test_loader.py | 89 ++++ tests/_save/stubs/test_lazy_codecs.py | 37 ++ tests/_save/stubs/test_unhashable_stub.py | 19 +- 11 files changed, 1009 insertions(+), 87 deletions(-) create mode 100644 marimo/_save/stores/dict_store.py create mode 100644 tests/_save/loaders/test_lazy_wasm.py create mode 100644 tests/_save/stubs/test_lazy_codecs.py diff --git a/marimo/_save/loaders/__init__.py b/marimo/_save/loaders/__init__.py index 55e6b5984a9..c603b7488e7 100644 --- a/marimo/_save/loaders/__init__.py +++ b/marimo/_save/loaders/__init__.py @@ -1,10 +1,11 @@ # Copyright 2026 Marimo. All rights reserved. from __future__ import annotations +from dataclasses import dataclass from typing import Literal from marimo._save.loaders.json import JsonLoader -from marimo._save.loaders.lazy import LazyLoader +from marimo._save.loaders.lazy import LazyLoader, WasmLazyLoader from marimo._save.loaders.loader import ( BasePersistenceLoader, Loader, @@ -13,18 +14,42 @@ ) from marimo._save.loaders.memory import MemoryLoader from marimo._save.loaders.pickle import PickleLoader +from marimo._utils.platform import is_pyodide LoaderKey = Literal["memory", "pickle", "json", "lazy"] -PERSISTENT_LOADERS: dict[LoaderKey, LoaderType] = { + +@dataclass(frozen=True) +class DualLoader: + """A loader registered as a native/WASM pair under one name. + + `resolve()` performs the *single* environment check and returns the + concrete loader class, so nothing downstream re-checks the platform. + Any loader can opt into dual behavior by registering one of these. + """ + + native: LoaderType + wasm: LoaderType + + def resolve(self) -> LoaderType: + return self.wasm if is_pyodide() else self.native + + +def resolve_loader(entry: LoaderType | DualLoader) -> LoaderType: + """Resolve a registry entry to a concrete loader class.""" + return entry.resolve() if isinstance(entry, DualLoader) else entry + + +PERSISTENT_LOADERS: dict[LoaderKey, LoaderType | DualLoader] = { "pickle": PickleLoader, "json": JsonLoader, - "lazy": LazyLoader, + "lazy": DualLoader(native=LazyLoader, wasm=WasmLazyLoader), } __all__ = [ "PERSISTENT_LOADERS", "BasePersistenceLoader", + "DualLoader", "JsonLoader", "LazyLoader", "Loader", @@ -33,4 +58,6 @@ "LoaderType", "MemoryLoader", "PickleLoader", + "WasmLazyLoader", + "resolve_loader", ] diff --git a/marimo/_save/loaders/lazy.py b/marimo/_save/loaders/lazy.py index 0e2824f5d0d..0a48b54dc4f 100644 --- a/marimo/_save/loaders/lazy.py +++ b/marimo/_save/loaders/lazy.py @@ -7,8 +7,8 @@ import queue import threading from enum import Enum, auto -from pathlib import Path -from typing import Any +from pathlib import Path, PurePosixPath +from typing import TYPE_CHECKING, Any import msgspec @@ -20,6 +20,10 @@ from marimo._save.hash import HashKey from marimo._save.loaders.loader import BasePersistenceLoader from marimo._save.stores import FileStore, Store +from marimo._save.stores.store import WasmExportableStore + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable, Iterator from marimo._save.stubs import ( ClassStub, FunctionStub, @@ -36,9 +40,13 @@ Item, Meta, ReferenceStub, + UnhashableStub, ) from marimo._save.stubs.stubs import mro_lookup +if TYPE_CHECKING: + from collections.abc import Callable + LOGGER = _loggers.marimo_logger() @@ -139,7 +147,13 @@ def to_item( ) -def from_item(item: Item) -> Any: +def from_item(item: Item, var_name: str = "") -> Any: + if item.unserializable_type is not None: + # No blob was written for this def — rebuild the tripwire in-memory + # from the manifest marker. + return UnhashableStub( + var_name=var_name, type_name=item.unserializable_type + ) if item.reference is not None: return ImmediateReferenceStub( ReferenceStub(item.reference, hash_value=item.hash or "") @@ -165,19 +179,197 @@ def from_item(item: Item) -> Any: return None -class LazyStore(FileStore): ... +# Module-level singleton DictStore for WASM. Shared across all LazyStore +# instances so cached data survives loader recreation (State GC, partial +# reconstruction, etc.). +_WASM_DICT_STORE: Store | None = None + + +def _get_wasm_dict_store() -> Store: + global _WASM_DICT_STORE + if _WASM_DICT_STORE is None: + from marimo._save.stores.dict_store import DictStore + + _WASM_DICT_STORE = DictStore() + return _WASM_DICT_STORE + + +# Keys whose deserialization failed — never re-fetched over HTTP. Survives +# across LazyStore instances. +_POISONED_KEYS: set[str] = set() + + +class LazyStore(WasmExportableStore): + """Native store for `LazyLoader`. + + Delegates to an inner store (default `FileStore` at `__marimo__/cache/`) + and tracks which keys were written or read so `--execute` export knows + exactly which blobs to bundle. Does no environment checks; the WASM + variant (`WasmLazyStore`) adds HTTP-backed reads. + """ + + def __init__(self, inner: Store | None = None) -> None: + self._inner = inner if inner is not None else FileStore() + self._written_keys: set[str] = set() + # Keys read this session. A warm re-export hits the cache rather + # than re-writing it, so the export manifest must cover reads too + # or the bundle ships incomplete. + self._touched_keys: set[str] = set() + + def get(self, key: str) -> bytes | None: + result = self._inner.get(key) + if result is not None: + self._touched_keys.add(key) + return result + + def put(self, key: str, value: bytes) -> bool: + self._written_keys.add(key) + return self._inner.put(key, value) + + def hit(self, key: str) -> bool: + result = self._inner.hit(key) + if result: + self._touched_keys.add(key) + return result + + def clear(self, key: str) -> bool: + self._written_keys.discard(key) + self._touched_keys.discard(key) + return self._inner.clear(key) + + def get_batch( + self, keys: Iterable[str] + ) -> Iterator[tuple[str, bytes | None]]: + for k in keys: + yield k, self.get(k) + + def export_manifest(self) -> list[str]: + return sorted(self._written_keys | self._touched_keys) + + +class WasmLazyStore(LazyStore): + """WASM store: writes to a shared in-session `DictStore`; reads fall + through to HTTP fetch from `notebook_location()/public/cache/`. + + Instantiated only by `WasmLazyLoader` (selected once via the dual-loader + registry), so it never needs to re-check the environment. + """ + + def __init__(self, inner: Store | None = None) -> None: + super().__init__( + inner if inner is not None else _get_wasm_dict_store() + ) + + def get(self, key: str) -> bytes | None: + result = super().get(key) + if result is not None: + return result + if key not in _POISONED_KEYS: + return self._http_get(key) + return None + + def get_batch( + self, keys: Iterable[str] + ) -> Iterator[tuple[str, bytes | None]]: + # Inner store first; HTTP-fetch only the (unpoisoned) misses, fired + # concurrently. + http_keys: list[str] = [] + for k in keys: + data = self._inner.get(k) + if data is not None: + self._touched_keys.add(k) + yield k, data + elif k not in _POISONED_KEYS: + http_keys.append(k) + else: + yield k, None + if http_keys: + yield from self._http_get_batch(http_keys) + + # -- HTTP internals -- + + def _base_url(self) -> str: + from marimo._runtime.runtime import notebook_location + + loc = notebook_location() + return f"{loc}/public/cache" if loc else "public/cache" + + @staticmethod + def _sanitize_key(key: str) -> str: + """Prevent path traversal in HTTP fetch keys.""" + clean = PurePosixPath(key) + if ".." in clean.parts or clean.is_absolute(): + raise ValueError(f"Invalid cache key: {key}") + return str(clean) + + def _http_get(self, key: str) -> bytes | None: + """Single sync fetch via pyodide_http-patched urllib.""" + import urllib.request + + key = self._sanitize_key(key) + url = f"{self._base_url()}/{key}" + try: + with urllib.request.urlopen(url) as resp: + return resp.read() if resp.status == 200 else None + except Exception: + return None + + def _http_get_batch( + self, keys: Iterable[str] + ) -> Iterator[tuple[str, bytes | None]]: + """Fire all fetches concurrently via JS fetch + asyncio.gather.""" + import asyncio + + from js import fetch # type: ignore + + base = self._base_url() + keys_list = [self._sanitize_key(k) for k in keys] + loop = asyncio.get_event_loop() + + async def _fetch_one(key: str) -> tuple[str, bytes | None]: + resp = await fetch(f"{base}/{key}") + if resp.ok: + buf = await resp.arrayBuffer() + return key, buf.to_bytes() + return key, None + + async def _fetch_all() -> list[tuple[str, bytes | None]]: + return await asyncio.gather(*(_fetch_one(k) for k in keys_list)) + + try: + results = loop.run_until_complete(_fetch_all()) + except Exception: + # run_until_complete on the live pyodide loop requires JSPI + # (WebAssembly stack switching), which e.g. Firefox lacks. Fall + # back to sequential synchronous XHR via the pyodide_http-patched + # urllib — legal in a worker. + results = [(k, self._http_get(k)) for k in keys_list] + yield from results + + +# Registry of active loaders by name, so a recreated loader reuses the same +# store (and so `flush_all`/export can reach every session loader). +_ACTIVE_LAZY_LOADERS: dict[str, LazyLoader] = {} class LazyLoader(BasePersistenceLoader): + # Default store class, overridden by the WASM variant. The single + # native/WASM decision is made once in the dual-loader registry, not here. + _store_cls: type[Store] = LazyStore + def __init__( self, name: str, store: Store | None = None, ) -> None: if store is None: - store = LazyStore() + # Reuse the same store across recreations of a named loader + # (State GC, partial reconstruction) so cached data survives. + prev = _ACTIVE_LAZY_LOADERS.get(name) + store = prev.store if prev is not None else self._store_cls() super().__init__(name, "jsonl", store) self._pending: list[threading.Thread] = [] + _ACTIVE_LAZY_LOADERS[name] = self def flush(self) -> None: """Wait for all pending background writes to complete.""" @@ -185,21 +377,35 @@ def flush(self) -> None: t.join() self._pending.clear() + @classmethod + def flush_all(cls) -> None: + """Flush all active LazyLoader instances.""" + for loader in list(_ACTIVE_LAZY_LOADERS.values()): + loader.flush() + def load_cache( self, key: HashKey, glbls: dict[str, Any] | None = None, ) -> Cache | None: del glbls + blob: bytes | None = None try: - blob: bytes | None = self.store.get(str(self.build_path(key))) + blob = self.store.get(str(self.build_path(key))) if not blob: return None return self.restore_cache(key, blob) except Exception as e: LOGGER.warning("Failed to restore lazy cache: %s", e) + self._on_restore_failure(key, blob) return None + def _on_restore_failure( + self, key: HashKey, manifest_blob: bytes | None + ) -> None: + """Hook after a failed restore. No-op natively; the WASM variant + evicts and poisons the bad keys so HTTP won't re-fetch them.""" + def restore_cache(self, _key: HashKey, blob: bytes) -> Cache: cache_data = msgspec.json.decode(blob, type=CacheSchema) base = Path(self.name) / cache_data.hash @@ -237,51 +443,15 @@ def restore_cache(self, _key: HashKey, blob: bytes) -> Cache: return_ref = cache_data.meta.return_value.reference return_type_hint = cache_data.meta.return_value.type_hint - # Read + deserialize in parallel, stream results via queue. - # Every thread unconditionally puts exactly one item — either the - # deserialized value or _BlobStatus.MISSING — so queue.get() needs - # no timeout. - results: queue.Queue[tuple[str, Any]] = queue.Queue() unique_keys = set(ref_vars.values()) if return_ref: unique_keys.add(return_ref) - - def _load_blob(key: str) -> None: - try: - data = self.store.get(key) - if data: - ext = Path(key).suffix - deserialize = BLOB_DESERIALIZERS.get( - ext, BLOB_DESERIALIZERS[".pickle"] - ) - type_hint = ref_type_hints.get(key) or ( - return_type_hint if key == return_ref else None - ) - results.put((key, deserialize(data, type_hint))) - else: - results.put((key, _BlobStatus.MISSING)) - except Exception as e: - LOGGER.warning("Failed to deserialize blob %s: %s", key, e) - results.put((key, _BlobStatus.MISSING)) - - threads = [ - threading.Thread(target=_load_blob, args=(key,)) - for key in unique_keys - ] - for t in threads: - t.start() - - # N threads → N results guaranteed; no timeout needed. - unpickled: dict[str, Any] = {} - try: - for _ in unique_keys: - key, val = results.get() - if val is _BlobStatus.MISSING: - raise FileNotFoundError("Incomplete cache: missing blobs") - unpickled[key] = val - finally: - for t in threads: - t.join() + # Read + deserialize the blobs. The native loader parallelizes via + # threads; the WASM loader overrides this to fetch via the store's + # concurrent get_batch (threads are unavailable in Pyodide). + unpickled = self._read_blobs( + unique_keys, ref_type_hints, return_ref, return_type_hint + ) # Distribute to defs defs: dict[str, Any] = {} @@ -307,12 +477,12 @@ def _load_blob(key: str) -> None: else: defs[var_name] = val else: - defs[var_name] = from_item(item) + defs[var_name] = from_item(item, var_name) if return_ref and return_ref in unpickled: return_item = unpickled[return_ref] elif cache_data.meta.return_value: - return_item = from_item(cache_data.meta.return_value) + return_item = from_item(cache_data.meta.return_value, "return") else: return_item = None @@ -329,6 +499,78 @@ def _load_blob(key: str) -> None: hit=True, ) + def _deserialize_blob( + self, + key: str, + data: bytes, + ref_type_hints: dict[str, str | None], + return_ref: str | None, + return_type_hint: str | None, + ) -> Any: + ext = Path(key).suffix + deserialize = BLOB_DESERIALIZERS.get( + ext, BLOB_DESERIALIZERS[".pickle"] + ) + type_hint = ref_type_hints.get(key) or ( + return_type_hint if key == return_ref else None + ) + return deserialize(data, type_hint) + + def _read_blobs( + self, + unique_keys: set[str], + ref_type_hints: dict[str, str | None], + return_ref: str | None, + return_type_hint: str | None, + ) -> dict[str, Any]: + """Read + deserialize blobs in parallel via threads. + + Every thread puts exactly one item — value or `_BlobStatus.MISSING` + — so `queue.get()` needs no timeout. `WasmLazyLoader` overrides this + to use the store's concurrent `get_batch` (no threads in Pyodide). + """ + results: queue.Queue[tuple[str, Any]] = queue.Queue() + + def _load_blob(key: str) -> None: + try: + data = self.store.get(key) + if data: + results.put( + ( + key, + self._deserialize_blob( + key, + data, + ref_type_hints, + return_ref, + return_type_hint, + ), + ) + ) + else: + results.put((key, _BlobStatus.MISSING)) + except Exception as e: + LOGGER.warning("Failed to deserialize blob %s: %s", key, e) + results.put((key, _BlobStatus.MISSING)) + + threads = [ + threading.Thread(target=_load_blob, args=(key,)) + for key in unique_keys + ] + for t in threads: + t.start() + unpickled: dict[str, Any] = {} + try: + for _ in unique_keys: + key, val = results.get() + if val is _BlobStatus.MISSING: + raise FileNotFoundError("Incomplete cache: missing blobs") + unpickled[key] = val + finally: + for t in threads: + t.join() + return unpickled + def save_cache(self, cache: Cache) -> bool: # Reap completed threads self._pending = [t for t in self._pending if t.is_alive()] @@ -374,19 +616,24 @@ def save_cache(self, cache: Cache) -> bool: elif loader not in ("inline",): format_vars.setdefault(loader, {})[var] = obj - manifest = msgspec.json.encode( - CacheSchema( - hash=cache.hash, - cache_type=cache_type_enum, - stateful_refs=list(cache.stateful_refs), - defs=defs_dict, - meta=Meta( - version=cache.meta.get("version", MARIMO_CACHE_VERSION), - return_value=return_item, - ), - ui_defs=ui_defs_list, + version = cache.meta.get("version", MARIMO_CACHE_VERSION) + + def _encode_manifest() -> bytes: + # Encoded after the blobs so any `unserializable_type` marks set + # by `_put_or_mark_unserializable` are reflected in the manifest. + return msgspec.json.encode( + CacheSchema( + hash=cache.hash, + cache_type=cache_type_enum, + stateful_refs=list(cache.stateful_refs), + defs=defs_dict, + meta=Meta( + version=version, + return_value=return_item, + ), + ui_defs=ui_defs_list, + ) ) - ) # Capture values for the background thread store = self.store @@ -399,6 +646,43 @@ def save_cache(self, cache: Cache) -> bool: ) manifest_key = str(self.build_path(cache.key)) + def _put_or_mark_unserializable( + key: str, + value: Any, + serialize: Callable[[Any], bytes], + items: list[Item], + var_name: str = "", + ) -> bool: + """Store one blob; on serialization failure write no blob and + instead mark each manifest `Item` with `unserializable_type`. + + A later load reconstructs an `UnhashableStub` tripwire from the + marker (see `from_item`) rather than reading a placeholder blob + off disk — so unserializable values never get dumped to disk and + the cached-execution pre-flight still re-runs the producing cell. + + Returns `True` when the blob was stored, `False` when it was + marked unserializable. + """ + try: + store.put(key, serialize(value)) + except Exception as e: + LOGGER.warning( + "Failed to serialize %s for cache; marking " + "unserializable: %s", + var_name or key, + e, + ) + fallback = f"{type(value).__module__}.{type(value).__name__}" + for item in items: + type_name = item.type_hint or fallback + item.reference = None + item.hash = None + item.type_hint = None + item.unserializable_type = type_name + return False + return True + def _serialize_and_write() -> None: """Serialize and write all blobs + manifest in background.""" try: @@ -406,30 +690,128 @@ def _serialize_and_write() -> None: serialize = BLOB_SERIALIZERS.get( return_loader, pickle.dumps ) - store.put(return_ref, serialize(return_value)) + _put_or_mark_unserializable( + return_ref, + return_value, + serialize, + [return_item], + "return", + ) if ui_vars: - store.put( - (path / "ui.pickle").as_posix(), - pickle.dumps(ui_vars), + ui_key = (path / "ui.pickle").as_posix() + ui_ok = _put_or_mark_unserializable( + ui_key, + ui_vars, + pickle.dumps, + [defs_dict[v] for v in ui_defs_list], + "ui", ) + if not ui_ok: + # UI defs restore via `ui_defs` → `ui.pickle`, + # bypassing the per-Item marks. Drop them from + # `ui_defs` so restore routes through the now-marked + # Items instead, and clear any stale blob lingering + # from a prior run at this same hash path (which would + # otherwise load as a phantom hit). + ui_defs_list.clear() + try: + store.clear(ui_key) + except Exception: + LOGGER.warning( + "Failed to clear stale ui blob %s", ui_key + ) for loader, vars_dict in format_vars.items(): serialize = BLOB_SERIALIZERS.get(loader, pickle.dumps) for var, obj in vars_dict.items(): - store.put( + _put_or_mark_unserializable( (path / f"{var}.{loader}").as_posix(), - serialize(obj), + obj, + serialize, + [defs_dict[var]], + var, ) - # Manifest last — readers check for it to detect complete writes - store.put(manifest_key, manifest) + # Manifest last — readers check for it to detect complete + # writes, and it now carries any unserializable marks set above. + store.put(manifest_key, _encode_manifest()) except Exception: LOGGER.exception("Failed to write cache blobs for %s", path) - t = threading.Thread(target=_serialize_and_write, daemon=False) + self._dispatch_write(_serialize_and_write) + return True + + def _dispatch_write(self, write_fn: Callable[[], None]) -> None: + """Run the blob+manifest write on a background thread (native). + + `WasmLazyLoader` overrides this to run synchronously, since threads + are unavailable in Pyodide. + """ + t = threading.Thread(target=write_fn, daemon=False) t.start() self._pending.append(t) - return True def to_blob(self, cache: Cache) -> bytes | None: # Not used — save_cache is overridden. Kept for interface compliance. del cache return None + + +class WasmLazyLoader(LazyLoader): + """WASM variant of `LazyLoader`, selected once via the dual-loader + registry (so the environment is never re-checked below). + + - reads blobs through the store's concurrent `get_batch` (HTTP fetch in + Pyodide) since threads are unavailable; + - writes synchronously for the same reason; + - on a corrupt restore, evicts the blobs and poisons their keys so the + store won't HTTP-re-fetch the same broken data. + """ + + _store_cls = WasmLazyStore + + def _read_blobs( + self, + unique_keys: set[str], + ref_type_hints: dict[str, str | None], + return_ref: str | None, + return_type_hint: str | None, + ) -> dict[str, Any]: + unpickled: dict[str, Any] = {} + # The store handles concurrency (HTTP batch fetch in WASM). The WASM + # variant always pairs with a WasmExportableStore (see `_store_cls`). + store = self.store + assert isinstance(store, WasmExportableStore) + for key, data in store.get_batch(unique_keys): + if not data: + raise FileNotFoundError("Incomplete cache: missing blobs") + unpickled[key] = self._deserialize_blob( + key, data, ref_type_hints, return_ref, return_type_hint + ) + return unpickled + + def _dispatch_write(self, write_fn: Callable[[], None]) -> None: + write_fn() # synchronous — no threads in Pyodide + + def _on_restore_failure( + self, key: HashKey, manifest_blob: bytes | None + ) -> None: + manifest_path = str(self.build_path(key)) + blob_keys: list[str] = [] + if manifest_blob: + try: + cache_data = msgspec.json.decode( + manifest_blob, type=CacheSchema + ) + for item in cache_data.defs.values(): + if item.reference: + blob_keys.append(item.reference) + if cache_data.meta.return_value and ( + ref := cache_data.meta.return_value.reference + ): + blob_keys.append(ref) + except Exception: + pass + self.store.clear(manifest_path) + _POISONED_KEYS.add(manifest_path) + for blob_key in blob_keys: + self.store.clear(blob_key) + _POISONED_KEYS.add(blob_key) diff --git a/marimo/_save/save.py b/marimo/_save/save.py index 342931a52d2..d498d68c74c 100644 --- a/marimo/_save/save.py +++ b/marimo/_save/save.py @@ -62,6 +62,7 @@ LoaderPartial, LoaderType, MemoryLoader, + resolve_loader, ) from marimo._save.stores.file import FileStore from marimo._save.toplevel import get_cell_id_from_scope, graph_from_scope @@ -1327,7 +1328,8 @@ def my_expensive_function_cached_in_a_certain_location(): "provide one or the other." ) - # Providing a save_path forces the store to be a FileStore + # Providing a save_path forces the store to be a FileStore (works in + # Pyodide too, via its in-memory MEMFS, for in-session caching). if save_path is not None: store = FileStore(save_path) @@ -1335,7 +1337,7 @@ def my_expensive_function_cached_in_a_certain_location(): if store is not None: partial_args["store"] = store - loader = PERSISTENT_LOADERS[method].partial(**partial_args) + loader = resolve_loader(PERSISTENT_LOADERS[method]).partial(**partial_args) # Injection hook for testing if "_loader" in kwargs: loader = kwargs.pop("_loader") diff --git a/marimo/_save/stores/__init__.py b/marimo/_save/stores/__init__.py index 333600dcef6..e114de59a55 100644 --- a/marimo/_save/stores/__init__.py +++ b/marimo/_save/stores/__init__.py @@ -10,7 +10,7 @@ from marimo._save.stores.file import FileStore from marimo._save.stores.redis import RedisStore from marimo._save.stores.rest import RestStore -from marimo._save.stores.store import Store, StoreType +from marimo._save.stores.store import Store, StoreType, WasmExportableStore from marimo._save.stores.tiered import TieredStore LOGGER = _loggers.marimo_logger() @@ -88,5 +88,6 @@ def _get_store_from_config( "StoreKey", "StoreType", "TieredStore", + "WasmExportableStore", "get_store", ] diff --git a/marimo/_save/stores/dict_store.py b/marimo/_save/stores/dict_store.py new file mode 100644 index 00000000000..0804164a146 --- /dev/null +++ b/marimo/_save/stores/dict_store.py @@ -0,0 +1,32 @@ +# Copyright 2026 Marimo. All rights reserved. +from __future__ import annotations + +from marimo._save.stores.store import Store + + +class DictStore(Store): + """A minimal dict-backed store for in-session caching. + + Used as the inner store for LazyStore in WASM (Pyodide), where + the filesystem is unavailable. Writes persist for the lifetime + of the session; reads fall through to HTTP in the outer LazyStore. + """ + + def __init__(self) -> None: + self._data: dict[str, bytes] = {} + + def get(self, key: str) -> bytes | None: + return self._data.get(key) + + def put(self, key: str, value: bytes) -> bool: + self._data[key] = value + return True + + def hit(self, key: str) -> bool: + return key in self._data + + def clear(self, key: str) -> bool: + if key in self._data: + del self._data[key] + return True + return False diff --git a/marimo/_save/stores/store.py b/marimo/_save/stores/store.py index 9ffe3be4a29..7cf535c3286 100644 --- a/marimo/_save/stores/store.py +++ b/marimo/_save/stores/store.py @@ -2,6 +2,10 @@ from __future__ import annotations from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator class Store(ABC): @@ -23,4 +27,28 @@ def clear(self, key: str) -> bool: return False +class WasmExportableStore(Store): + """Interface for stores that support WASM export. + + Provides concurrent multi-key fetch (for WASM blob loading) + and export manifest tracking (for --execute bundling). + """ + + @abstractmethod + def get_batch( + self, keys: Iterable[str] + ) -> Iterator[tuple[str, bytes | None]]: + """Yield (key, data) pairs, potentially fetching concurrently.""" + ... + + @abstractmethod + def export_manifest(self) -> list[str]: + """Return all store keys this session wrote or read. + + Used by --execute export to know exactly which files to copy + to public/cache/ for WASM bundling. + """ + ... + + StoreType = type[Store] diff --git a/marimo/_save/stubs/lazy_stub.py b/marimo/_save/stubs/lazy_stub.py index a6f3348ac25..368e241c2ab 100644 --- a/marimo/_save/stubs/lazy_stub.py +++ b/marimo/_save/stubs/lazy_stub.py @@ -52,6 +52,11 @@ class Item(msgspec.Struct): # Fully-qualified class name of the original value — used by format-aware # deserializers (e.g. to distinguish pandas from polars Arrow blobs). type_hint: str | None = None + # Fully-qualified class name of a def whose value could not be serialized. + # No blob is written for it; the loader reconstructs an `UnhashableStub` + # tripwire from this marker (see `from_item`), so a missing-blob load never + # masquerades as a clean cache hit. + unserializable_type: str | None = None def __post_init__(self) -> None: fields_set = sum( @@ -63,6 +68,7 @@ def __post_init__(self) -> None: self.import_ref, self.function, self.class_def, + self.unserializable_type, ] if field is not None ) @@ -306,13 +312,15 @@ def to_bytes(self) -> bytes: class UnhashableStub(CustomStub): """Marker + tripwire for a def that could not be serialized for caching. - Written to the cache as a placeholder when per-def pickling fails (e.g. - a lambda, a closure over an unpicklable object). The marker is placed - in scope as-is by `Cache.restore` (no `.load()` call). It is - harmless when the consumer cell never touches it; any meaningful access - (call) raises `MarimoUnhashableCacheError` carrying - `variables=[var_name]` so the runner can identify the defining cell, - invalidate its manifest, and re-queue. + When per-def serialization fails (e.g. a lambda, a closure over an + unpicklable object), no blob is written; instead the manifest `Item` + records `unserializable_type`, and the loader reconstructs this stub + in-memory from that marker (see `from_item`). The marker is placed in + scope as-is by `Cache.restore` (no `.load()` call). It is harmless when + the consumer cell never touches it; any meaningful access (call) raises + `MarimoUnhashableCacheError` carrying `variables=[var_name]` so the + runner can identify the defining cell, invalidate its manifest, and + re-queue. Detection happens at use-site, not at pre-execution. Bodies that don't touch the stub run normally; closure-captured stubs surface through @@ -337,9 +345,14 @@ def __init__( _obj: Any = None, var_name: str = "", error_msg: str = "", + type_name: str | None = None, ) -> None: self.var_name = var_name - if _obj is not None: + if type_name is not None: + # Explicit fq name — used when rebuilding from a manifest marker, + # where the original value object is no longer available. + self.type_name = type_name + elif _obj is not None: value_type = type(_obj) self.type_name = ( f"{getattr(value_type, '__module__', '')}." diff --git a/tests/_save/loaders/test_lazy_wasm.py b/tests/_save/loaders/test_lazy_wasm.py new file mode 100644 index 00000000000..c420128fd33 --- /dev/null +++ b/tests/_save/loaders/test_lazy_wasm.py @@ -0,0 +1,294 @@ +# Copyright 2026 Marimo. All rights reserved. +from __future__ import annotations + +import pickle +import tempfile +from pathlib import Path +from unittest import mock + +import msgspec +import pytest + +from marimo._save.cache import MARIMO_CACHE_VERSION, Cache +from marimo._save.loaders.lazy import ( + _ACTIVE_LAZY_LOADERS, + LazyLoader, + LazyStore, + WasmLazyLoader, + WasmLazyStore, +) +from marimo._save.stores.dict_store import DictStore +from marimo._save.stores.file import FileStore +from marimo._save.stores.store import Store, WasmExportableStore +from marimo._save.stubs.lazy_stub import ( + Cache as CacheSchema, + CacheType, + Item, + Meta, +) + + +class TestDictStore: + def test_get_put_hit(self) -> None: + store = DictStore() + assert store.get("k") is None + assert not store.hit("k") + + assert store.put("k", b"v") + assert store.get("k") == b"v" + assert store.hit("k") + + def test_clear(self) -> None: + store = DictStore() + store.put("k", b"v") + assert store.clear("k") + assert not store.hit("k") + assert not store.clear("k") # already gone + + +class TestLazyStoreNative: + """Test LazyStore in native (non-Pyodide) mode.""" + + def test_isinstance_wasm_exportable(self) -> None: + store = LazyStore(inner=DictStore()) + assert isinstance(store, WasmExportableStore) + assert isinstance(store, Store) + + def test_delegates_to_inner(self) -> None: + with tempfile.TemporaryDirectory() as td: + inner = FileStore(td) + store = LazyStore(inner=inner) + + assert store.put("blob.bin", b"data") + assert store.get("blob.bin") == b"data" + assert store.hit("blob.bin") + + # Verify it was written to the inner FileStore + assert (Path(td) / "blob.bin").read_bytes() == b"data" + + def test_get_returns_none_for_missing(self) -> None: + store = LazyStore(inner=DictStore()) + assert store.get("nonexistent") is None + + def test_get_batch_sequential(self) -> None: + inner = DictStore() + store = LazyStore(inner=inner) + store.put("a.bin", b"aa") + store.put("b.bin", b"bb") + store.put("c.bin", b"cc") + + results = dict(store.get_batch(["a.bin", "b.bin", "c.bin"])) + assert results == { + "a.bin": b"aa", + "b.bin": b"bb", + "c.bin": b"cc", + } + + def test_get_batch_missing_key(self) -> None: + store = LazyStore(inner=DictStore()) + store.put("a.bin", b"aa") + results = dict(store.get_batch(["a.bin", "missing"])) + assert results["a.bin"] == b"aa" + assert results["missing"] is None + + def test_export_manifest_tracks_puts(self) -> None: + store = LazyStore(inner=DictStore()) + assert store.export_manifest() == [] + + store.put("x.bin", b"x") + store.put("y.bin", b"y") + assert store.export_manifest() == ["x.bin", "y.bin"] + + def test_export_manifest_clear_removes(self) -> None: + store = LazyStore(inner=DictStore()) + store.put("x.bin", b"x") + store.clear("x.bin") + assert store.export_manifest() == [] + + def test_default_inner_is_filestore(self) -> None: + """Without Pyodide, default inner store is FileStore.""" + store = LazyStore() + assert isinstance(store._inner, FileStore) + + +class TestLazyStoreWasm: + """Test the WASM store variant (`WasmLazyStore`). + + The native/WASM decision is made once via the dual-loader registry + (`resolve_loader` in `loaders/__init__.py`), so the WASM behaviour lives + in the `WasmLazyStore` subclass rather than runtime `is_pyodide()` + branching inside `LazyStore`. + """ + + def test_wasm_default_inner_is_dictstore(self) -> None: + store = WasmLazyStore() + assert isinstance(store._inner, DictStore) + + def test_wasm_put_writes_to_dict(self) -> None: + store = WasmLazyStore(inner=DictStore()) + assert store.put("k", b"v") + # Read from inner directly (no HTTP) + assert store._inner.get("k") == b"v" + + def test_wasm_get_inner_first(self) -> None: + store = WasmLazyStore(inner=DictStore()) + store._inner.put("k", b"cached") + # Should return from inner without HTTP + assert store.get("k") == b"cached" + + @mock.patch("urllib.request.urlopen") + def test_wasm_get_falls_back_to_http( + self, mock_urlopen: mock.Mock + ) -> None: + store = WasmLazyStore(inner=DictStore()) + + # Mock HTTP response + mock_resp = mock.MagicMock() + mock_resp.__enter__ = mock.Mock(return_value=mock_resp) + mock_resp.__exit__ = mock.Mock(return_value=False) + mock_resp.status = 200 + mock_resp.read.return_value = b"from_http" + mock_urlopen.return_value = mock_resp + + # Mock notebook_location for _base_url + with mock.patch( + "marimo._save.loaders.lazy.WasmLazyStore._base_url", + return_value="http://example.com/public/cache", + ): + result = store.get("some/blob.bin") + + assert result == b"from_http" + mock_urlopen.assert_called_once() + + def test_wasm_http_error_returns_none(self) -> None: + store = WasmLazyStore(inner=DictStore()) + with ( + mock.patch( + "urllib.request.urlopen", + side_effect=Exception("network"), + ), + mock.patch( + "marimo._save.loaders.lazy.WasmLazyStore._base_url", + return_value="http://example.com/public/cache", + ), + ): + assert store.get("missing.bin") is None + + +class TestKeySanitization: + def test_valid_key(self) -> None: + assert ( + WasmLazyStore._sanitize_key("hash1/var.pickle") + == "hash1/var.pickle" + ) + + def test_rejects_parent_traversal(self) -> None: + with pytest.raises(ValueError, match="Invalid cache key"): + WasmLazyStore._sanitize_key("../etc/passwd") + + def test_rejects_absolute_path(self) -> None: + with pytest.raises(ValueError, match="Invalid cache key"): + WasmLazyStore._sanitize_key("/etc/passwd") + + def test_rejects_embedded_dotdot(self) -> None: + with pytest.raises(ValueError, match="Invalid cache key"): + WasmLazyStore._sanitize_key("foo/../../bar") + + +class TestLazyLoaderBatchPath: + """Test that LazyLoader uses get_batch when store is WasmExportableStore.""" + + def test_restore_uses_batch_path(self) -> None: + inner = DictStore() + store = LazyStore(inner=inner) + loader = LazyLoader("test_batch", store=store) + + # Seed a cache manually + base = Path("test_batch") / "hash1" + var_ref = (base / "var1.pickle").as_posix() + store.put(var_ref, pickle.dumps("value1")) + + manifest = msgspec.json.encode( + CacheSchema( + hash="hash1", + cache_type=CacheType("Pure"), + defs={"var1": Item(reference=var_ref)}, + stateful_refs=[], + meta=Meta(version=MARIMO_CACHE_VERSION), + ) + ) + cache_path = loader.build_path( + type("Key", (), {"hash": "hash1", "cache_type": "Pure"})() + ) + store.put(str(cache_path), manifest) + + # Load — should use get_batch path (no threads) + from marimo._save.hash import HashKey + + key = HashKey(hash="hash1", cache_type="Pure") + loaded = loader.load_cache(key) + assert loaded is not None + assert loaded.defs["var1"] == "value1" + + def test_save_cache_sync_in_wasm(self) -> None: + # `WasmLazyLoader` is the WASM variant: it writes synchronously + # (no threads in Pyodide) via `_dispatch_write`. + inner = DictStore() + store = WasmLazyStore(inner=inner) + loader = WasmLazyLoader("test_sync", store=store) + + cache = Cache( + defs={"x": 42, "y": "hello"}, + hash="sync_hash", + cache_type="Pure", + stateful_refs=set(), + hit=False, + meta={"version": MARIMO_CACHE_VERSION}, + ) + assert loader.save_cache(cache) + + # Verify blobs were written to the DictStore + assert store.export_manifest() # something was written + + # Load back in native mode (so we don't need js module). + # The data is in the DictStore inner, and get_batch uses + # the native sequential path. + from marimo._save.hash import HashKey + + key = HashKey(hash="sync_hash", cache_type="Pure") + loaded = loader.load_cache(key) + assert loaded is not None + assert loaded.defs["x"] == 42 + assert loaded.defs["y"] == "hello" + + +class TestFlushAll: + def test_flush_all_flushes_active_loaders(self) -> None: + inner = DictStore() + store1 = LazyStore(inner=inner) + store2 = LazyStore(inner=DictStore()) + loader1 = LazyLoader("flush_a", store=store1) + loader2 = LazyLoader("flush_b", store=store2) + + assert _ACTIVE_LAZY_LOADERS.get("flush_a") is loader1 + assert _ACTIVE_LAZY_LOADERS.get("flush_b") is loader2 + + # flush_all should not raise + LazyLoader.flush_all() + + def test_loaders_tracked_by_name(self) -> None: + """Loaders are tracked in the active dict by name.""" + loader = LazyLoader("track_test", store=LazyStore(inner=DictStore())) + assert _ACTIVE_LAZY_LOADERS["track_test"] is loader + + def test_store_reused_across_recreations(self) -> None: + """When a loader is recreated with the same name, it reuses the + previous loader's store (preserving DictStore data).""" + store1 = LazyStore(inner=DictStore()) + loader1 = LazyLoader("reuse_test", store=store1) + store1.put("key1", b"data1") + + # Recreate without explicit store — should reuse store1 + loader2 = LazyLoader("reuse_test") + assert loader2.store is store1 + assert loader2.store.get("key1") == b"data1" diff --git a/tests/_save/loaders/test_loader.py b/tests/_save/loaders/test_loader.py index 4310089e97a..95d2f922196 100644 --- a/tests/_save/loaders/test_loader.py +++ b/tests/_save/loaders/test_loader.py @@ -30,6 +30,7 @@ CacheType, Item, Meta, + UnhashableStub, ) from tests._save.loaders.mocks import MockLoader @@ -289,6 +290,94 @@ def test_import_reference_stays_inline(self) -> None: assert loaded.defs["Optional"] is Optional assert loaded.defs["OrderedDict"] is OrderedDict + def test_unserializable_def_marks_manifest_no_blob(self) -> None: + """A def that can't be serialized (a lambda) writes no blob; the + manifest `Item` carries `unserializable_type`, and load reconstructs + an `UnhashableStub` tripwire in-memory.""" + loader = self.instance() + cache = Cache( + # `ok` is a list (pickle-blob path); `f` is a lambda (unpicklable). + defs={"ok": [1, 2, 3], "f": lambda x: x + 1}, + hash="unserializable_hash", + cache_type="Pure", + stateful_refs=set(), + hit=False, + meta={"version": MARIMO_CACHE_VERSION}, + ) + assert loader.save_cache(cache) + loader.flush() + + # No blob written for the lambda — only the serializable `ok` blob + # and the manifest exist. + blob_names = sorted( + p.name + for p in Path(self.store.save_path).rglob("*") + if p.is_file() and p.suffix != ".jsonl" + ) + assert blob_names == ["ok.pickle"], blob_names + + # Manifest marks the lambda's Item, with no dangling reference. + cache_path = loader.build_path(key("unserializable_hash", "Pure")) + manifest = self.store.get(str(cache_path)) + decoded = msgspec.json.decode(manifest, type=CacheSchema) + f_item = decoded.defs["f"] + assert f_item.reference is None + assert f_item.unserializable_type is not None + assert "function" in f_item.unserializable_type.lower() + + # Load: `ok` restores; `f` becomes an in-memory tripwire. + loaded = loader.load_cache(key("unserializable_hash", "Pure")) + assert loaded is not None + assert loaded.defs["ok"] == [1, 2, 3] + assert isinstance(loaded.defs["f"], UnhashableStub) + assert loaded.defs["f"].var_name == "f" + + def test_unserializable_ui_clears_ui_defs_and_stale_blob(self) -> None: + """A UI def that can't be pickled must not leave `ui_defs` pointing + at `ui.pickle`: restore loads UI defs via `ui_defs`, bypassing the + per-Item marks, so a stale `ui.pickle` from a prior run would load as + a phantom hit. On failure `ui_defs` is cleared, each UI Item is + marked, and the stale blob is removed.""" + from marimo._save.stubs.ui_element_stub import UIElementStub + + loader = self.instance() + # A UIElementStub instance routes to the "ui" loader by type; give it + # an unpicklable attribute so the shared ui.pickle write fails. + ui_stub = UIElementStub.__new__(UIElementStub) + ui_stub.data = {"bad": lambda: 1} # type: ignore[attr-defined] + + # Pre-seed a stale ui.pickle at the hash path from a "prior run". + base = Path("test") / "ui_fail_hash" + ui_key = (base / "ui.pickle").as_posix() + self.store.put(ui_key, pickle.dumps({"u": "stale"})) + + cache = Cache( + defs={"ok": [1, 2, 3], "u": ui_stub}, + hash="ui_fail_hash", + cache_type="Pure", + stateful_refs=set(), + hit=False, + meta={"version": MARIMO_CACHE_VERSION}, + ) + assert loader.save_cache(cache) + loader.flush() + + # Stale ui.pickle removed; the serializable `ok` blob remains. + assert self.store.get(ui_key) is None + cache_path = loader.build_path(key("ui_fail_hash", "Pure")) + decoded = msgspec.json.decode( + self.store.get(str(cache_path)), type=CacheSchema + ) + assert decoded.ui_defs == [] + assert decoded.defs["u"].reference is None + assert decoded.defs["u"].unserializable_type is not None + + # Load: `u` is an in-memory tripwire, not the stale value. + loaded = loader.load_cache(key("ui_fail_hash", "Pure")) + assert loaded is not None + assert loaded.defs["ok"] == [1, 2, 3] + assert isinstance(loaded.defs["u"], UnhashableStub) + def test_corrupt_cache_returns_none(self) -> None: """Corrupt manifest triggers cache miss, not crash.""" loader = self.instance() diff --git a/tests/_save/stubs/test_lazy_codecs.py b/tests/_save/stubs/test_lazy_codecs.py new file mode 100644 index 00000000000..aa7dada9dae --- /dev/null +++ b/tests/_save/stubs/test_lazy_codecs.py @@ -0,0 +1,37 @@ +# Copyright 2026 Marimo. All rights reserved. +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") + +from marimo._save.loaders.lazy import maybe_update_lazy_stub # noqa: E402 +from marimo._save.stubs.lazy_stub import ( # noqa: E402 + BLOB_DESERIALIZERS, + BLOB_SERIALIZERS, +) + + +def test_tensor_resolves_to_pt_codec() -> None: + assert maybe_update_lazy_stub(torch.ones(4)) == "pt" + + +def test_parameter_resolves_to_pt_codec_via_mro() -> None: + param = torch.nn.Parameter(torch.ones(4)) + assert maybe_update_lazy_stub(param) == "pt" + + +def test_pt_round_trip() -> None: + tensor = torch.randn(8, 3) + data = BLOB_SERIALIZERS["pt"](tensor) + restored = BLOB_DESERIALIZERS[".pt"](data, "torch.Tensor") + assert torch.equal(tensor, restored) + assert restored.dtype == tensor.dtype + + +def test_pt_round_trip_preserves_parameter_subclass() -> None: + param = torch.nn.Parameter(torch.randn(4)) + data = BLOB_SERIALIZERS["pt"](param) + restored = BLOB_DESERIALIZERS[".pt"](data, "torch.nn.parameter.Parameter") + assert isinstance(restored, torch.nn.Parameter) + assert torch.equal(param.data, restored.data) diff --git a/tests/_save/stubs/test_unhashable_stub.py b/tests/_save/stubs/test_unhashable_stub.py index d1df7acb129..c1e65eecd76 100644 --- a/tests/_save/stubs/test_unhashable_stub.py +++ b/tests/_save/stubs/test_unhashable_stub.py @@ -16,7 +16,8 @@ MarimoUnhashableCacheError, ) from marimo._save.cache import Cache -from marimo._save.stubs.lazy_stub import UnhashableStub +from marimo._save.loaders.lazy import from_item +from marimo._save.stubs.lazy_stub import Item, UnhashableStub # --------------------------------------------------------------------------- # UnhashableStub: data + tripwire semantics @@ -55,6 +56,22 @@ def test_pickle_roundtrip(self) -> None: assert round_tripped.type_name == "builtins.function" assert round_tripped.error_msg == "oops" + def test_type_name_param_overrides_derivation(self) -> None: + """An explicit `type_name` (used when rebuilding from a manifest + marker, with no original object) wins over `_obj` derivation.""" + stub = UnhashableStub(var_name="f", type_name="numpy.ndarray") + assert stub.type_name == "numpy.ndarray" + assert stub.var_name == "f" + + def test_from_item_reconstructs_marker(self) -> None: + """`from_item` rebuilds the tripwire in-memory from an Item carrying + `unserializable_type` — no blob is read off disk.""" + item = Item(unserializable_type="builtins.function") + stub = from_item(item, "f") + assert isinstance(stub, UnhashableStub) + assert stub.var_name == "f" + assert stub.type_name == "builtins.function" + def test_isinstance_works(self) -> None: """Tripwires must not interfere with isinstance — the restore path uses isinstance to decide whether to install the marker.""" From 4f4ad20533d4cfdd1923ae47a471f7bbd2fac8f0 Mon Sep 17 00:00:00 2001 From: "dmadisetti@coreweave.com" Date: Mon, 29 Jun 2026 10:30:51 -0700 Subject: [PATCH 2/6] =?UTF-8?q?refactor(save):=20address=20review=20?= =?UTF-8?q?=E2=80=94=20fold=20export=20API=20into=20Store,=20cache=20HTTP?= =?UTF-8?q?=20reads?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Move get_batch (sequential default) and export_keys (empty default, renamed from export_manifest) onto the base Store; drop the WasmExportableStore ABC. LazyStore overrides export_keys to report tracked keys; WasmLazyStore overrides get_batch for concurrent fetch. - WasmLazyStore._http_get/_http_get_batch now cache successful fetches into the inner store and mark them touched, so repeat reads stay local and the export manifest covers HTTP-fetched blobs. - _sanitize_key validates (.. / absolute) but returns the original key instead of a normalized string, avoiding requested/stored key desync. - Remove a duplicate TYPE_CHECKING import block. --- marimo/_save/loaders/lazy.py | 56 ++++++++++++++++----------- marimo/_save/stores/__init__.py | 3 +- marimo/_save/stores/store.py | 30 +++++++------- tests/_save/loaders/test_lazy_wasm.py | 38 +++++++++++++----- 4 files changed, 75 insertions(+), 52 deletions(-) diff --git a/marimo/_save/loaders/lazy.py b/marimo/_save/loaders/lazy.py index 0a48b54dc4f..0aacf38d7dd 100644 --- a/marimo/_save/loaders/lazy.py +++ b/marimo/_save/loaders/lazy.py @@ -20,7 +20,6 @@ from marimo._save.hash import HashKey from marimo._save.loaders.loader import BasePersistenceLoader from marimo._save.stores import FileStore, Store -from marimo._save.stores.store import WasmExportableStore if TYPE_CHECKING: from collections.abc import Callable, Iterable, Iterator @@ -44,9 +43,6 @@ ) from marimo._save.stubs.stubs import mro_lookup -if TYPE_CHECKING: - from collections.abc import Callable - LOGGER = _loggers.marimo_logger() @@ -199,7 +195,7 @@ def _get_wasm_dict_store() -> Store: _POISONED_KEYS: set[str] = set() -class LazyStore(WasmExportableStore): +class LazyStore(Store): """Native store for `LazyLoader`. Delegates to an inner store (default `FileStore` at `__marimo__/cache/`) @@ -237,13 +233,7 @@ def clear(self, key: str) -> bool: self._touched_keys.discard(key) return self._inner.clear(key) - def get_batch( - self, keys: Iterable[str] - ) -> Iterator[tuple[str, bytes | None]]: - for k in keys: - yield k, self.get(k) - - def export_manifest(self) -> list[str]: + def export_keys(self) -> list[str]: return sorted(self._written_keys | self._touched_keys) @@ -296,23 +286,36 @@ def _base_url(self) -> str: @staticmethod def _sanitize_key(key: str) -> str: - """Prevent path traversal in HTTP fetch keys.""" + """Validate a fetch key against path traversal. + + Returns `key` unchanged so it still matches the stored key; + only rejects `..` segments and absolute paths (normalizing the + string here could desync requested and stored keys). + """ clean = PurePosixPath(key) if ".." in clean.parts or clean.is_absolute(): raise ValueError(f"Invalid cache key: {key}") - return str(clean) + return key def _http_get(self, key: str) -> bytes | None: - """Single sync fetch via pyodide_http-patched urllib.""" + """Single sync fetch via pyodide_http-patched urllib. + + A successful response is written back to the inner store and + marked touched, so repeat reads stay local and the export + manifest covers HTTP-fetched blobs. + """ import urllib.request - key = self._sanitize_key(key) - url = f"{self._base_url()}/{key}" + url = f"{self._base_url()}/{self._sanitize_key(key)}" try: with urllib.request.urlopen(url) as resp: - return resp.read() if resp.status == 200 else None + data = resp.read() if resp.status == 200 else None except Exception: return None + if data is not None: + self._inner.put(key, data) + self._touched_keys.add(key) + return data def _http_get_batch( self, keys: Iterable[str] @@ -344,7 +347,14 @@ async def _fetch_all() -> list[tuple[str, bytes | None]]: # back to sequential synchronous XHR via the pyodide_http-patched # urllib — legal in a worker. results = [(k, self._http_get(k)) for k in keys_list] - yield from results + # Cache successful fetches in-session (idempotent for the fallback + # path, which already wrote via `_http_get`) so repeat reads stay + # local and the export manifest covers HTTP-fetched blobs. + for key, data in results: + if data is not None: + self._inner.put(key, data) + self._touched_keys.add(key) + yield key, data # Registry of active loaders by name, so a recreated loader reuses the same @@ -777,10 +787,10 @@ def _read_blobs( ) -> dict[str, Any]: unpickled: dict[str, Any] = {} # The store handles concurrency (HTTP batch fetch in WASM). The WASM - # variant always pairs with a WasmExportableStore (see `_store_cls`). - store = self.store - assert isinstance(store, WasmExportableStore) - for key, data in store.get_batch(unique_keys): + # variant pairs with a `WasmLazyStore` (see `_store_cls`), whose + # `get_batch` fetches concurrently; `get_batch` is defined on every + # `Store` so no narrowing is needed. + for key, data in self.store.get_batch(unique_keys): if not data: raise FileNotFoundError("Incomplete cache: missing blobs") unpickled[key] = self._deserialize_blob( diff --git a/marimo/_save/stores/__init__.py b/marimo/_save/stores/__init__.py index e114de59a55..333600dcef6 100644 --- a/marimo/_save/stores/__init__.py +++ b/marimo/_save/stores/__init__.py @@ -10,7 +10,7 @@ from marimo._save.stores.file import FileStore from marimo._save.stores.redis import RedisStore from marimo._save.stores.rest import RestStore -from marimo._save.stores.store import Store, StoreType, WasmExportableStore +from marimo._save.stores.store import Store, StoreType from marimo._save.stores.tiered import TieredStore LOGGER = _loggers.marimo_logger() @@ -88,6 +88,5 @@ def _get_store_from_config( "StoreKey", "StoreType", "TieredStore", - "WasmExportableStore", "get_store", ] diff --git a/marimo/_save/stores/store.py b/marimo/_save/stores/store.py index 7cf535c3286..c086ba65c76 100644 --- a/marimo/_save/stores/store.py +++ b/marimo/_save/stores/store.py @@ -26,29 +26,25 @@ def clear(self, key: str) -> bool: del key return False - -class WasmExportableStore(Store): - """Interface for stores that support WASM export. - - Provides concurrent multi-key fetch (for WASM blob loading) - and export manifest tracking (for --execute bundling). - """ - - @abstractmethod def get_batch( self, keys: Iterable[str] ) -> Iterator[tuple[str, bytes | None]]: - """Yield (key, data) pairs, potentially fetching concurrently.""" - ... + """Yield `(key, data)` pairs for `keys`. - @abstractmethod - def export_manifest(self) -> list[str]: - """Return all store keys this session wrote or read. + Defaults to a sequential `get` per key. Stores that can fetch + concurrently (e.g. the WASM HTTP store) override this. + """ + for key in keys: + yield key, self.get(key) + + def export_keys(self) -> list[str]: + """Return the keys this session wrote or read that should be + bundled on `--execute` export. - Used by --execute export to know exactly which files to copy - to public/cache/ for WASM bundling. + Defaults to none; stores that track usage (e.g. `LazyStore`) + override this. Returning `[]` keeps non-tracking stores inert. """ - ... + return [] StoreType = type[Store] diff --git a/tests/_save/loaders/test_lazy_wasm.py b/tests/_save/loaders/test_lazy_wasm.py index c420128fd33..fa0adfe6b83 100644 --- a/tests/_save/loaders/test_lazy_wasm.py +++ b/tests/_save/loaders/test_lazy_wasm.py @@ -19,7 +19,7 @@ ) from marimo._save.stores.dict_store import DictStore from marimo._save.stores.file import FileStore -from marimo._save.stores.store import Store, WasmExportableStore +from marimo._save.stores.store import Store from marimo._save.stubs.lazy_stub import ( Cache as CacheSchema, CacheType, @@ -45,13 +45,27 @@ def test_clear(self) -> None: assert not store.hit("k") assert not store.clear("k") # already gone + def test_get_batch_default_sequential(self) -> None: + # DictStore inherits the base Store sequential get_batch. + store = DictStore() + store.put("a", b"1") + assert dict(store.get_batch(["a", "missing"])) == { + "a": b"1", + "missing": None, + } + + def test_export_keys_defaults_empty(self) -> None: + # Non-tracking stores inherit the inert base Store default. + store = DictStore() + store.put("a", b"1") + assert store.export_keys() == [] + class TestLazyStoreNative: """Test LazyStore in native (non-Pyodide) mode.""" - def test_isinstance_wasm_exportable(self) -> None: + def test_is_store(self) -> None: store = LazyStore(inner=DictStore()) - assert isinstance(store, WasmExportableStore) assert isinstance(store, Store) def test_delegates_to_inner(self) -> None: @@ -91,19 +105,19 @@ def test_get_batch_missing_key(self) -> None: assert results["a.bin"] == b"aa" assert results["missing"] is None - def test_export_manifest_tracks_puts(self) -> None: + def test_export_keys_tracks_puts(self) -> None: store = LazyStore(inner=DictStore()) - assert store.export_manifest() == [] + assert store.export_keys() == [] store.put("x.bin", b"x") store.put("y.bin", b"y") - assert store.export_manifest() == ["x.bin", "y.bin"] + assert store.export_keys() == ["x.bin", "y.bin"] - def test_export_manifest_clear_removes(self) -> None: + def test_export_keys_clear_removes(self) -> None: store = LazyStore(inner=DictStore()) store.put("x.bin", b"x") store.clear("x.bin") - assert store.export_manifest() == [] + assert store.export_keys() == [] def test_default_inner_is_filestore(self) -> None: """Without Pyodide, default inner store is FileStore.""" @@ -159,6 +173,10 @@ def test_wasm_get_falls_back_to_http( assert result == b"from_http" mock_urlopen.assert_called_once() + # Successful fetch is cached in-session and recorded for export, + # so repeat reads stay local and the bundle ships it. + assert store._inner.get("some/blob.bin") == b"from_http" + assert "some/blob.bin" in store.export_keys() def test_wasm_http_error_returns_none(self) -> None: store = WasmLazyStore(inner=DictStore()) @@ -196,7 +214,7 @@ def test_rejects_embedded_dotdot(self) -> None: class TestLazyLoaderBatchPath: - """Test that LazyLoader uses get_batch when store is WasmExportableStore.""" + """Test that LazyLoader uses get_batch via the store.get_batch path.""" def test_restore_uses_batch_path(self) -> None: inner = DictStore() @@ -248,7 +266,7 @@ def test_save_cache_sync_in_wasm(self) -> None: assert loader.save_cache(cache) # Verify blobs were written to the DictStore - assert store.export_manifest() # something was written + assert store.export_keys() # something was written # Load back in native mode (so we don't need js module). # The data is in the DictStore inner, and get_batch uses From 2ebdf7c3b37319790e96ab43e7fafde46b5d544c Mon Sep 17 00:00:00 2001 From: "dmadisetti@coreweave.com" Date: Tue, 30 Jun 2026 10:06:18 -0700 Subject: [PATCH 3/6] test(save): cover WasmLazyLoader._on_restore_failure The corrupt-restore eviction + key-poisoning path was untested. Covers: evicting and poisoning the manifest plus its referenced blobs, the return_value reference, None and undecodable manifests, and an end-to-end corrupt restore through load_cache. --- tests/_save/loaders/test_lazy_wasm.py | 118 ++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) diff --git a/tests/_save/loaders/test_lazy_wasm.py b/tests/_save/loaders/test_lazy_wasm.py index fa0adfe6b83..9f5c7ab36e7 100644 --- a/tests/_save/loaders/test_lazy_wasm.py +++ b/tests/_save/loaders/test_lazy_wasm.py @@ -4,14 +4,21 @@ import pickle import tempfile from pathlib import Path +from typing import TYPE_CHECKING from unittest import mock import msgspec import pytest +if TYPE_CHECKING: + from collections.abc import Iterator + + from marimo._save.hash import HashKey + from marimo._save.cache import MARIMO_CACHE_VERSION, Cache from marimo._save.loaders.lazy import ( _ACTIVE_LAZY_LOADERS, + _POISONED_KEYS, LazyLoader, LazyStore, WasmLazyLoader, @@ -310,3 +317,114 @@ def test_store_reused_across_recreations(self) -> None: loader2 = LazyLoader("reuse_test") assert loader2.store is store1 assert loader2.store.get("key1") == b"data1" + + +class TestOnRestoreFailure: + """`WasmLazyLoader._on_restore_failure`: on a corrupt restore, evict the + manifest and its referenced blobs from the store and poison their keys so + the HTTP fallback never re-fetches the same broken data.""" + + @pytest.fixture(autouse=True) + def _isolate_poisoned_keys(self) -> Iterator[None]: + # _POISONED_KEYS is a module global; snapshot/restore so these tests + # neither leak into nor depend on others. + snapshot = set(_POISONED_KEYS) + yield + _POISONED_KEYS.clear() + _POISONED_KEYS.update(snapshot) + + @staticmethod + def _manifest( + refs: list[str], return_ref: str | None = None + ) -> bytes: + meta = Meta(version=MARIMO_CACHE_VERSION) + if return_ref is not None: + meta = Meta( + version=MARIMO_CACHE_VERSION, + return_value=Item(reference=return_ref), + ) + return msgspec.json.encode( + CacheSchema( + hash="h", + cache_type=CacheType("Pure"), + defs={ + f"v{i}": Item(reference=r) for i, r in enumerate(refs) + }, + stateful_refs=[], + meta=meta, + ) + ) + + def _loader_and_key( + self, + ) -> tuple[WasmLazyStore, WasmLazyLoader, HashKey]: + from marimo._save.hash import HashKey + + store = WasmLazyStore(inner=DictStore()) + loader = WasmLazyLoader("restore_fail", store=store) + return store, loader, HashKey(hash="h", cache_type="Pure") + + def test_evicts_and_poisons_manifest_and_blobs(self) -> None: + store, loader, key = self._loader_and_key() + manifest_path = str(loader.build_path(key)) + blob_a, blob_b = "h/a.pickle", "h/b.pickle" + for k in (manifest_path, blob_a, blob_b): + store.put(k, b"x") + + loader._on_restore_failure(key, self._manifest([blob_a, blob_b])) + + for k in (manifest_path, blob_a, blob_b): + assert not store.hit(k), f"{k} not evicted" + assert k in _POISONED_KEYS, f"{k} not poisoned" + + def test_poisons_return_value_reference(self) -> None: + store, loader, key = self._loader_and_key() + ret_ref = "h/return.pickle" + store.put(ret_ref, b"x") + + loader._on_restore_failure( + key, self._manifest([], return_ref=ret_ref) + ) + + assert not store.hit(ret_ref) + assert ret_ref in _POISONED_KEYS + + def test_none_manifest_poisons_only_manifest_path(self) -> None: + store, loader, key = self._loader_and_key() + manifest_path = str(loader.build_path(key)) + + loader._on_restore_failure(key, None) + + assert manifest_path in _POISONED_KEYS + # No blob keys to discover, so nothing else is poisoned. + assert _POISONED_KEYS == {manifest_path} + + def test_undecodable_manifest_poisons_only_manifest_path(self) -> None: + store, loader, key = self._loader_and_key() + manifest_path = str(loader.build_path(key)) + + # Garbage bytes must not raise; manifest path is still poisoned. + loader._on_restore_failure(key, b"not valid msgpack/json") + + assert manifest_path in _POISONED_KEYS + assert _POISONED_KEYS == {manifest_path} + + def test_corrupt_restore_via_load_cache_poisons_keys(self) -> None: + # End-to-end: a manifest referencing a missing blob makes restore + # raise; load_cache swallows it and invokes _on_restore_failure. + store, loader, key = self._loader_and_key() + manifest_path = str(loader.build_path(key)) + missing_ref = "h/missing.pickle" + store._inner.put(manifest_path, self._manifest([missing_ref])) + + # The missing blob would otherwise be fetched over HTTP; stub it to + # a miss so the restore fails deterministically without network. + with mock.patch.object( + store, + "_http_get_batch", + return_value=iter([(missing_ref, None)]), + ): + assert loader.load_cache(key) is None + + assert manifest_path in _POISONED_KEYS + assert missing_ref in _POISONED_KEYS From 76768b2756dbcb6e1e217c4d3a37b7c45d3ef958 Mon Sep 17 00:00:00 2001 From: "dmadisetti@coreweave.com" Date: Tue, 30 Jun 2026 11:08:02 -0700 Subject: [PATCH 4/6] refactor(save): scope lazy-store session state to CacheState Move the _ACTIVE_LAZY_LOADERS, _POISONED_KEYS, and _WASM_DICT_STORE module globals onto the per-session CacheState via a _cache_state() accessor, falling back to a process-local CacheState when no runtime context is installed (standalone @mo.cache/persistent_cache). The accessor walks to the root context so embedded apps (child kernel contexts) keep sharing the session's loader registry, poison set, and WASM store rather than reinitializing them per child. --- marimo/_save/cache.py | 5 ++ marimo/_save/loaders/lazy.py | 73 +++++++++++++++++---------- tests/_save/loaders/test_lazy_wasm.py | 67 ++++++++++++++---------- 3 files changed, 90 insertions(+), 55 deletions(-) diff --git a/marimo/_save/cache.py b/marimo/_save/cache.py index c07159bcfea..91444a21103 100644 --- a/marimo/_save/cache.py +++ b/marimo/_save/cache.py @@ -41,6 +41,7 @@ from marimo._runtime.state import State from marimo._save.hash import HashKey from marimo._save.loaders import Loader + from marimo._save.loaders.lazy import LazyLoader from marimo._save.stores import Store # NB. Increment on cache breaking changes. @@ -76,6 +77,10 @@ class CacheState: store: Store hash_memo: dict[str, bytes] = field(default_factory=dict) + # Lazy-store session state; see `loaders/lazy.py:_cache_state`. + active_lazy_loaders: dict[str, LazyLoader] = field(default_factory=dict) + poisoned_keys: set[str] = field(default_factory=set) + wasm_dict_store: Store | None = None def is_memoizable(self, value: Any) -> bool: """Whether *value* is eligible for content-hash memoization. diff --git a/marimo/_save/loaders/lazy.py b/marimo/_save/loaders/lazy.py index 0aacf38d7dd..5cda39bebd0 100644 --- a/marimo/_save/loaders/lazy.py +++ b/marimo/_save/loaders/lazy.py @@ -13,13 +13,15 @@ import msgspec from marimo import _loggers +from marimo._runtime.context import safe_get_context from marimo._save.cache import ( MARIMO_CACHE_VERSION, Cache, + CacheState, ) from marimo._save.hash import HashKey from marimo._save.loaders.loader import BasePersistenceLoader -from marimo._save.stores import FileStore, Store +from marimo._save.stores import DEFAULT_STORE, FileStore, Store if TYPE_CHECKING: from collections.abc import Callable, Iterable, Iterator @@ -175,24 +177,36 @@ def from_item(item: Item, var_name: str = "") -> Any: return None -# Module-level singleton DictStore for WASM. Shared across all LazyStore -# instances so cached data survives loader recreation (State GC, partial -# reconstruction, etc.). -_WASM_DICT_STORE: Store | None = None +# Fallback to a process-local CacheState. +_FALLBACK_CACHE_STATE: CacheState | None = None -def _get_wasm_dict_store() -> Store: - global _WASM_DICT_STORE - if _WASM_DICT_STORE is None: - from marimo._save.stores.dict_store import DictStore +def _cache_state() -> CacheState: + """Cache state for the current session, scoped to the root context. - _WASM_DICT_STORE = DictStore() - return _WASM_DICT_STORE + Embedded apps run in child contexts (`create_kernel_context(parent=...)`); + walk to the root so the loader registry, poison set, and WASM store are + shared across the session rather than per child. + """ + ctx = safe_get_context() + if ctx is not None: + while ctx.parent is not None: + ctx = ctx.parent + return ctx.cache + global _FALLBACK_CACHE_STATE + if _FALLBACK_CACHE_STATE is None: + _FALLBACK_CACHE_STATE = CacheState(store=DEFAULT_STORE()) + return _FALLBACK_CACHE_STATE -# Keys whose deserialization failed — never re-fetched over HTTP. Survives -# across LazyStore instances. -_POISONED_KEYS: set[str] = set() +def _get_wasm_dict_store() -> Store: + """The session's shared in-memory WASM store, created on first use.""" + cache_state = _cache_state() + if cache_state.wasm_dict_store is None: + from marimo._save.stores.dict_store import DictStore + + cache_state.wasm_dict_store = DictStore() + return cache_state.wasm_dict_store class LazyStore(Store): @@ -254,7 +268,7 @@ def get(self, key: str) -> bytes | None: result = super().get(key) if result is not None: return result - if key not in _POISONED_KEYS: + if key not in _cache_state().poisoned_keys: return self._http_get(key) return None @@ -263,13 +277,14 @@ def get_batch( ) -> Iterator[tuple[str, bytes | None]]: # Inner store first; HTTP-fetch only the (unpoisoned) misses, fired # concurrently. + poisoned = _cache_state().poisoned_keys http_keys: list[str] = [] for k in keys: data = self._inner.get(k) if data is not None: self._touched_keys.add(k) yield k, data - elif k not in _POISONED_KEYS: + elif k not in poisoned: http_keys.append(k) else: yield k, None @@ -357,11 +372,6 @@ async def _fetch_all() -> list[tuple[str, bytes | None]]: yield key, data -# Registry of active loaders by name, so a recreated loader reuses the same -# store (and so `flush_all`/export can reach every session loader). -_ACTIVE_LAZY_LOADERS: dict[str, LazyLoader] = {} - - class LazyLoader(BasePersistenceLoader): # Default store class, overridden by the WASM variant. The single # native/WASM decision is made once in the dual-loader registry, not here. @@ -372,14 +382,15 @@ def __init__( name: str, store: Store | None = None, ) -> None: + loaders = _cache_state().active_lazy_loaders if store is None: - # Reuse the same store across recreations of a named loader - # (State GC, partial reconstruction) so cached data survives. - prev = _ACTIVE_LAZY_LOADERS.get(name) + # Reuse the store across recreations of a named loader (State GC, + # partial reconstruction) so cached data survives. + prev = loaders.get(name) store = prev.store if prev is not None else self._store_cls() super().__init__(name, "jsonl", store) self._pending: list[threading.Thread] = [] - _ACTIVE_LAZY_LOADERS[name] = self + loaders[name] = self def flush(self) -> None: """Wait for all pending background writes to complete.""" @@ -390,9 +401,14 @@ def flush(self) -> None: @classmethod def flush_all(cls) -> None: """Flush all active LazyLoader instances.""" - for loader in list(_ACTIVE_LAZY_LOADERS.values()): + for loader in cls.active_loaders(): loader.flush() + @classmethod + def active_loaders(cls) -> list[LazyLoader]: + """Loaders the current session has created (for flush/export).""" + return list(_cache_state().active_lazy_loaders.values()) + def load_cache( self, key: HashKey, @@ -820,8 +836,9 @@ def _on_restore_failure( blob_keys.append(ref) except Exception: pass + poisoned = _cache_state().poisoned_keys self.store.clear(manifest_path) - _POISONED_KEYS.add(manifest_path) + poisoned.add(manifest_path) for blob_key in blob_keys: self.store.clear(blob_key) - _POISONED_KEYS.add(blob_key) + poisoned.add(blob_key) diff --git a/tests/_save/loaders/test_lazy_wasm.py b/tests/_save/loaders/test_lazy_wasm.py index 9f5c7ab36e7..4d9261d3440 100644 --- a/tests/_save/loaders/test_lazy_wasm.py +++ b/tests/_save/loaders/test_lazy_wasm.py @@ -17,12 +17,11 @@ from marimo._save.cache import MARIMO_CACHE_VERSION, Cache from marimo._save.loaders.lazy import ( - _ACTIVE_LAZY_LOADERS, - _POISONED_KEYS, LazyLoader, LazyStore, WasmLazyLoader, WasmLazyStore, + _cache_state, ) from marimo._save.stores.dict_store import DictStore from marimo._save.stores.file import FileStore @@ -295,8 +294,9 @@ def test_flush_all_flushes_active_loaders(self) -> None: loader1 = LazyLoader("flush_a", store=store1) loader2 = LazyLoader("flush_b", store=store2) - assert _ACTIVE_LAZY_LOADERS.get("flush_a") is loader1 - assert _ACTIVE_LAZY_LOADERS.get("flush_b") is loader2 + loaders = _cache_state().active_lazy_loaders + assert loaders.get("flush_a") is loader1 + assert loaders.get("flush_b") is loader2 # flush_all should not raise LazyLoader.flush_all() @@ -304,7 +304,7 @@ def test_flush_all_flushes_active_loaders(self) -> None: def test_loaders_tracked_by_name(self) -> None: """Loaders are tracked in the active dict by name.""" loader = LazyLoader("track_test", store=LazyStore(inner=DictStore())) - assert _ACTIVE_LAZY_LOADERS["track_test"] is loader + assert _cache_state().active_lazy_loaders["track_test"] is loader def test_store_reused_across_recreations(self) -> None: """When a loader is recreated with the same name, it reuses the @@ -319,6 +319,25 @@ def test_store_reused_across_recreations(self) -> None: assert loader2.store.get("key1") == b"data1" +class TestCacheStateResolution: + def test_resolves_to_root_context(self) -> None: + """Child contexts (embedded apps) share the root's cache state.""" + from types import SimpleNamespace + + from marimo._save.cache import CacheState + + root_cache = CacheState(store=DictStore()) + root = SimpleNamespace(parent=None, cache=root_cache) + child = SimpleNamespace( + parent=root, cache=CacheState(store=DictStore()) + ) + + with mock.patch( + "marimo._save.loaders.lazy.safe_get_context", return_value=child + ): + assert _cache_state() is root_cache + + class TestOnRestoreFailure: """`WasmLazyLoader._on_restore_failure`: on a corrupt restore, evict the manifest and its referenced blobs from the store and poison their keys so @@ -326,17 +345,15 @@ class TestOnRestoreFailure: @pytest.fixture(autouse=True) def _isolate_poisoned_keys(self) -> Iterator[None]: - # _POISONED_KEYS is a module global; snapshot/restore so these tests - # neither leak into nor depend on others. - snapshot = set(_POISONED_KEYS) + # Snapshot/restore so poison doesn't leak across tests. + poisoned = _cache_state().poisoned_keys + snapshot = set(poisoned) yield - _POISONED_KEYS.clear() - _POISONED_KEYS.update(snapshot) + poisoned.clear() + poisoned.update(snapshot) @staticmethod - def _manifest( - refs: list[str], return_ref: str | None = None - ) -> bytes: + def _manifest(refs: list[str], return_ref: str | None = None) -> bytes: meta = Meta(version=MARIMO_CACHE_VERSION) if return_ref is not None: meta = Meta( @@ -347,9 +364,7 @@ def _manifest( CacheSchema( hash="h", cache_type=CacheType("Pure"), - defs={ - f"v{i}": Item(reference=r) for i, r in enumerate(refs) - }, + defs={f"v{i}": Item(reference=r) for i, r in enumerate(refs)}, stateful_refs=[], meta=meta, ) @@ -375,19 +390,17 @@ def test_evicts_and_poisons_manifest_and_blobs(self) -> None: for k in (manifest_path, blob_a, blob_b): assert not store.hit(k), f"{k} not evicted" - assert k in _POISONED_KEYS, f"{k} not poisoned" + assert k in _cache_state().poisoned_keys, f"{k} not poisoned" def test_poisons_return_value_reference(self) -> None: store, loader, key = self._loader_and_key() ret_ref = "h/return.pickle" store.put(ret_ref, b"x") - loader._on_restore_failure( - key, self._manifest([], return_ref=ret_ref) - ) + loader._on_restore_failure(key, self._manifest([], return_ref=ret_ref)) assert not store.hit(ret_ref) - assert ret_ref in _POISONED_KEYS + assert ret_ref in _cache_state().poisoned_keys def test_none_manifest_poisons_only_manifest_path(self) -> None: store, loader, key = self._loader_and_key() @@ -395,9 +408,9 @@ def test_none_manifest_poisons_only_manifest_path(self) -> None: loader._on_restore_failure(key, None) - assert manifest_path in _POISONED_KEYS + assert manifest_path in _cache_state().poisoned_keys # No blob keys to discover, so nothing else is poisoned. - assert _POISONED_KEYS == {manifest_path} + assert _cache_state().poisoned_keys == {manifest_path} def test_undecodable_manifest_poisons_only_manifest_path(self) -> None: store, loader, key = self._loader_and_key() @@ -406,8 +419,8 @@ def test_undecodable_manifest_poisons_only_manifest_path(self) -> None: # Garbage bytes must not raise; manifest path is still poisoned. loader._on_restore_failure(key, b"not valid msgpack/json") - assert manifest_path in _POISONED_KEYS - assert _POISONED_KEYS == {manifest_path} + assert manifest_path in _cache_state().poisoned_keys + assert _cache_state().poisoned_keys == {manifest_path} def test_corrupt_restore_via_load_cache_poisons_keys(self) -> None: # End-to-end: a manifest referencing a missing blob makes restore @@ -426,5 +439,5 @@ def test_corrupt_restore_via_load_cache_poisons_keys(self) -> None: ): assert loader.load_cache(key) is None - assert manifest_path in _POISONED_KEYS - assert missing_ref in _POISONED_KEYS + assert manifest_path in _cache_state().poisoned_keys + assert missing_ref in _cache_state().poisoned_keys From 4065bb68e05a0e6330d661955bd5b0b396910da5 Mon Sep 17 00:00:00 2001 From: "dmadisetti@coreweave.com" Date: Tue, 30 Jun 2026 12:02:08 -0700 Subject: [PATCH 5/6] comments: cleanup brittle langauge --- marimo/_save/loaders/lazy.py | 62 ++++--------------------------- marimo/_save/stores/dict_store.py | 7 +--- marimo/_save/stores/store.py | 4 +- 3 files changed, 10 insertions(+), 63 deletions(-) diff --git a/marimo/_save/loaders/lazy.py b/marimo/_save/loaders/lazy.py index 5cda39bebd0..56285ae5109 100644 --- a/marimo/_save/loaders/lazy.py +++ b/marimo/_save/loaders/lazy.py @@ -58,7 +58,7 @@ def maybe_update_lazy_stub(value: Any) -> str: """Return the loader strategy string for *value*, caching the result. Walks the MRO of `type(value)` against `LAZY_STUB_LOOKUP` (a - fq-class-name → loader-string registry). Falls back to `"pickle"` + fq-class-name -> loader-string registry). Falls back to `"pickle"` when no match is found. """ value_type = type(value) @@ -182,12 +182,7 @@ def from_item(item: Item, var_name: str = "") -> Any: def _cache_state() -> CacheState: - """Cache state for the current session, scoped to the root context. - - Embedded apps run in child contexts (`create_kernel_context(parent=...)`); - walk to the root so the loader registry, poison set, and WASM store are - shared across the session rather than per child. - """ + """Cache state for the current session, scoped to the root context.""" ctx = safe_get_context() if ctx is not None: while ctx.parent is not None: @@ -213,10 +208,7 @@ class LazyStore(Store): """Native store for `LazyLoader`. Delegates to an inner store (default `FileStore` at `__marimo__/cache/`) - and tracks which keys were written or read so `--execute` export knows - exactly which blobs to bundle. Does no environment checks; the WASM - variant (`WasmLazyStore`) adds HTTP-backed reads. - """ + and tracks which keys were written or read.""" def __init__(self, inner: Store | None = None) -> None: self._inner = inner if inner is not None else FileStore() @@ -254,9 +246,6 @@ def export_keys(self) -> list[str]: class WasmLazyStore(LazyStore): """WASM store: writes to a shared in-session `DictStore`; reads fall through to HTTP fetch from `notebook_location()/public/cache/`. - - Instantiated only by `WasmLazyLoader` (selected once via the dual-loader - registry), so it never needs to re-check the environment. """ def __init__(self, inner: Store | None = None) -> None: @@ -291,7 +280,6 @@ def get_batch( if http_keys: yield from self._http_get_batch(http_keys) - # -- HTTP internals -- def _base_url(self) -> str: from marimo._runtime.runtime import notebook_location @@ -301,24 +289,14 @@ def _base_url(self) -> str: @staticmethod def _sanitize_key(key: str) -> str: - """Validate a fetch key against path traversal. - - Returns `key` unchanged so it still matches the stored key; - only rejects `..` segments and absolute paths (normalizing the - string here could desync requested and stored keys). - """ + """Validate a fetch key against path traversal.""" clean = PurePosixPath(key) if ".." in clean.parts or clean.is_absolute(): raise ValueError(f"Invalid cache key: {key}") return key def _http_get(self, key: str) -> bytes | None: - """Single sync fetch via pyodide_http-patched urllib. - - A successful response is written back to the inner store and - marked touched, so repeat reads stay local and the export - manifest covers HTTP-fetched blobs. - """ + """Single sync fetch via pyodide_http-patched urllib.""" import urllib.request url = f"{self._base_url()}/{self._sanitize_key(key)}" @@ -373,8 +351,6 @@ async def _fetch_all() -> list[tuple[str, bytes | None]]: class LazyLoader(BasePersistenceLoader): - # Default store class, overridden by the WASM variant. The single - # native/WASM decision is made once in the dual-loader registry, not here. _store_cls: type[Store] = LazyStore def __init__( @@ -442,7 +418,6 @@ def restore_cache(self, _key: HashKey, blob: bytes) -> Cache: variable_hashes: dict[str, str] = {} # Instances of cell-defined (__main__) classes are deferred: their # class must be re-exec'd into __main__ before the blob can unpickle. - # Cache.restore orders these after their class via `requires`. deferred: dict[str, tuple[str, str]] = {} for var_name, item in cache_data.defs.items(): if var_name in cache_data.ui_defs: @@ -472,9 +447,6 @@ def restore_cache(self, _key: HashKey, blob: bytes) -> Cache: unique_keys = set(ref_vars.values()) if return_ref: unique_keys.add(return_ref) - # Read + deserialize the blobs. The native loader parallelizes via - # threads; the WASM loader overrides this to fetch via the store's - # concurrent get_batch (threads are unavailable in Pyodide). unpickled = self._read_blobs( unique_keys, ref_type_hints, return_ref, return_type_hint ) @@ -549,12 +521,7 @@ def _read_blobs( return_ref: str | None, return_type_hint: str | None, ) -> dict[str, Any]: - """Read + deserialize blobs in parallel via threads. - - Every thread puts exactly one item — value or `_BlobStatus.MISSING` - — so `queue.get()` needs no timeout. `WasmLazyLoader` overrides this - to use the store's concurrent `get_batch` (no threads in Pyodide). - """ + """Read + deserialize blobs in parallel via threads.""" results: queue.Queue[tuple[str, Any]] = queue.Queue() def _load_blob(key: str) -> None: @@ -682,11 +649,6 @@ def _put_or_mark_unserializable( """Store one blob; on serialization failure write no blob and instead mark each manifest `Item` with `unserializable_type`. - A later load reconstructs an `UnhashableStub` tripwire from the - marker (see `from_item`) rather than reading a placeholder blob - off disk — so unserializable values never get dumped to disk and - the cached-execution pre-flight still re-runs the producing cell. - Returns `True` when the blob was stored, `False` when it was marked unserializable. """ @@ -767,9 +729,6 @@ def _serialize_and_write() -> None: def _dispatch_write(self, write_fn: Callable[[], None]) -> None: """Run the blob+manifest write on a background thread (native). - - `WasmLazyLoader` overrides this to run synchronously, since threads - are unavailable in Pyodide. """ t = threading.Thread(target=write_fn, daemon=False) t.start() @@ -783,14 +742,7 @@ def to_blob(self, cache: Cache) -> bytes | None: class WasmLazyLoader(LazyLoader): """WASM variant of `LazyLoader`, selected once via the dual-loader - registry (so the environment is never re-checked below). - - - reads blobs through the store's concurrent `get_batch` (HTTP fetch in - Pyodide) since threads are unavailable; - - writes synchronously for the same reason; - - on a corrupt restore, evicts the blobs and poisons their keys so the - store won't HTTP-re-fetch the same broken data. - """ + registry (so the environment is never re-checked below).""" _store_cls = WasmLazyStore diff --git a/marimo/_save/stores/dict_store.py b/marimo/_save/stores/dict_store.py index 0804164a146..f9fcf4b2665 100644 --- a/marimo/_save/stores/dict_store.py +++ b/marimo/_save/stores/dict_store.py @@ -5,12 +5,7 @@ class DictStore(Store): - """A minimal dict-backed store for in-session caching. - - Used as the inner store for LazyStore in WASM (Pyodide), where - the filesystem is unavailable. Writes persist for the lifetime - of the session; reads fall through to HTTP in the outer LazyStore. - """ + """A minimal dict-backed store for in-session caching.""" def __init__(self) -> None: self._data: dict[str, bytes] = {} diff --git a/marimo/_save/stores/store.py b/marimo/_save/stores/store.py index c086ba65c76..b0c156a56a2 100644 --- a/marimo/_save/stores/store.py +++ b/marimo/_save/stores/store.py @@ -41,8 +41,8 @@ def export_keys(self) -> list[str]: """Return the keys this session wrote or read that should be bundled on `--execute` export. - Defaults to none; stores that track usage (e.g. `LazyStore`) - override this. Returning `[]` keeps non-tracking stores inert. + Defaults to none; stores that track usage override this. Returning `[]` + keeps non-tracking stores inert. """ return [] From 6b1e21f7c86e6383b3570ed96c5d53226d20d5c9 Mon Sep 17 00:00:00 2001 From: "dmadisetti@coreweave.com" Date: Tue, 30 Jun 2026 12:02:49 -0700 Subject: [PATCH 6/6] ruff: format --- marimo/_save/loaders/lazy.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/marimo/_save/loaders/lazy.py b/marimo/_save/loaders/lazy.py index 56285ae5109..aad98fc7586 100644 --- a/marimo/_save/loaders/lazy.py +++ b/marimo/_save/loaders/lazy.py @@ -280,7 +280,6 @@ def get_batch( if http_keys: yield from self._http_get_batch(http_keys) - def _base_url(self) -> str: from marimo._runtime.runtime import notebook_location @@ -728,8 +727,7 @@ def _serialize_and_write() -> None: return True def _dispatch_write(self, write_fn: Callable[[], None]) -> None: - """Run the blob+manifest write on a background thread (native). - """ + """Run the blob+manifest write on a background thread (native).""" t = threading.Thread(target=write_fn, daemon=False) t.start() self._pending.append(t)