Skip to content

Use marimo export session for session snapshots (and make embedding notebook GPU friendly) #63

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
akshayka merged 1 commit into from
May 29, 2026
Merged

Use marimo export session for session snapshots (and make embedding notebook GPU friendly) #63

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 24 additions & 42 deletions notebooks/algorithms/__marimo__/session/visualizing-embeddings.py.json
View file
Open in desktop

Large diffs are not rendered by default.

5 changes: 3 additions & 2 deletions notebooks/algorithms/visualizing-embeddings.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@

import marimo

__generated_with = "0.20.2"
__generated_with = "0.23.8"
Comment thread
akshayka marked this conversation as resolved.
app = marimo.App(width="medium")

with app.setup:
Expand Down Expand Up @@ -71,13 +71,14 @@ def _():

@app.cell
def _(constraint, embedding_dimension):
embedding = compute_embedding(embedding_dimension, constraint)
embedding = compute_embedding(embedding_dimension, constraint).cpu()
return (embedding,)


@app.cell
def _(embedding):
ax = pymde.plot(embedding, color_by=mnist.attributes["digits"])
plt.tight_layout()
ax = mo.ui.matplotlib(ax)
ax
return (ax,)
Expand Down
12 changes: 9 additions & 3 deletions scripts/create-sessions-changed.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#!/usr/bin/env bash
# Run create-sessions.py in parallel for marimo notebooks.
# Run create-sessions.py (a thin wrapper around `marimo export session`) in
# parallel for marimo notebooks.
#
# Usage:
# bash scripts/create-sessions-changed.sh # only git-changed notebooks
Expand Down Expand Up @@ -29,10 +30,15 @@ if [ "$ALL" = true ]; then
changed_files+=("$f")
done < <(grep -rl "import marimo" --include='*.py' "$REPO_ROOT/notebooks" | sed "s|^$REPO_ROOT/||" | sort)
else
# Collect notebooks with uncommitted changes (staged + unstaged), deduplicated
# Collect changed notebooks (staged + unstaged), deduplicated. Scoped to
# notebooks/ and filtered to existing marimo notebooks so unrelated changed
# .py files (e.g. scripts/) aren't mistaken for notebooks.
while IFS= read -r f; do
case "$f" in *.py) ;; *) continue ;; esac
[ -f "$REPO_ROOT/$f" ] || continue
grep -q "import marimo" "$REPO_ROOT/$f" || continue
changed_files+=("$f")
done < <({ git diff --name-only -- '*.py'; git diff --cached --name-only -- '*.py'; } | sort -u)
done < <({ git diff --name-only -- notebooks; git diff --cached --name-only -- notebooks; } | sort -u)
fi

if [ ${#changed_files[@]} -eq 0 ]; then
Expand Down
208 changes: 42 additions & 166 deletions scripts/create-sessions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,208 +1,84 @@
# /// script
# requires-python = ">=3.12"
# dependencies = [
# "marimo",
# "marimo>=0.23.8",
# ]
# ///
"""Generate session snapshots for marimo notebooks.

Runs marimo notebooks to completion and writes session snapshots to the
__marimo__/session/ directory next to each notebook file. This enables
instant loading of pre-computed outputs.
Thin wrapper around `marimo export session`, the supported CLI for executing
notebooks and writing session snapshots to `__marimo__/session/<notebook>.py.json`
next to each notebook. These snapshots enable instant loading of pre-computed
outputs in the gallery.

Each notebook is executed in its own uv sandbox, using the PEP 723
inline script metadata to resolve dependencies automatically.
Each notebook runs in its own uv sandbox (`--sandbox`, i.e. `uv run --isolated`),
so its PEP 723 inline dependencies are resolved automatically.

Usage:
# Single file
uv run scripts/create-sessions.py notebook.py

# Multiple files/folders
# Multiple files or directories
uv run scripts/create-sessions.py notebooks/ examples/my_notebook.py

# With CLI args passed to notebooks
uv run scripts/create-sessions.py --cli-args '{"key": "value"}' notebooks/
uv run scripts/create-sessions.py --argv '--flag value' notebooks/
# Forward extra args to the notebook(s), parsed via mo.cli_args()
uv run scripts/create-sessions.py notebook.py -- --key value
"""

from __future__ import annotations

import argparse
import atexit
import json
import os
import subprocess
import sys
import tempfile
from pathlib import Path

from marimo._server.files.directory_scanner import is_marimo_app
from marimo._utils.files import expand_file_patterns


def run_single(notebook_path: str, cli_args: dict, argv: list[str] | None) -> None:
"""Run a single notebook and write its session snapshot (in-process)."""
from marimo._server.export import run_app_until_completion
from marimo._server.file_router import AppFileRouter
from marimo._server.utils import asyncio_run
from marimo._session.state.serialize import (
get_session_cache_file,
serialize_session_view,
)
from marimo._utils.marimo_path import MarimoPath

marimo_path = MarimoPath(notebook_path)
file_router = AppFileRouter.from_filename(marimo_path)
file_key = file_router.get_unique_file_key()
assert file_key is not None
file_manager = file_router.get_file_manager(file_key)

session_view, did_error = asyncio_run(
run_app_until_completion(file_manager, cli_args, argv)
)

if did_error:
print(" Warning: notebook had errors during execution")

cell_ids = list(file_manager.app.cell_manager.cell_ids())
session_data = serialize_session_view(
session_view, cell_ids, drop_virtual_file_outputs=True
)

# Treat ModuleNotFoundError as a hard failure — the session cache
# would be useless if a dependency is missing.
serialized = json.dumps(session_data, indent=2)
if "ModuleNotFoundError" in serialized:
print(" Error: notebook has ModuleNotFoundError — missing dependency")
sys.exit(1)

cache_file = get_session_cache_file(Path(notebook_path).resolve())
cache_file.parent.mkdir(parents=True, exist_ok=True)
cache_file.write_text(serialized)

status = "with errors" if did_error else "ok"
print(f" -> {cache_file} ({status})")

if did_error:
sys.exit(2)


def process_notebook_in_sandbox(
notebook_path: Path, cli_args: dict, argv: list[str] | None
) -> tuple[bool, bool]:
"""Spawn a subprocess that runs this script with --single inside a uv sandbox.

Uses marimo's own sandbox helpers to build the uv command, so the
notebook's PEP 723 inline dependencies are resolved automatically.

Returns (succeeded, had_errors).
"""
from marimo._cli.sandbox import construct_uv_flags
from marimo._utils.inline_script_metadata import PyProjectReader
from marimo._utils.uv import find_uv_bin

pyproject = PyProjectReader.from_filename(str(notebook_path))

with tempfile.NamedTemporaryFile(
mode="w", delete=False, suffix=".txt", encoding="utf-8"
) as tmp:
tmp_path = tmp.name
uv_flags = construct_uv_flags(pyproject, tmp, [], [])
atexit.register(lambda p=tmp_path: os.unlink(p))

# Build: uv run <sandbox-flags> python <this-script> --single <notebook> ...
uv_cmd = [find_uv_bin(), "run"] + uv_flags + [
"python",
__file__,
"--single",
str(notebook_path.resolve()),
"--cli-args",
json.dumps(cli_args),
def export_session(path: str, notebook_args: list[str]) -> int:
"""Run `marimo export session --sandbox` for a single file or directory."""
cmd = [
sys.executable,
"-m",
"marimo",
"export",
"session",
"--sandbox",
# Always (re)generate the requested snapshots; freshness gating lives
# in validate-sessions.py.
"--force-overwrite",
path,
*notebook_args,
]
if argv is not None:
uv_cmd.extend(["--argv", " ".join(argv)])
return subprocess.run(cmd).returncode

result = subprocess.run(uv_cmd)

if result.returncode == 0:
return True, False
elif result.returncode == 2:
return True, True # succeeded with errors
else:
return False, False
def main() -> None:
argv = sys.argv[1:]

# Everything after a literal "--" is forwarded to the notebook(s) as argv.
notebook_args: list[str] = []
if "--" in argv:
sep = argv.index("--")
notebook_args = argv[sep + 1 :]
argv = argv[:sep]

def main() -> None:
parser = argparse.ArgumentParser(
description="Generate session snapshots for marimo notebooks.",
)
parser.add_argument(
"paths",
nargs="*",
help="One or more files, directories, or glob patterns",
)
parser.add_argument(
"--cli-args",
default="{}",
help='JSON string of CLI args to pass to notebooks (default: "{}")',
)
parser.add_argument(
"--argv",
default=None,
help="Space-separated argv to pass to notebooks (default: None)",
)
parser.add_argument(
"--single",
default=None,
help="(internal) Run a single notebook in-process and exit.",
nargs="+",
help="One or more notebook files or directories",
)
args = parser.parse_args()
args = parser.parse_args(argv)

# Parse CLI args
try:
cli_args: dict = json.loads(args.cli_args) # type: ignore[type-arg]
except json.JSONDecodeError as e:
print(f"Error: Invalid JSON for --cli-args: {e}", file=sys.stderr)
sys.exit(1)

argv: list[str] | None = None
if args.argv is not None:
argv = args.argv.split()

# --single mode: run one notebook in-process (called from the sandbox subprocess)
if args.single:
run_single(args.single, cli_args, argv)
return

if not args.paths:
parser.error("the following arguments are required: paths")

# Discover files
all_files = expand_file_patterns(tuple(args.paths))
notebooks = [f for f in all_files if is_marimo_app(str(f))]

if not notebooks:
print("No marimo notebooks found in the given paths.")
sys.exit(0)

print(f"Found {len(notebooks)} marimo notebook(s)\n")

succeeded = 0
failed = 0

for notebook_path in notebooks:
print(f"Processing: {notebook_path}")
ok, had_errors = process_notebook_in_sandbox(
Path(notebook_path), cli_args, argv
)
if ok:
succeeded += 1
else:
print(f" Error: notebook processing failed", file=sys.stderr)
for path in args.paths:
print(f"Processing: {path}")
if export_session(path, notebook_args) != 0:
print(f" Error: failed to export session for {path}", file=sys.stderr)
failed += 1

print(f"\nDone: {succeeded} succeeded, {failed} failed")
if failed > 0:
if failed:
print(f"\n{failed} path(s) failed.")
sys.exit(1)


Expand Down
Loading
Loading