Skip to content

fix(save): keep lazy cache restore robust when a cached dep is absent#10065

Closed
dmadisetti wants to merge 2 commits into
mainfrom
dm/lazy-restore-robustness
Closed

fix(save): keep lazy cache restore robust when a cached dep is absent#10065
dmadisetti wants to merge 2 commits into
mainfrom
dm/lazy-restore-robustness

Conversation

@dmadisetti

Copy link
Copy Markdown
Member

This pull request was authored by a coding agent.

Follow-up to #9898 (merged). Three related fixes to the lazy-store restore path, extracted and hardened from WASM cache-restore validation (torch → KANNs demo). All three are independent robustness improvements to the save layer.

1. Restore tripwire

LazyLoader._deserialize_blob now catches ModuleNotFoundError and binds the def as an UnhashableStub instead of aborting the entire restore. A cached value whose codec needs an absent package (e.g. a torch tensor restored where torch isn't installed) stays inert until touched, then raises a clear unhydratable error naming the def. The return value is excluded — a stubbed return fails the restore as a clean miss. Each def sharing a blob (all UI vars share ui.pickle) gets its own re-labeled stub.

2. Module-version pin survives restore

pin_modules=True captured a module's __version__ at cache time but discarded it on restore, so a version-pinned content hash recomputed to module:<ref>: and missed. The version is now persisted (Item.module_version) and replayed onto the MissingModule placeholder; hash.py falls back to the in-scope placeholder when the module isn't in sys.modules. Version is str()-coerced because torch's TorchVersion isn't encodable by the manifest codec.

3. Stale-key invalidation mechanism

Adds CacheState.stale_keys + Loader.mark_stale() (no-op default): a manifest can be forced to miss for the session without being served or re-fetched. This is the primitive the cell-cache preflight (stacked PR #9895) needs — the WASM HTTP store would otherwise re-fetch a cleared key and re-hit the same unusable value in a loop. save_cache clears the mark so a recomputed value self-heals to a hit.

Also resolves nested stubs on restore (an ImmediateReferenceStub whose blob is itself a CustomStub).

Test plan

  • tests/_save/loaders/test_lazy_wasm.py: new TestRestoreTripwire, TestModuleVersionPin, TestStaleKeys (42 pass)
  • tests/_save/ full suite: 327 pass
  • ruff / format / mypy clean on all changed files

🤖 Generated with Claude Code

Three related fixes to the lazy-store restore path, extracted from the
cache-e2e WASM validation work:

- Restore tripwire: `_deserialize_blob` catches `ModuleNotFoundError` and
  binds the def as an `UnhashableStub` instead of aborting the whole
  restore — a torch tensor restored where torch is absent stays inert
  until touched, then raises a clear unhydratable error naming the def.
  The return value is excluded (a stubbed return fails the restore as a
  clean miss). Each def sharing a blob gets its own re-labeled stub.

- Module-version pin survives restore: the pinned `__version__` captured
  at cache time is now persisted (`Item.module_version`) and replayed onto
  the `MissingModule` placeholder, so a version-pinned content hash
  reproduces on restore instead of collapsing to an empty version and
  missing. `hash.py` falls back to the in-scope placeholder when the
  module is not in `sys.modules`. `str()`-coerce the version because
  torch's `TorchVersion` is not encodable by the manifest codec.

- Stale-key invalidation: `CacheState.stale_keys` + `Loader.mark_stale()`
  let a manifest be forced to miss for the session without being served or
  re-fetched (the WASM HTTP store would otherwise just re-fetch a cleared
  key). `save_cache` clears the mark so a recomputed value self-heals to a
  hit. No-op for loaders without a session-scoped store.

Also resolves nested stubs on restore (an `ImmediateReferenceStub` whose
blob is itself a `CustomStub`).
@vercel

vercel Bot commented Jul 2, 2026

Copy link
Copy Markdown

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
marimo-docs Ready Ready Preview, Comment Jul 2, 2026 4:06pm

Request Review

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant