From 071913c1bcbbce9832c4733e21ce4d16a506b42f Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Wed, 13 May 2026 12:16:25 -0400 Subject: [PATCH 01/22] add marimo MNIST + W&B Registry example Introduces the first marimo notebook in this repo. Trains a small PyTorch CNN on MNIST, logs the run with wandb, saves the trained model as an Artifact, and links it to a W&B Registry collection via the cross-org link_artifact API. The notebook is interactive: hyperparameters live in a marimo form and training is gated by an explicit Train button so slider tweaks do not trigger runs. PEP 723 inline script metadata makes it self-installing under `uvx marimo edit --sandbox`. Co-Authored-By: Claude Opus 4.7 (1M context) --- examples/marimo/mnist-registry/README.md | 90 +++ .../marimo/mnist-registry/mnist_registry.py | 647 ++++++++++++++++++ .../marimo/mnist-registry/requirements.txt | 7 + 3 files changed, 744 insertions(+) create mode 100644 examples/marimo/mnist-registry/README.md create mode 100644 examples/marimo/mnist-registry/mnist_registry.py create mode 100644 examples/marimo/mnist-registry/requirements.txt diff --git a/examples/marimo/mnist-registry/README.md b/examples/marimo/mnist-registry/README.md new file mode 100644 index 00000000..e78cc6eb --- /dev/null +++ b/examples/marimo/mnist-registry/README.md @@ -0,0 +1,90 @@ +# MNIST -> W&B Registry (marimo) + +A [marimo](https://marimo.io) notebook that trains a small CNN on MNIST with +PyTorch, tracks the run in Weights & Biases, saves the trained model as a W&B +Artifact, and links that Artifact to a collection in the **W&B Registry**. + +The notebook is the first marimo example in this repo and is intentionally +self-contained: dependencies are declared in a [PEP 723](https://peps.python.org/pep-0723/) +inline-script block at the top of `mnist_registry.py`, so [`uv`](https://docs.astral.sh/uv/) +can resolve them automatically. + +## Prerequisites + +- Python 3.10 or newer. +- A W&B account. Run `wandb login` once in your shell before launching the + notebook — this notebook does not prompt for an API key interactively. +- A W&B **Registry** must exist in your org for the final linking step. The + built-in Model registry is provisioned automatically in newer orgs. If + linking fails, the notebook surfaces a remediation message in the last + Registry cell instead of crashing. +- GPU is optional. Defaults are tuned to finish in roughly two minutes on CPU. + +## Run + +The recommended entry point is `uvx` with marimo's sandbox mode — it +creates an isolated venv from the inline dependencies in the notebook: + +```bash +uvx marimo edit mnist_registry.py --sandbox +``` + +Marimo opens in your browser. Adjust hyperparameters in the form, then click +**Train model** to start the run. The run URL appears inline as soon as +training begins. + +If you prefer pip: + +```bash +pip install -r requirements.txt +marimo edit mnist_registry.py +``` + +The notebook is interactive-only by design: training is gated by a button +click, so `marimo run` will render the form but never start training without +an explicit click. + +## What you get + +After a successful run: + +- A W&B run with training and test metrics, gradient histograms (`wandb.watch`), + and up to 16 example test-set predictions logged as images. +- A model Artifact named `mnist-cnn-` of type `model` with metadata + for test accuracy, parameter count, dataset sizes, and the full + hyperparameter dict. Tagged with the `latest` alias. +- A version of that Artifact linked into the configured Registry collection + (default: `wandb-registry-model/MNIST Classifiers`). + +To consume the registered model from another script or notebook: + +```python +import wandb +api = wandb.Api() +art = api.artifact("wandb-registry-model/MNIST Classifiers:latest") +art.download() # writes mnist_cnn.pt under ./artifacts/ +``` + +## Design notes + +- **Training is gated by a button.** Marimo cells re-run reactively when their + inputs change. Before the first click of **Train model**, slider changes do + not start a run. After a run has completed, clicking **Train model** again + starts a new run with whatever the form values are at that moment; the + previous run is finished cleanly first. +- **`wandb.run` is finished defensively** at the top of the training cell so + the second click of **Train model** does not nest runs in the same marimo + kernel. +- **`logged.wait()` is called** after `log_artifact` and before + `link_artifact` to avoid a race where the link tries to resolve a version + that has not finished committing server-side. +- **Registry failures soft-fail.** If `link_artifact` raises — usually + because the Registry does not exist in your org — the notebook + surfaces remediation guidance via `mo.callout` rather than aborting. + +## Reference + +The CNN architecture and training loop mirror +[`examples/pytorch/pytorch-cnn-mnist/main.py`](../../pytorch/pytorch-cnn-mnist/main.py). +The Registry linking pattern follows +[`colabs/wandb_registry/zoo_wandb.ipynb`](../../../colabs/wandb_registry/zoo_wandb.ipynb). diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py new file mode 100644 index 00000000..fdb2eb93 --- /dev/null +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -0,0 +1,647 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "marimo>=0.9", +# "torch>=2.1", +# "torchvision>=0.16", +# "wandb>=0.18", +# "tqdm", +# ] +# /// +"""Train an MNIST CNN with PyTorch, track the run with Weights & Biases, +and link the resulting model artifact to a W&B Registry collection. + +Run: + + uvx marimo edit mnist_registry.py --sandbox + +This notebook is interactive: hyperparameters live in a form, and training is +gated by a button so slider changes do not trigger runs. +""" + +import marimo + +__generated_with = "0.9.0" +app = marimo.App(width="medium", app_title="MNIST -> W&B Registry") + + +@app.cell(hide_code=True) +def _(): + import marimo as mo + + mo.md( + """ + # MNIST -> W&B Run -> Registry + + ## What you will build + + - A **W&B run** with training and test metrics, gradient histograms, + and example test-set predictions logged as images. + - A **model Artifact** named `mnist-cnn-` of type `model`, + carrying metadata (test accuracy, parameter count, hyperparameters). + - A version of that Artifact **linked into a W&B Registry collection** + so it appears under registered models org-wide. + + ## Prerequisites + + - **`wandb login`** completed in your shell before starting marimo. + This notebook will not prompt for an API key interactively. + - A W&B entity (your user or a team) the run will be written to. + - A **W&B Registry** must exist in your org. The built-in Model + registry is provisioned automatically in newer orgs. If linking + fails, the Registry step surfaces remediation guidance inline + instead of crashing. + - A GPU is optional. The defaults are tuned to finish in ~2 minutes + on CPU. + """ + ) + return (mo,) + + +@app.cell +def _(): + import os + + import torch + import torch.nn as nn + import torch.nn.functional as F + import torch.optim as optim + from torch.utils.data import DataLoader + from torchvision import datasets, transforms + + import wandb + from tqdm.auto import tqdm + + return ( + DataLoader, + F, + datasets, + nn, + optim, + os, + torch, + tqdm, + transforms, + wandb, + ) + + +@app.cell +def _(mo, torch): + if torch.cuda.is_available(): + device = torch.device("cuda") + device_note = "CUDA GPU detected. Training will be fast." + callout_kind = "success" + elif torch.backends.mps.is_available(): + device = torch.device("mps") + device_note = "Apple MPS detected. Training will run on the GPU." + callout_kind = "success" + else: + device = torch.device("cpu") + device_note = ( + "No GPU detected. Training will run on CPU. With the default " + "hyperparameters this takes ~2 minutes." + ) + callout_kind = "warn" + + mo.callout(mo.md(f"**Device:** `{device}` — {device_note}"), kind=callout_kind) + return (device,) + + +@app.cell(hide_code=True) +def _(mo): + mo.md( + """ + ## Hyperparameters + + Configure the training run and the Registry target. The defaults reach + roughly 98% test accuracy in about two minutes on CPU. The **Registry** + section controls where the trained model is linked after training + finishes. + """ + ) + return + + +@app.cell +def _(mo): + epochs = mo.ui.slider(start=1, stop=10, step=1, value=3, label="Epochs") + batch_size = mo.ui.dropdown( + options=["32", "64", "128", "256"], value="64", label="Batch size" + ) + lr = mo.ui.slider( + start=0.001, stop=0.1, step=0.001, value=0.01, label="Learning rate", show_value=True + ) + momentum = mo.ui.slider( + start=0.0, stop=0.99, step=0.01, value=0.5, label="SGD momentum", show_value=True + ) + seed = mo.ui.number(start=0, stop=99999, value=42, label="Random seed") + + project = mo.ui.text(value="marimo-mnist-registry", label="W&B project") + entity = mo.ui.text(value="", label="W&B entity (blank uses your default)") + run_name = mo.ui.text(value="", label="Run name (blank auto-generates)") + + registry_name = mo.ui.text(value="model", label="W&B Registry name") + collection_name = mo.ui.text(value="MNIST Classifiers", label="Registry collection") + link_to_registry = mo.ui.checkbox(value=True, label="Link artifact to Registry") + + form = mo.vstack( + [ + mo.md("### Training"), + mo.hstack([epochs, batch_size]), + mo.hstack([lr, momentum]), + seed, + mo.md("### W&B run"), + mo.hstack([project, entity, run_name]), + mo.md("### Registry"), + mo.hstack([registry_name, collection_name, link_to_registry]), + ] + ) + form + return ( + batch_size, + collection_name, + entity, + epochs, + link_to_registry, + lr, + momentum, + project, + registry_name, + run_name, + seed, + ) + + +@app.cell +def _( + batch_size, + collection_name, + entity, + epochs, + link_to_registry, + lr, + momentum, + project, + registry_name, + run_name, + seed, +): + config = { + "epochs": epochs.value, + "batch_size": int(batch_size.value), + "lr": lr.value, + "momentum": momentum.value, + "seed": seed.value, + "architecture": "CNN", + "dataset": "MNIST", + } + wandb_project = project.value or None + wandb_entity = entity.value or None + wandb_run_name = run_name.value or None + registry_name_v = registry_name.value.strip() + collection_name_v = collection_name.value.strip() + link_to_registry_v = link_to_registry.value + return ( + collection_name_v, + config, + link_to_registry_v, + registry_name_v, + wandb_entity, + wandb_project, + wandb_run_name, + ) + + +@app.cell(hide_code=True) +def _(mo): + mo.md( + """ + ## Train + + Training is gated by an explicit button click so changing a + hyperparameter does not by itself start a run. Click **Train model** + to begin. Once the first run has completed, clicking the button again + starts a new run with whatever the form values are at that moment; + the previous run is finished cleanly first. + """ + ) + return + + +@app.cell +def _(mo): + train_button = mo.ui.run_button(label="Train model", kind="success") + train_button + return (train_button,) + + +@app.cell +def _(F, nn): + class Net(nn.Module): + """The same small CNN used in examples/pytorch/pytorch-cnn-mnist. + + Two convolutional layers (10 and 20 filters, 5x5 kernels) feed into two + fully connected layers (50 hidden units, 10 outputs). Roughly 21k + parameters. + """ + + def __init__(self): + super().__init__() + self.conv1 = nn.Conv2d(1, 10, kernel_size=5) + self.conv2 = nn.Conv2d(10, 20, kernel_size=5) + self.conv2_drop = nn.Dropout2d() + self.fc1 = nn.Linear(320, 50) + self.fc2 = nn.Linear(50, 10) + + def forward(self, x): + x = F.relu(F.max_pool2d(self.conv1(x), 2)) + x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) + x = x.view(-1, 320) + x = F.relu(self.fc1(x)) + x = F.dropout(x, training=self.training) + x = self.fc2(x) + return F.log_softmax(x, dim=1) + + return (Net,) + + +@app.cell +def _(DataLoader, batch_size, datasets, device, mo, transforms): + transform = transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] + ) + train_ds = datasets.MNIST("./data", train=True, download=True, transform=transform) + test_ds = datasets.MNIST("./data", train=False, download=True, transform=transform) + + # Only batch_size and device affect the loaders, so we depend on them + # directly rather than the full config dict; this avoids re-creating the + # loaders whenever an unrelated hyperparameter changes. + bs = int(batch_size.value) + loader_kwargs = ( + {"num_workers": 2, "pin_memory": True} if device.type == "cuda" else {} + ) + train_loader = DataLoader( + train_ds, batch_size=bs, shuffle=True, **loader_kwargs + ) + test_loader = DataLoader( + test_ds, batch_size=1000, shuffle=False, **loader_kwargs + ) + + mo.md( + f"**Train:** {len(train_ds):,} examples · " + f"**Test:** {len(test_ds):,} examples · " + f"**Batch size:** {bs}" + ) + return test_ds, test_loader, train_ds, train_loader + + +@app.cell +def _( + Net, + config, + device, + mo, + optim, + torch, + train_button, + wandb, + wandb_entity, + wandb_project, + wandb_run_name, +): + mo.stop(not train_button.value, mo.md("Click **Train model** to begin.")) + + # Defensive: finish any prior run still attached to this Python process. + # marimo keeps the kernel alive across re-clicks, so a second click — or + # a slider change after the first click — re-executes this cell. Without + # this guard `wandb.init` would warn about a run already being active. + # `wandb.finish` blocks until the prior run's tail logs are uploaded. + if wandb.run is not None: + wandb.finish() + + torch.manual_seed(config["seed"]) + + run = wandb.init( + project=wandb_project, + entity=wandb_entity, + name=wandb_run_name, + config=config, + job_type="train", + ) + + # Use `epoch` as the x-axis for train and test metrics in the W&B UI. + wandb.define_metric("epoch") + wandb.define_metric("train/*", step_metric="epoch") + wandb.define_metric("test/*", step_metric="epoch") + + model = Net().to(device) + # `log="gradients"` is the conventional choice for didactic examples; + # `log="all"` would additionally log parameter histograms at extra cost. + wandb.watch(model, log="gradients", log_freq=100) + optimizer = optim.SGD( + model.parameters(), lr=config["lr"], momentum=config["momentum"] + ) + + mo.md(f"Run started: [`{run.name}`]({run.url})") + return model, optimizer, run + + +@app.cell +def _( + F, + config, + device, + mo, + model, + optimizer, + test_loader, + torch, + tqdm, + train_button, + train_loader, + wandb, +): + mo.stop(not train_button.value, mo.md("")) + + history = [] + best_acc = 0.0 + final_acc = 0.0 + final_loss = 0.0 + + for epoch in range(1, config["epochs"] + 1): + # ---- train ---- + model.train() + for batch_idx, (data, target) in enumerate( + tqdm(train_loader, desc=f"epoch {epoch}/{config['epochs']}") + ): + data, target = data.to(device), target.to(device) + optimizer.zero_grad() + output = model(data) + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % 50 == 0: + wandb.log({"train/loss": loss.item(), "epoch": epoch}) + + # ---- test ---- + model.eval() + test_loss = 0.0 + correct = 0 + example_images = [] + with torch.no_grad(): + for data, target in test_loader: + data, target = data.to(device), target.to(device) + output = model(data) + test_loss += F.nll_loss(output, target, reduction="sum").item() + pred = output.argmax(dim=1, keepdim=True) + correct += pred.eq(target.view_as(pred)).sum().item() + # Pull up to 16 example predictions from the first batch we see. + while len(example_images) < 16 and len(example_images) < data.size(0): + j = len(example_images) + example_images.append( + wandb.Image( + data[j], + caption=( + f"pred={pred[j].item()} " + f"true={target[j].item()}" + ), + ) + ) + + test_loss /= len(test_loader.dataset) + test_acc = correct / len(test_loader.dataset) + best_acc = max(best_acc, test_acc) + final_acc = test_acc + final_loss = test_loss + wandb.log( + { + "test/loss": test_loss, + "test/accuracy": test_acc, + "epoch": epoch, + "examples": example_images, + } + ) + history.append( + {"epoch": epoch, "test_loss": test_loss, "test_acc": test_acc} + ) + + return best_acc, final_acc, final_loss, history + + +@app.cell +def _(final_acc, final_loss, history, mo, train_button): + mo.stop(not train_button.value, mo.md("")) + mo.vstack( + [ + mo.md("### Training summary"), + mo.ui.table(history, selection=None), + mo.md( + f"**Final test accuracy:** {final_acc:.2%} · " + f"**Final test loss:** {final_loss:.4f}" + ), + ] + ) + return + + +@app.cell +def _(mo, model, os, torch, train_button): + mo.stop(not train_button.value, mo.md("")) + + model_path = "mnist_cnn.pt" + torch.save(model.state_dict(), model_path) + + mo.md( + f"Saved `{model_path}` ({os.path.getsize(model_path) / 1024:.1f} KB)" + ) + return (model_path,) + + +@app.cell +def _( + best_acc, + config, + final_acc, + mo, + model, + model_path, + run, + test_ds, + train_button, + train_ds, + wandb, +): + mo.stop(not train_button.value, mo.md("")) + + num_params = sum(p.numel() for p in model.parameters()) + + artifact = wandb.Artifact( + name=f"mnist-cnn-{run.id}", + type="model", + description=( + "Small CNN trained on MNIST. Architecture: 2 conv layers " + "(10 and 20 filters, 5x5 kernels) + 2 FC layers (50, 10)." + ), + metadata={ + "framework": "pytorch", + "architecture": "CNN", + "num_parameters": num_params, + "dataset": "MNIST", + "train_size": len(train_ds), + "test_size": len(test_ds), + "test_accuracy": final_acc, + "best_test_accuracy": best_acc, + "hyperparameters": dict(config), + }, + ) + artifact.add_file(model_path) + + # We only log a single artifact per run (the final-epoch weights), so we + # tag it `latest` unconditionally. Use the Registry UI or the API to + # promote a specific version with aliases like `best` or `production` + # after comparing across runs. + aliases = ["latest"] + + logged = run.log_artifact(artifact, aliases=aliases) + # Block until the artifact has fully committed server-side. Without this, + # link_artifact below may race on the version reference. + logged.wait() + + mo.md( + f"Artifact logged: `{artifact.name}` with aliases `{aliases}`" + ) + return (logged,) + + +@app.cell +def _( + collection_name_v, + link_to_registry_v, + logged, + mo, + registry_name_v, + run, + train_button, +): + mo.stop(not train_button.value, mo.md("")) + mo.stop( + not link_to_registry_v, + mo.md( + "_Registry linking is disabled (checkbox unchecked). " + "The artifact is logged to the run but not linked to a Registry " + "collection._" + ), + ) + + target_path = f"wandb-registry-{registry_name_v}/{collection_name_v}" + + try: + run.link_artifact(artifact=logged, target_path=target_path) + link_result = mo.callout( + mo.md( + f"**Linked to Registry:** `{target_path}`\n\n" + f"Open the Registry at " + f"[https://wandb.ai/registry](https://wandb.ai/registry) to " + f"see the version." + ), + kind="success", + ) + except Exception as exc: # noqa: BLE001 - we want to surface any failure to the reader + link_result = mo.callout( + mo.md( + f"**Registry link failed.**\n\n" + f"Target path: `{target_path}`\n\n" + f"Error: `{exc}`\n\n" + f"Common causes:\n\n" + f"- The Registry `{registry_name_v}` does not exist in your " + f"org. An org admin can create the Model registry from the " + f"W&B Registry UI.\n" + f"- Your account lacks Registry write permission.\n" + f"- Your org is on the legacy Model Registry. In that case " + f"use the legacy pattern:\n\n" + f" ```python\n" + f" run.link_artifact(\n" + f" logged,\n" + f" target_path='model-registry/{collection_name_v}',\n" + f" )\n" + f" ```" + ), + kind="danger", + ) + link_result + return + + +@app.cell(hide_code=True) +def _(collection_name_v, mo, registry_name_v, run, train_button): + mo.stop(not train_button.value, mo.md("")) + mo.md( + f""" + ## Verify + + 1. Open the run page: [{run.name}]({run.url}). Confirm the + **Charts**, **System**, and **Examples** panels populated. + 2. Click **Artifacts** in the run's left nav. Confirm the + `mnist-cnn-{run.id}` model artifact is listed with metadata + (test accuracy, hyperparameters, number of parameters). + 3. Go to [wandb.ai/registry](https://wandb.ai/registry), open the + **{registry_name_v.title()}** registry, then the + **{collection_name_v}** collection. Confirm the linked version + is present. + + ## Consume the registered model + + From any other script or notebook, fetch the latest registered version: + + ```python + import wandb + api = wandb.Api() + art = api.artifact( + "wandb-registry-{registry_name_v}/{collection_name_v}:latest" + ) + art.download() # writes mnist_cnn.pt under ./artifacts/.../ + ``` + + ## Next steps + + - **Promote a version.** From the Registry UI, add the `production` + alias to the version you want consumers to pick up. The same + collection path with `:production` will then resolve to it. + - **Compare runs.** Re-run with a deeper architecture or a different + learning rate. Group runs in the W&B UI to compare test accuracy + across configurations. + - **Automate on promotion.** Configure a W&B Automation on the + collection to trigger evaluation jobs or webhooks when a new + version is linked. + """ + ) + return + + +@app.cell(hide_code=True) +def _(mo): + mo.md( + """ + ## Finish + + Closes the W&B run so the run summary and the Registry version + finalize on the server. + """ + ) + return + + +@app.cell +def _(mo, train_button, wandb): + mo.stop(not train_button.value, mo.md("")) + # Mirror of the defensive `wandb.finish` at the top of the training cell: + # leaves the kernel in a clean state for the next Train click. + if wandb.run is not None: + wandb.finish() + mo.md("Run finished. Click **Train model** again to start a new run.") + return + + +if __name__ == "__main__": + app.run() diff --git a/examples/marimo/mnist-registry/requirements.txt b/examples/marimo/mnist-registry/requirements.txt new file mode 100644 index 00000000..69a57c0a --- /dev/null +++ b/examples/marimo/mnist-registry/requirements.txt @@ -0,0 +1,7 @@ +# Mirror of the PEP 723 inline dependency block in mnist_registry.py. +# Keep these two in sync. +marimo>=0.9 +torch>=2.1 +torchvision>=0.16 +wandb>=0.18 +tqdm From a690e06aa81f4dc2d75a88e398756e7141d97cd6 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Wed, 13 May 2026 15:16:23 -0400 Subject: [PATCH 02/22] apply Google Developer Style Guide pass to marimo MNIST example Active voice, present tense, plain English, and consistent brand capitalization across the README and the notebook's prose cells. Code identifiers wrapped in backticks where missing; a sentence fragment in the verify step rewritten as a complete clause. Co-Authored-By: Claude Opus 4.7 (1M context) --- examples/marimo/mnist-registry/README.md | 30 +++++------ .../marimo/mnist-registry/mnist_registry.py | 52 +++++++++---------- 2 files changed, 40 insertions(+), 42 deletions(-) diff --git a/examples/marimo/mnist-registry/README.md b/examples/marimo/mnist-registry/README.md index e78cc6eb..ad281ee5 100644 --- a/examples/marimo/mnist-registry/README.md +++ b/examples/marimo/mnist-registry/README.md @@ -15,21 +15,21 @@ can resolve them automatically. - A W&B account. Run `wandb login` once in your shell before launching the notebook — this notebook does not prompt for an API key interactively. - A W&B **Registry** must exist in your org for the final linking step. The - built-in Model registry is provisioned automatically in newer orgs. If + built-in Model Registry is provisioned automatically in newer orgs. If linking fails, the notebook surfaces a remediation message in the last Registry cell instead of crashing. - GPU is optional. Defaults are tuned to finish in roughly two minutes on CPU. ## Run -The recommended entry point is `uvx` with marimo's sandbox mode — it -creates an isolated venv from the inline dependencies in the notebook: +Use `uvx` with marimo's sandbox mode — it creates an isolated virtual +environment from the inline dependencies in the notebook: ```bash uvx marimo edit mnist_registry.py --sandbox ``` -Marimo opens in your browser. Adjust hyperparameters in the form, then click +marimo opens in your browser. Adjust hyperparameters in the form, then click **Train model** to start the run. The run URL appears inline as soon as training begins. @@ -41,7 +41,7 @@ marimo edit mnist_registry.py ``` The notebook is interactive-only by design: training is gated by a button -click, so `marimo run` will render the form but never start training without +click, so `marimo run` renders the form but never starts training without an explicit click. ## What you get @@ -67,20 +67,20 @@ art.download() # writes mnist_cnn.pt under ./artifacts/ ## Design notes -- **Training is gated by a button.** Marimo cells re-run reactively when their +- **Training is gated by a button.** marimo cells re-run reactively when their inputs change. Before the first click of **Train model**, slider changes do - not start a run. After a run has completed, clicking **Train model** again - starts a new run with whatever the form values are at that moment; the - previous run is finished cleanly first. -- **`wandb.run` is finished defensively** at the top of the training cell so - the second click of **Train model** does not nest runs in the same marimo + not start a run. After a run completes, clicking **Train model** again + starts a new run with the current form values; the previous run finishes + cleanly first. +- **`wandb.run` finishes defensively** at the top of the training cell so + a second click of **Train model** does not nest runs in the same marimo kernel. -- **`logged.wait()` is called** after `log_artifact` and before - `link_artifact` to avoid a race where the link tries to resolve a version - that has not finished committing server-side. +- **`logged.wait()` runs** after `log_artifact` and before `link_artifact` + to avoid a race where the link tries to resolve a version that has not + finished committing server-side. - **Registry failures soft-fail.** If `link_artifact` raises — usually because the Registry does not exist in your org — the notebook - surfaces remediation guidance via `mo.callout` rather than aborting. + surfaces remediation guidance through `mo.callout` rather than aborting. ## Reference diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index fdb2eb93..6217d8aa 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -44,15 +44,14 @@ def _(): ## Prerequisites - - **`wandb login`** completed in your shell before starting marimo. - This notebook will not prompt for an API key interactively. + - Run **`wandb login`** in your shell before starting marimo. + This notebook does not prompt for an API key interactively. - A W&B entity (your user or a team) the run will be written to. - A **W&B Registry** must exist in your org. The built-in Model registry is provisioned automatically in newer orgs. If linking - fails, the Registry step surfaces remediation guidance inline + fails, the Registry step shows remediation guidance inline instead of crashing. - - A GPU is optional. The defaults are tuned to finish in ~2 minutes - on CPU. + - A GPU is optional. The defaults finish in about 2 minutes on CPU. """ ) return (mo,) @@ -100,7 +99,7 @@ def _(mo, torch): device = torch.device("cpu") device_note = ( "No GPU detected. Training will run on CPU. With the default " - "hyperparameters this takes ~2 minutes." + "hyperparameters this takes about 2 minutes." ) callout_kind = "warn" @@ -219,11 +218,10 @@ def _(mo): """ ## Train - Training is gated by an explicit button click so changing a - hyperparameter does not by itself start a run. Click **Train model** - to begin. Once the first run has completed, clicking the button again - starts a new run with whatever the form values are at that moment; - the previous run is finished cleanly first. + Click **Train model** to begin. Changing a hyperparameter does not + start a run by itself — the button gates execution. Once a run + completes, clicking the button again starts a new run using the + current form values; the previous run finishes cleanly first. """ ) return @@ -274,8 +272,8 @@ def _(DataLoader, batch_size, datasets, device, mo, transforms): train_ds = datasets.MNIST("./data", train=True, download=True, transform=transform) test_ds = datasets.MNIST("./data", train=False, download=True, transform=transform) - # Only batch_size and device affect the loaders, so we depend on them - # directly rather than the full config dict; this avoids re-creating the + # Only `batch_size` and `device` affect the loaders, so we depend on them + # directly rather than the full `config` dict; this avoids re-creating the # loaders whenever an unrelated hyperparameter changes. bs = int(batch_size.value) loader_kwargs = ( @@ -312,10 +310,10 @@ def _( ): mo.stop(not train_button.value, mo.md("Click **Train model** to begin.")) - # Defensive: finish any prior run still attached to this Python process. + # Finish any prior run still attached to this Python process. # marimo keeps the kernel alive across re-clicks, so a second click — or # a slider change after the first click — re-executes this cell. Without - # this guard `wandb.init` would warn about a run already being active. + # this guard, `wandb.init` warns about a run already being active. # `wandb.finish` blocks until the prior run's tail logs are uploaded. if wandb.run is not None: wandb.finish() @@ -330,14 +328,14 @@ def _( job_type="train", ) - # Use `epoch` as the x-axis for train and test metrics in the W&B UI. + # Use `epoch` as the x-axis for training and test metrics in the W&B UI. wandb.define_metric("epoch") wandb.define_metric("train/*", step_metric="epoch") wandb.define_metric("test/*", step_metric="epoch") model = Net().to(device) - # `log="gradients"` is the conventional choice for didactic examples; - # `log="all"` would additionally log parameter histograms at extra cost. + # `log="gradients"` is the standard choice for tutorial examples; + # `log="all"` also logs parameter histograms at extra cost. wandb.watch(model, log="gradients", log_freq=100) optimizer = optim.SGD( model.parameters(), lr=config["lr"], momentum=config["momentum"] @@ -497,15 +495,15 @@ def _( ) artifact.add_file(model_path) - # We only log a single artifact per run (the final-epoch weights), so we - # tag it `latest` unconditionally. Use the Registry UI or the API to - # promote a specific version with aliases like `best` or `production` - # after comparing across runs. + # Log a single artifact per run (the final-epoch weights) and tag it + # `latest` unconditionally. Use the Registry UI or the API to promote a + # specific version with aliases like `best` or `production` after + # comparing runs. aliases = ["latest"] logged = run.log_artifact(artifact, aliases=aliases) # Block until the artifact has fully committed server-side. Without this, - # link_artifact below may race on the version reference. + # `link_artifact` below may race on the version reference. logged.wait() mo.md( @@ -581,7 +579,7 @@ def _(collection_name_v, mo, registry_name_v, run, train_button): ## Verify 1. Open the run page: [{run.name}]({run.url}). Confirm the - **Charts**, **System**, and **Examples** panels populated. + **Charts**, **System**, and **Examples** panels are populated. 2. Click **Artifacts** in the run's left nav. Confirm the `mnist-cnn-{run.id}` model artifact is listed with metadata (test accuracy, hyperparameters, number of parameters). @@ -625,8 +623,8 @@ def _(mo): """ ## Finish - Closes the W&B run so the run summary and the Registry version - finalize on the server. + This cell closes the W&B run so the run summary and the Registry + version finalize on the server. """ ) return @@ -635,7 +633,7 @@ def _(mo): @app.cell def _(mo, train_button, wandb): mo.stop(not train_button.value, mo.md("")) - # Mirror of the defensive `wandb.finish` at the top of the training cell: + # Mirrors the `wandb.finish` guard at the top of the training cell: # leaves the kernel in a clean state for the next Train click. if wandb.run is not None: wandb.finish() From c4c72880f3a8b04dd600079ada7b80bd9738ff4e Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 13:03:20 -0400 Subject: [PATCH 03/22] add workflow posting molab run-links for modified marimo notebooks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the existing "Get Colabs" PR-comment flow (which serves .ipynb files via nb_helpers) for marimo .py notebooks, which nb_helpers does not handle. On each PR that touches a .py file, the job detects marimo notebooks by their `marimo.App(` marker and upserts a single comment linking each to molab — marimo's hosted runtime that loads any public notebook on GitHub. Uses pull_request_target with pull-requests:write so the comment also lands on fork PRs, and never checks out or runs PR code (reads file content as text only). Links pin to the head commit SHA because branch names with slashes make the molab blob// split ambiguous. Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/marimo_molab.yml | 135 +++++++++++++++++++++++++++++ 1 file changed, 135 insertions(+) create mode 100644 .github/workflows/marimo_molab.yml diff --git a/.github/workflows/marimo_molab.yml b/.github/workflows/marimo_molab.yml new file mode 100644 index 00000000..e2184827 --- /dev/null +++ b/.github/workflows/marimo_molab.yml @@ -0,0 +1,135 @@ +name: marimo molab links + +# Posts — and keeps updated — a PR comment linking each modified marimo +# notebook to molab (https://molab.marimo.io), which runs any public marimo +# notebook on GitHub in a hosted environment with no local setup. +# +# Security note: this uses `pull_request_target` so the comment can also be +# posted on PRs from forks (a plain `pull_request` event gives fork PRs a +# read-only token that cannot comment). The job NEVER checks out or executes +# PR code — it only reads changed-file metadata and file contents as text +# through the API, then posts a comment. Do not add a checkout of the PR head +# or run any PR-provided code in this workflow. + +on: + pull_request_target: + types: [opened, synchronize, reopened] + paths: + - '**.py' + +permissions: + contents: read + pull-requests: write + +jobs: + molab-links: + runs-on: ubuntu-latest + steps: + - name: Comment molab links for modified marimo notebooks + uses: actions/github-script@v7 + with: + script: | + const pr = context.payload.pull_request; + const headOwner = pr.head.repo.owner.login; + const headRepo = pr.head.repo.name; + const headSha = pr.head.sha; + const marker = ''; + + // 1. List the files changed in this PR. + const files = await github.paginate(github.rest.pulls.listFiles, { + owner: context.repo.owner, + repo: context.repo.repo, + pull_number: pr.number, + per_page: 100, + }); + + // 2. Keep added/modified .py files and decide whether each is a + // marimo notebook by inspecting its content (never executing it). + // Every marimo notebook constructs `marimo.App(...)`. + const isMarimo = /\bmarimo\.App\s*\(/; + const notebooks = []; + for (const f of files) { + if (f.status === 'removed') continue; + if (!f.filename.endsWith('.py')) continue; + try { + const res = await github.rest.repos.getContent({ + owner: headOwner, + repo: headRepo, + path: f.filename, + ref: headSha, + }); + if (!res.data.content) { + core.warning(`Skipping ${f.filename}: content not inlined (file too large?).`); + continue; + } + const content = Buffer.from(res.data.content, res.data.encoding).toString('utf8'); + if (isMarimo.test(content)) notebooks.push(f.filename); + } catch (err) { + core.warning(`Could not read ${f.filename}: ${err.message}`); + } + } + + // 3. Find any prior comment so we update it in place instead of + // posting a new one on every push. + const comments = await github.paginate(github.rest.issues.listComments, { + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: pr.number, + per_page: 100, + }); + const existing = comments.find(c => c.body && c.body.includes(marker)); + + // 4. No marimo notebooks: clear a stale comment if present, else exit. + if (notebooks.length === 0) { + if (existing) { + await github.rest.issues.updateComment({ + owner: context.repo.owner, + repo: context.repo.repo, + comment_id: existing.id, + body: `${marker}\n_No marimo notebooks in the current changes._`, + }); + } + core.info('No marimo notebooks found; nothing to link.'); + return; + } + + // 5. Build the comment. Links pin to the head commit SHA so they + // always resolve — branch names containing slashes make the + // molab `blob//` split ambiguous. + const shortSha = headSha.slice(0, 7); + const rows = notebooks.map((path) => { + const url = `https://molab.marimo.io/github/${headOwner}/${headRepo}/blob/${headSha}/${path}`; + return `| \`${path}\` | [Open in molab ▶](${url}) |`; + }).join('\n'); + + const body = [ + marker, + '### ▶️ Run the marimo notebook(s) in this PR', + '', + '[molab](https://molab.marimo.io) launches any public marimo notebook on ' + + 'GitHub in a hosted environment — no local setup required.', + '', + '| Notebook | molab |', + '| --- | --- |', + rows, + '', + `_Links pin to commit \`${shortSha}\` and refresh on every push._`, + ].join('\n'); + + // 6. Upsert the comment. + if (existing) { + await github.rest.issues.updateComment({ + owner: context.repo.owner, + repo: context.repo.repo, + comment_id: existing.id, + body, + }); + } else { + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: pr.number, + body, + }); + } + core.info(`Linked ${notebooks.length} marimo notebook(s).`); From 3d812beeb5838af6c27a89030e0364fe86c6ad45 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 13:08:41 -0400 Subject: [PATCH 04/22] use official molab shield badge in PR comment link Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/marimo_molab.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/marimo_molab.yml b/.github/workflows/marimo_molab.yml index e2184827..14e27236 100644 --- a/.github/workflows/marimo_molab.yml +++ b/.github/workflows/marimo_molab.yml @@ -99,7 +99,7 @@ jobs: const shortSha = headSha.slice(0, 7); const rows = notebooks.map((path) => { const url = `https://molab.marimo.io/github/${headOwner}/${headRepo}/blob/${headSha}/${path}`; - return `| \`${path}\` | [Open in molab ▶](${url}) |`; + return `| \`${path}\` | [![Open in molab](https://marimo.io/molab-shield.svg)](${url}) |`; }).join('\n'); const body = [ From 5d9d4b8aef384557118888f37e9869a8606b8eff Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 13:16:58 -0400 Subject: [PATCH 05/22] add W&B API key field to marimo notebook Adds a masked mo.ui.text(kind="password") field so users can authenticate by pasting a key instead of running `wandb login` in a shell. The resolved value is kept out of the run config (never logged) and feeds wandb.login(key=...) in the gated training cell; blank falls back to ambient auth (shell login, WANDB_API_KEY, or netrc). Downstream cells read the key via the returned wandb_api_key. Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/marimo/mnist-registry/README.md | 6 +++-- .../marimo/mnist-registry/mnist_registry.py | 24 +++++++++++++++++-- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/examples/marimo/mnist-registry/README.md b/examples/marimo/mnist-registry/README.md index ad281ee5..dbf5d0c3 100644 --- a/examples/marimo/mnist-registry/README.md +++ b/examples/marimo/mnist-registry/README.md @@ -12,8 +12,10 @@ can resolve them automatically. ## Prerequisites - Python 3.10 or newer. -- A W&B account. Run `wandb login` once in your shell before launching the - notebook — this notebook does not prompt for an API key interactively. +- A W&B account, authenticated one of two ways: run `wandb login` in your + shell before launching the notebook, or paste your key into the **W&B API + key** field in the form. Get your key from + [wandb.ai/authorize](https://wandb.ai/authorize). - A W&B **Registry** must exist in your org for the final linking step. The built-in Model Registry is provisioned automatically in newer orgs. If linking fails, the notebook surfaces a remediation message in the last diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index 6217d8aa..2dc1bedd 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -44,8 +44,10 @@ def _(): ## Prerequisites - - Run **`wandb login`** in your shell before starting marimo. - This notebook does not prompt for an API key interactively. + - Authenticate with W&B one of two ways: run **`wandb login`** in + your shell before starting marimo, or paste your key into the + **W&B API key** field in the form below. Get your key from + [wandb.ai/authorize](https://wandb.ai/authorize). - A W&B entity (your user or a team) the run will be written to. - A **W&B Registry** must exist in your org. The built-in Model registry is provisioned automatically in newer orgs. If linking @@ -139,6 +141,9 @@ def _(mo): project = mo.ui.text(value="marimo-mnist-registry", label="W&B project") entity = mo.ui.text(value="", label="W&B entity (blank uses your default)") run_name = mo.ui.text(value="", label="Run name (blank auto-generates)") + api_key = mo.ui.text( + value="", kind="password", label="W&B API key (blank uses your shell login)" + ) registry_name = mo.ui.text(value="model", label="W&B Registry name") collection_name = mo.ui.text(value="MNIST Classifiers", label="Registry collection") @@ -151,6 +156,7 @@ def _(mo): mo.hstack([lr, momentum]), seed, mo.md("### W&B run"), + api_key, mo.hstack([project, entity, run_name]), mo.md("### Registry"), mo.hstack([registry_name, collection_name, link_to_registry]), @@ -158,6 +164,7 @@ def _(mo): ) form return ( + api_key, batch_size, collection_name, entity, @@ -174,6 +181,7 @@ def _(mo): @app.cell def _( + api_key, batch_size, collection_name, entity, @@ -198,6 +206,10 @@ def _( wandb_project = project.value or None wandb_entity = entity.value or None wandb_run_name = run_name.value or None + # Resolve the key but keep it out of `config` so it is never logged to + # W&B. Blank falls back to your ambient login (shell `wandb login`, the + # WANDB_API_KEY env var, or netrc). + wandb_api_key = api_key.value or None registry_name_v = registry_name.value.strip() collection_name_v = collection_name.value.strip() link_to_registry_v = link_to_registry.value @@ -206,6 +218,7 @@ def _( config, link_to_registry_v, registry_name_v, + wandb_api_key, wandb_entity, wandb_project, wandb_run_name, @@ -304,6 +317,7 @@ def _( torch, train_button, wandb, + wandb_api_key, wandb_entity, wandb_project, wandb_run_name, @@ -318,6 +332,12 @@ def _( if wandb.run is not None: wandb.finish() + # Authenticate. A key pasted into the form takes precedence; otherwise + # fall back to your ambient login (shell `wandb login`, the WANDB_API_KEY + # env var, or netrc). The key is never written to the run config. + if wandb_api_key: + wandb.login(key=wandb_api_key) + torch.manual_seed(config["seed"]) run = wandb.init( From 360c5b93b168bc7d2df9d18f1e3506b0b23640ef Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 13:26:38 -0400 Subject: [PATCH 06/22] clarify Registry view-only permission error in marimo notebook The link step's "common causes" did not name the most common real failure: a view-only org seat (`view-only member cannot write to project`), where the run and artifact succeed but linking is blocked. Lead the remediation with that case and how to resolve it, and note the write-access requirement up front in the prerequisites. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../marimo/mnist-registry/mnist_registry.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index 2dc1bedd..fff3bdc9 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -49,10 +49,11 @@ def _(): **W&B API key** field in the form below. Get your key from [wandb.ai/authorize](https://wandb.ai/authorize). - A W&B entity (your user or a team) the run will be written to. - - A **W&B Registry** must exist in your org. The built-in Model - registry is provisioned automatically in newer orgs. If linking - fails, the Registry step shows remediation guidance inline - instead of crashing. + - A **W&B Registry** must exist in your org, and your account needs + **write access** to it. The built-in Model registry is provisioned + automatically in newer orgs. If linking fails (for example, from a + view-only seat), the Registry step shows remediation guidance + inline instead of crashing. - A GPU is optional. The defaults finish in about 2 minutes on CPU. """ ) @@ -572,10 +573,16 @@ def _( f"Target path: `{target_path}`\n\n" f"Error: `{exc}`\n\n" f"Common causes:\n\n" + f"- Your account lacks **write access** to the Registry. The " + f"error `view-only member cannot write to project` means you " + f"are signed in as a view-only member of the org that owns the " + f"Registry: the run and artifact succeed, but linking is " + f"blocked. Ask an org admin for a non-view-only role, or set " + f"the **W&B entity** field to an org or team where you have " + f"write access (and sign in with that account).\n" f"- The Registry `{registry_name_v}` does not exist in your " f"org. An org admin can create the Model registry from the " f"W&B Registry UI.\n" - f"- Your account lacks Registry write permission.\n" f"- Your org is on the legacy Model Registry. In that case " f"use the legacy pattern:\n\n" f" ```python\n" From da3b0ba6fb94b571f7d55222fb0242ac3ce44cc4 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 13:31:41 -0400 Subject: [PATCH 07/22] use branch ref instead of commit SHA for molab links GitHub resolves multi-segment (slashed) branch names in blob// and molab fetches from GitHub, so a branch-based link points at the latest revision and needs no per-commit comment maintenance. Content detection still pins to the head SHA. Skip the comment write when the body is unchanged so no-op pushes don't churn it. Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/marimo_molab.yml | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/.github/workflows/marimo_molab.yml b/.github/workflows/marimo_molab.yml index 14e27236..4e19ecbe 100644 --- a/.github/workflows/marimo_molab.yml +++ b/.github/workflows/marimo_molab.yml @@ -32,7 +32,8 @@ jobs: const pr = context.payload.pull_request; const headOwner = pr.head.repo.owner.login; const headRepo = pr.head.repo.name; - const headSha = pr.head.sha; + const headSha = pr.head.sha; // pin content detection to this PR revision + const headRef = pr.head.ref; // branch name for the (auto-tracking) links const marker = ''; // 1. List the files changed in this PR. @@ -93,12 +94,13 @@ jobs: return; } - // 5. Build the comment. Links pin to the head commit SHA so they - // always resolve — branch names containing slashes make the - // molab `blob//` split ambiguous. - const shortSha = headSha.slice(0, 7); + // 5. Build the comment. Links use the branch ref, not a commit + // SHA, so they always point at the latest revision without the + // comment needing an update on every push. GitHub resolves + // multi-segment (slashed) branch names in `blob//`, + // and molab fetches from GitHub, so slashed branches are fine. const rows = notebooks.map((path) => { - const url = `https://molab.marimo.io/github/${headOwner}/${headRepo}/blob/${headSha}/${path}`; + const url = `https://molab.marimo.io/github/${headOwner}/${headRepo}/blob/${headRef}/${path}`; return `| \`${path}\` | [![Open in molab](https://marimo.io/molab-shield.svg)](${url}) |`; }).join('\n'); @@ -113,11 +115,16 @@ jobs: '| --- | --- |', rows, '', - `_Links pin to commit \`${shortSha}\` and refresh on every push._`, + `_Links track the head of \`${headRef}\`._`, ].join('\n'); - // 6. Upsert the comment. + // 6. Upsert the comment (skip the write when nothing changed, so + // pushes that add no new notebook don't churn the comment). if (existing) { + if (existing.body === body) { + core.info('Comment already up to date.'); + return; + } await github.rest.issues.updateComment({ owner: context.repo.owner, repo: context.repo.repo, From 2d65a7c5f016271cf84da0bd437902b9ff3a83f5 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 13:37:08 -0400 Subject: [PATCH 08/22] add concrete Registry access-grant remediation to view-only error The view-only link failure now names the required role (Member) and the three ways an admin can grant it: the Registry Members UI, the Python SDK (wandb.Api().registry(...) then add_member()/update_member()), and SCIM (PATCH /scim/Users/{id} with registryRoles), with a link to the configure registry access docs. README prerequisite updated to match. Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/marimo/mnist-registry/README.md | 11 +++++++---- examples/marimo/mnist-registry/mnist_registry.py | 15 ++++++++++----- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/examples/marimo/mnist-registry/README.md b/examples/marimo/mnist-registry/README.md index dbf5d0c3..cb3b89db 100644 --- a/examples/marimo/mnist-registry/README.md +++ b/examples/marimo/mnist-registry/README.md @@ -16,10 +16,13 @@ can resolve them automatically. shell before launching the notebook, or paste your key into the **W&B API key** field in the form. Get your key from [wandb.ai/authorize](https://wandb.ai/authorize). -- A W&B **Registry** must exist in your org for the final linking step. The - built-in Model Registry is provisioned automatically in newer orgs. If - linking fails, the notebook surfaces a remediation message in the last - Registry cell instead of crashing. +- A W&B **Registry** must exist in your org, and your account needs at least + the **Member** role on it for the final linking step (linking an artifact is + a write action). The built-in Model registry is provisioned automatically in + newer orgs. If linking fails (for example, from a view-only seat), the + notebook surfaces a remediation message in the last Registry cell instead of + crashing. See + [configuring registry access](https://docs.wandb.ai/guides/registry/configure_registry/). - GPU is optional. Defaults are tuned to finish in roughly two minutes on CPU. ## Run diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index fff3bdc9..b051218a 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -575,11 +575,16 @@ def _( f"Common causes:\n\n" f"- Your account lacks **write access** to the Registry. The " f"error `view-only member cannot write to project` means you " - f"are signed in as a view-only member of the org that owns the " - f"Registry: the run and artifact succeed, but linking is " - f"blocked. Ask an org admin for a non-view-only role, or set " - f"the **W&B entity** field to an org or team where you have " - f"write access (and sign in with that account).\n" + f"are signed in as a view-only member: the run and artifact " + f"succeed, but linking needs at least the **Member** role on " + f"the Registry. An admin can grant it from the Registry's " + f"**Members** settings, with the Python SDK " + f"(`wandb.Api().registry(...)` then `add_member()` or " + f"`update_member()`), or via SCIM (`PATCH /scim/Users/{{id}}` " + f"with `registryRoles`) — see " + f"https://docs.wandb.ai/guides/registry/configure_registry/. " + f"Alternatively, set the **W&B entity** field to an org or " + f"team where you already have write access.\n" f"- The Registry `{registry_name_v}` does not exist in your " f"org. An org admin can create the Model registry from the " f"W&B Registry UI.\n" From 545798438ce698a5c7f46a5da08cbf14d6d9bf52 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 14:53:22 -0400 Subject: [PATCH 09/22] open molab links in /server (hosted runtime) mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The bare molab GitHub URL renders a static preview. Appending /server opens the notebook in a hosted runtime so reviewers can actually run it — needed here since the notebook depends on torch, which the in-browser /wasm mode cannot run. Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/marimo_molab.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/marimo_molab.yml b/.github/workflows/marimo_molab.yml index 4e19ecbe..aa2f5020 100644 --- a/.github/workflows/marimo_molab.yml +++ b/.github/workflows/marimo_molab.yml @@ -100,7 +100,9 @@ jobs: // multi-segment (slashed) branch names in `blob//`, // and molab fetches from GitHub, so slashed branches are fine. const rows = notebooks.map((path) => { - const url = `https://molab.marimo.io/github/${headOwner}/${headRepo}/blob/${headRef}/${path}`; + // The `/server` suffix opens the notebook in a hosted runtime; + // without it molab shows a static, non-runnable preview. + const url = `https://molab.marimo.io/github/${headOwner}/${headRepo}/blob/${headRef}/${path}/server`; return `| \`${path}\` | [![Open in molab](https://marimo.io/molab-shield.svg)](${url}) |`; }).join('\n'); From a5c3933b6eb002eb8842f3d2377da58db3133e26 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 14:59:36 -0400 Subject: [PATCH 10/22] export marimo notebooks to Markdown peers in CI Adds a push-triggered workflow that runs `marimo export md` on each notebook under examples/marimo/, writing a peer .md, and commits the result back to the branch so the rendered Markdown always tracks the notebook. marimo is pinned for byte-deterministic output (no version-only churn), and three guards prevent commit loops (GITHUB_TOKEN pushes don't re-trigger, *.py-only trigger vs *.md-only commits, and [skip ci]). Includes the initial export of mnist_registry.md. Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/marimo_export_md.yml | 76 +++ .../marimo/mnist-registry/mnist_registry.md | 523 ++++++++++++++++++ 2 files changed, 599 insertions(+) create mode 100644 .github/workflows/marimo_export_md.yml create mode 100644 examples/marimo/mnist-registry/mnist_registry.md diff --git a/.github/workflows/marimo_export_md.yml b/.github/workflows/marimo_export_md.yml new file mode 100644 index 00000000..130601ec --- /dev/null +++ b/.github/workflows/marimo_export_md.yml @@ -0,0 +1,76 @@ +name: marimo export markdown + +# On every push that changes a marimo notebook under examples/marimo/, export +# each notebook to a peer Markdown file (notebook.py -> notebook.md) and commit +# the result back to the branch, so the rendered Markdown always tracks the +# notebook. +# +# Loop safety (three independent guards): +# 1. Pushes made with GITHUB_TOKEN do not trigger new workflow runs — a +# GitHub Actions built-in, and the primary protection here. +# 2. The trigger watches only *.py; this job only ever commits *.md. +# 3. The commit message carries [skip ci]. +# +# marimo is pinned so exports are byte-deterministic (the front matter records +# the marimo version), which means an unchanged notebook never produces a +# spurious commit. Bump MARIMO_VERSION to refresh all exports on the next push. + +on: + push: + paths: + - 'examples/marimo/**/*.py' + +permissions: + contents: write + +concurrency: + group: marimo-export-md-${{ github.ref }} + cancel-in-progress: true + +env: + MARIMO_VERSION: "0.23.9" + +jobs: + export-md: + # Redundant with the GITHUB_TOKEN protection above, but keeps things safe + # if someone later swaps in a personal access token. + if: github.actor != 'github-actions[bot]' + runs-on: ubuntu-latest + steps: + - name: Checkout branch + uses: actions/checkout@v4 + with: + ref: ${{ github.ref_name }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install marimo + run: python -m pip install --quiet "marimo==${MARIMO_VERSION}" + + - name: Export marimo notebooks to Markdown + run: | + shopt -s globstar nullglob + for nb in examples/marimo/**/*.py; do + # Only real marimo notebooks construct marimo.App(...). + if grep -q 'marimo\.App(' "$nb"; then + echo "Exporting $nb -> ${nb%.py}.md" + marimo export md "$nb" -o "${nb%.py}.md" -f + fi + done + + - name: Commit and push if the Markdown changed + run: | + git config user.name 'github-actions[bot]' + git config user.email '41898282+github-actions[bot]@users.noreply.github.com' + # Only Markdown peers are generated, so staging the tree captures + # exactly the exported files (the notebooks themselves are untouched). + git add -A examples/marimo + if git diff --cached --quiet; then + echo "Markdown already up to date." + else + git commit -m "docs: export marimo notebook(s) to Markdown [skip ci]" + git push origin "HEAD:${{ github.ref_name }}" + fi diff --git a/examples/marimo/mnist-registry/mnist_registry.md b/examples/marimo/mnist-registry/mnist_registry.md new file mode 100644 index 00000000..e9989e88 --- /dev/null +++ b/examples/marimo/mnist-registry/mnist_registry.md @@ -0,0 +1,523 @@ +--- +title: MNIST -> W&B Registry +marimo-version: 0.23.9 +width: medium +header: |- + # /// script + # requires-python = ">=3.10" + # dependencies = [ + # "marimo>=0.9", + # "torch>=2.1", + # "torchvision>=0.16", + # "wandb>=0.18", + # "tqdm", + # ] + # /// + """Train an MNIST CNN with PyTorch, track the run with Weights & Biases, + and link the resulting model artifact to a W&B Registry collection. + + Run: + + uvx marimo edit mnist_registry.py --sandbox + + This notebook is interactive: hyperparameters live in a form, and training is + gated by a button so slider changes do not trigger runs. + """ +--- + +```python {.marimo hide_code="true"} +import marimo as mo + +mo.md( + """ + # MNIST -> W&B Run -> Registry + + ## What you will build + + - A **W&B run** with training and test metrics, gradient histograms, + and example test-set predictions logged as images. + - A **model Artifact** named `mnist-cnn-` of type `model`, + carrying metadata (test accuracy, parameter count, hyperparameters). + - A version of that Artifact **linked into a W&B Registry collection** + so it appears under registered models org-wide. + + ## Prerequisites + + - Authenticate with W&B one of two ways: run **`wandb login`** in + your shell before starting marimo, or paste your key into the + **W&B API key** field in the form below. Get your key from + [wandb.ai/authorize](https://wandb.ai/authorize). + - A W&B entity (your user or a team) the run will be written to. + - A **W&B Registry** must exist in your org, and your account needs + **write access** to it. The built-in Model registry is provisioned + automatically in newer orgs. If linking fails (for example, from a + view-only seat), the Registry step shows remediation guidance + inline instead of crashing. + - A GPU is optional. The defaults finish in about 2 minutes on CPU. + """ +) +``` + +```python {.marimo} +import os + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +import wandb +from tqdm.auto import tqdm +``` + +```python {.marimo} +if torch.cuda.is_available(): + device = torch.device("cuda") + device_note = "CUDA GPU detected. Training will be fast." + callout_kind = "success" +elif torch.backends.mps.is_available(): + device = torch.device("mps") + device_note = "Apple MPS detected. Training will run on the GPU." + callout_kind = "success" +else: + device = torch.device("cpu") + device_note = ( + "No GPU detected. Training will run on CPU. With the default " + "hyperparameters this takes about 2 minutes." + ) + callout_kind = "warn" + +mo.callout(mo.md(f"**Device:** `{device}` — {device_note}"), kind=callout_kind) +``` + +## Hyperparameters + +Configure the training run and the Registry target. The defaults reach +roughly 98% test accuracy in about two minutes on CPU. The **Registry** +section controls where the trained model is linked after training +finishes. + +```python {.marimo} +epochs = mo.ui.slider(start=1, stop=10, step=1, value=3, label="Epochs") +batch_size = mo.ui.dropdown( + options=["32", "64", "128", "256"], value="64", label="Batch size" +) +lr = mo.ui.slider( + start=0.001, stop=0.1, step=0.001, value=0.01, label="Learning rate", show_value=True +) +momentum = mo.ui.slider( + start=0.0, stop=0.99, step=0.01, value=0.5, label="SGD momentum", show_value=True +) +seed = mo.ui.number(start=0, stop=99999, value=42, label="Random seed") + +project = mo.ui.text(value="marimo-mnist-registry", label="W&B project") +entity = mo.ui.text(value="", label="W&B entity (blank uses your default)") +run_name = mo.ui.text(value="", label="Run name (blank auto-generates)") +api_key = mo.ui.text( + value="", kind="password", label="W&B API key (blank uses your shell login)" +) + +registry_name = mo.ui.text(value="model", label="W&B Registry name") +collection_name = mo.ui.text(value="MNIST Classifiers", label="Registry collection") +link_to_registry = mo.ui.checkbox(value=True, label="Link artifact to Registry") + +form = mo.vstack( + [ + mo.md("### Training"), + mo.hstack([epochs, batch_size]), + mo.hstack([lr, momentum]), + seed, + mo.md("### W&B run"), + api_key, + mo.hstack([project, entity, run_name]), + mo.md("### Registry"), + mo.hstack([registry_name, collection_name, link_to_registry]), + ] +) +form +``` + +```python {.marimo} +config = { + "epochs": epochs.value, + "batch_size": int(batch_size.value), + "lr": lr.value, + "momentum": momentum.value, + "seed": seed.value, + "architecture": "CNN", + "dataset": "MNIST", +} +wandb_project = project.value or None +wandb_entity = entity.value or None +wandb_run_name = run_name.value or None +# Resolve the key but keep it out of `config` so it is never logged to +# W&B. Blank falls back to your ambient login (shell `wandb login`, the +# WANDB_API_KEY env var, or netrc). +wandb_api_key = api_key.value or None +registry_name_v = registry_name.value.strip() +collection_name_v = collection_name.value.strip() +link_to_registry_v = link_to_registry.value +``` + +## Train + +Click **Train model** to begin. Changing a hyperparameter does not +start a run by itself — the button gates execution. Once a run +completes, clicking the button again starts a new run using the +current form values; the previous run finishes cleanly first. + +```python {.marimo} +train_button = mo.ui.run_button(label="Train model", kind="success") +train_button +``` + +```python {.marimo} +class Net(nn.Module): + """The same small CNN used in examples/pytorch/pytorch-cnn-mnist. + + Two convolutional layers (10 and 20 filters, 5x5 kernels) feed into two + fully connected layers (50 hidden units, 10 outputs). Roughly 21k + parameters. + """ + + def __init__(self): + super().__init__() + self.conv1 = nn.Conv2d(1, 10, kernel_size=5) + self.conv2 = nn.Conv2d(10, 20, kernel_size=5) + self.conv2_drop = nn.Dropout2d() + self.fc1 = nn.Linear(320, 50) + self.fc2 = nn.Linear(50, 10) + + def forward(self, x): + x = F.relu(F.max_pool2d(self.conv1(x), 2)) + x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) + x = x.view(-1, 320) + x = F.relu(self.fc1(x)) + x = F.dropout(x, training=self.training) + x = self.fc2(x) + return F.log_softmax(x, dim=1) +``` + +```python {.marimo} +transform = transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] +) +train_ds = datasets.MNIST("./data", train=True, download=True, transform=transform) +test_ds = datasets.MNIST("./data", train=False, download=True, transform=transform) + +# Only `batch_size` and `device` affect the loaders, so we depend on them +# directly rather than the full `config` dict; this avoids re-creating the +# loaders whenever an unrelated hyperparameter changes. +bs = int(batch_size.value) +loader_kwargs = ( + {"num_workers": 2, "pin_memory": True} if device.type == "cuda" else {} +) +train_loader = DataLoader( + train_ds, batch_size=bs, shuffle=True, **loader_kwargs +) +test_loader = DataLoader( + test_ds, batch_size=1000, shuffle=False, **loader_kwargs +) + +mo.md( + f"**Train:** {len(train_ds):,} examples · " + f"**Test:** {len(test_ds):,} examples · " + f"**Batch size:** {bs}" +) +``` + +```python {.marimo} +mo.stop(not train_button.value, mo.md("Click **Train model** to begin.")) + +# Finish any prior run still attached to this Python process. +# marimo keeps the kernel alive across re-clicks, so a second click — or +# a slider change after the first click — re-executes this cell. Without +# this guard, `wandb.init` warns about a run already being active. +# `wandb.finish` blocks until the prior run's tail logs are uploaded. +if wandb.run is not None: + wandb.finish() + +# Authenticate. A key pasted into the form takes precedence; otherwise +# fall back to your ambient login (shell `wandb login`, the WANDB_API_KEY +# env var, or netrc). The key is never written to the run config. +if wandb_api_key: + wandb.login(key=wandb_api_key) + +torch.manual_seed(config["seed"]) + +run = wandb.init( + project=wandb_project, + entity=wandb_entity, + name=wandb_run_name, + config=config, + job_type="train", +) + +# Use `epoch` as the x-axis for training and test metrics in the W&B UI. +wandb.define_metric("epoch") +wandb.define_metric("train/*", step_metric="epoch") +wandb.define_metric("test/*", step_metric="epoch") + +model = Net().to(device) +# `log="gradients"` is the standard choice for tutorial examples; +# `log="all"` also logs parameter histograms at extra cost. +wandb.watch(model, log="gradients", log_freq=100) +optimizer = optim.SGD( + model.parameters(), lr=config["lr"], momentum=config["momentum"] +) + +mo.md(f"Run started: [`{run.name}`]({run.url})") +``` + +```python {.marimo} +mo.stop(not train_button.value, mo.md("")) + +history = [] +best_acc = 0.0 +final_acc = 0.0 +final_loss = 0.0 + +for epoch in range(1, config["epochs"] + 1): + # ---- train ---- + model.train() + for batch_idx, (data, target) in enumerate( + tqdm(train_loader, desc=f"epoch {epoch}/{config['epochs']}") + ): + data, target = data.to(device), target.to(device) + optimizer.zero_grad() + output = model(data) + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % 50 == 0: + wandb.log({"train/loss": loss.item(), "epoch": epoch}) + + # ---- test ---- + model.eval() + test_loss = 0.0 + correct = 0 + example_images = [] + with torch.no_grad(): + for data, target in test_loader: + data, target = data.to(device), target.to(device) + output = model(data) + test_loss += F.nll_loss(output, target, reduction="sum").item() + pred = output.argmax(dim=1, keepdim=True) + correct += pred.eq(target.view_as(pred)).sum().item() + # Pull up to 16 example predictions from the first batch we see. + while len(example_images) < 16 and len(example_images) < data.size(0): + j = len(example_images) + example_images.append( + wandb.Image( + data[j], + caption=( + f"pred={pred[j].item()} " + f"true={target[j].item()}" + ), + ) + ) + + test_loss /= len(test_loader.dataset) + test_acc = correct / len(test_loader.dataset) + best_acc = max(best_acc, test_acc) + final_acc = test_acc + final_loss = test_loss + wandb.log( + { + "test/loss": test_loss, + "test/accuracy": test_acc, + "epoch": epoch, + "examples": example_images, + } + ) + history.append( + {"epoch": epoch, "test_loss": test_loss, "test_acc": test_acc} + ) +``` + +```python {.marimo} +mo.stop(not train_button.value, mo.md("")) +mo.vstack( + [ + mo.md("### Training summary"), + mo.ui.table(history, selection=None), + mo.md( + f"**Final test accuracy:** {final_acc:.2%} · " + f"**Final test loss:** {final_loss:.4f}" + ), + ] +) +``` + +```python {.marimo} +mo.stop(not train_button.value, mo.md("")) + +model_path = "mnist_cnn.pt" +torch.save(model.state_dict(), model_path) + +mo.md( + f"Saved `{model_path}` ({os.path.getsize(model_path) / 1024:.1f} KB)" +) +``` + +```python {.marimo} +mo.stop(not train_button.value, mo.md("")) + +num_params = sum(p.numel() for p in model.parameters()) + +artifact = wandb.Artifact( + name=f"mnist-cnn-{run.id}", + type="model", + description=( + "Small CNN trained on MNIST. Architecture: 2 conv layers " + "(10 and 20 filters, 5x5 kernels) + 2 FC layers (50, 10)." + ), + metadata={ + "framework": "pytorch", + "architecture": "CNN", + "num_parameters": num_params, + "dataset": "MNIST", + "train_size": len(train_ds), + "test_size": len(test_ds), + "test_accuracy": final_acc, + "best_test_accuracy": best_acc, + "hyperparameters": dict(config), + }, +) +artifact.add_file(model_path) + +# Log a single artifact per run (the final-epoch weights) and tag it +# `latest` unconditionally. Use the Registry UI or the API to promote a +# specific version with aliases like `best` or `production` after +# comparing runs. +aliases = ["latest"] + +logged = run.log_artifact(artifact, aliases=aliases) +# Block until the artifact has fully committed server-side. Without this, +# `link_artifact` below may race on the version reference. +logged.wait() + +mo.md( + f"Artifact logged: `{artifact.name}` with aliases `{aliases}`" +) +``` + +````python {.marimo} +mo.stop(not train_button.value, mo.md("")) +mo.stop( + not link_to_registry_v, + mo.md( + "_Registry linking is disabled (checkbox unchecked). " + "The artifact is logged to the run but not linked to a Registry " + "collection._" + ), +) + +target_path = f"wandb-registry-{registry_name_v}/{collection_name_v}" + +try: + run.link_artifact(artifact=logged, target_path=target_path) + link_result = mo.callout( + mo.md( + f"**Linked to Registry:** `{target_path}`\n\n" + f"Open the Registry at " + f"[https://wandb.ai/registry](https://wandb.ai/registry) to " + f"see the version." + ), + kind="success", + ) +except Exception as exc: # noqa: BLE001 - we want to surface any failure to the reader + link_result = mo.callout( + mo.md( + f"**Registry link failed.**\n\n" + f"Target path: `{target_path}`\n\n" + f"Error: `{exc}`\n\n" + f"Common causes:\n\n" + f"- Your account lacks **write access** to the Registry. The " + f"error `view-only member cannot write to project` means you " + f"are signed in as a view-only member: the run and artifact " + f"succeed, but linking needs at least the **Member** role on " + f"the Registry. An admin can grant it from the Registry's " + f"**Members** settings, with the Python SDK " + f"(`wandb.Api().registry(...)` then `add_member()` or " + f"`update_member()`), or via SCIM (`PATCH /scim/Users/{{id}}` " + f"with `registryRoles`) — see " + f"https://docs.wandb.ai/guides/registry/configure_registry/. " + f"Alternatively, set the **W&B entity** field to an org or " + f"team where you already have write access.\n" + f"- The Registry `{registry_name_v}` does not exist in your " + f"org. An org admin can create the Model registry from the " + f"W&B Registry UI.\n" + f"- Your org is on the legacy Model Registry. In that case " + f"use the legacy pattern:\n\n" + f" ```python\n" + f" run.link_artifact(\n" + f" logged,\n" + f" target_path='model-registry/{collection_name_v}',\n" + f" )\n" + f" ```" + ), + kind="danger", + ) +link_result +```` + +````python {.marimo hide_code="true"} +mo.stop(not train_button.value, mo.md("")) +mo.md( + f""" + ## Verify + + 1. Open the run page: [{run.name}]({run.url}). Confirm the + **Charts**, **System**, and **Examples** panels are populated. + 2. Click **Artifacts** in the run's left nav. Confirm the + `mnist-cnn-{run.id}` model artifact is listed with metadata + (test accuracy, hyperparameters, number of parameters). + 3. Go to [wandb.ai/registry](https://wandb.ai/registry), open the + **{registry_name_v.title()}** registry, then the + **{collection_name_v}** collection. Confirm the linked version + is present. + + ## Consume the registered model + + From any other script or notebook, fetch the latest registered version: + + ```python + import wandb + api = wandb.Api() + art = api.artifact( + "wandb-registry-{registry_name_v}/{collection_name_v}:latest" + ) + art.download() # writes mnist_cnn.pt under ./artifacts/.../ + ``` + + ## Next steps + + - **Promote a version.** From the Registry UI, add the `production` + alias to the version you want consumers to pick up. The same + collection path with `:production` will then resolve to it. + - **Compare runs.** Re-run with a deeper architecture or a different + learning rate. Group runs in the W&B UI to compare test accuracy + across configurations. + - **Automate on promotion.** Configure a W&B Automation on the + collection to trigger evaluation jobs or webhooks when a new + version is linked. + """ +) +```` + +## Finish + +This cell closes the W&B run so the run summary and the Registry +version finalize on the server. + +```python {.marimo} +mo.stop(not train_button.value, mo.md("")) +# Mirrors the `wandb.finish` guard at the top of the training cell: +# leaves the kernel in a clean state for the next Train click. +if wandb.run is not None: + wandb.finish() +mo.md("Run finished. Click **Train model** again to start a new run.") +``` \ No newline at end of file From 6690ce4501ce7f08ba951260c6d1bb3f6d7d0f15 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 15:17:18 -0400 Subject: [PATCH 11/22] consolidate marimo notebook into input -> button -> run cells MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Collapses 19 cells into 5, breaking only where marimo requires it or where the user actually provides input: - intro (markdown) - setup + form: imports, device detection, and all input widgets in one cell - train button: its own cell (a widget must be defined separately from the cell that reads its value) - run training: one gated cell that configs, authenticates, inits, builds and trains the model, logs, saves, and links the artifact — streaming each milestone with mo.output.append so it isn't a black box - verify + next steps (markdown), rendered once a run exists run_button.value resets to False after its cascade, so editing the form after a run re-runs the training cell but mo.stop halts it immediately — no silent retrain. Regenerates the Markdown peer. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../marimo/mnist-registry/mnist_registry.md | 441 +++++-------- .../marimo/mnist-registry/mnist_registry.py | 583 ++++++------------ 2 files changed, 367 insertions(+), 657 deletions(-) diff --git a/examples/marimo/mnist-registry/mnist_registry.md b/examples/marimo/mnist-registry/mnist_registry.md index e9989e88..df505c21 100644 --- a/examples/marimo/mnist-registry/mnist_registry.md +++ b/examples/marimo/mnist-registry/mnist_registry.md @@ -20,8 +20,9 @@ header: |- uvx marimo edit mnist_registry.py --sandbox - This notebook is interactive: hyperparameters live in a form, and training is - gated by a button so slider changes do not trigger runs. + The notebook has three interactive cells: fill in the form, click **Train + model**, then read the results. Everything between the inputs and the button + runs as a single step, so one click trains, logs, saves, and registers. """ --- @@ -48,17 +49,21 @@ mo.md( **W&B API key** field in the form below. Get your key from [wandb.ai/authorize](https://wandb.ai/authorize). - A W&B entity (your user or a team) the run will be written to. - - A **W&B Registry** must exist in your org, and your account needs - **write access** to it. The built-in Model registry is provisioned - automatically in newer orgs. If linking fails (for example, from a - view-only seat), the Registry step shows remediation guidance - inline instead of crashing. + - A **W&B Registry** must exist in your org, and your account needs at + least the **Member** role on it (linking an artifact is a write + action). The built-in Model registry is provisioned automatically in + newer orgs. If linking fails (for example, from a view-only seat), + the run still completes and the Registry step explains how to fix it. - A GPU is optional. The defaults finish in about 2 minutes on CPU. """ ) ``` ```python {.marimo} +# Imports, device detection, and the input form all live in one cell: this +# is "everything up to collecting your inputs". It defines the form widgets +# but never reads their `.value` — marimo only makes a widget reactive when +# a *different* cell consumes it, which the training cell below does. import os import torch @@ -70,36 +75,23 @@ from torchvision import datasets, transforms import wandb from tqdm.auto import tqdm -``` -```python {.marimo} if torch.cuda.is_available(): device = torch.device("cuda") device_note = "CUDA GPU detected. Training will be fast." - callout_kind = "success" + device_kind = "success" elif torch.backends.mps.is_available(): device = torch.device("mps") device_note = "Apple MPS detected. Training will run on the GPU." - callout_kind = "success" + device_kind = "success" else: device = torch.device("cpu") device_note = ( "No GPU detected. Training will run on CPU. With the default " "hyperparameters this takes about 2 minutes." ) - callout_kind = "warn" - -mo.callout(mo.md(f"**Device:** `{device}` — {device_note}"), kind=callout_kind) -``` - -## Hyperparameters + device_kind = "warn" -Configure the training run and the Registry target. The defaults reach -roughly 98% test accuracy in about two minutes on CPU. The **Registry** -section controls where the trained model is linked after training -finishes. - -```python {.marimo} epochs = mo.ui.slider(start=1, stop=10, step=1, value=3, label="Epochs") batch_size = mo.ui.dropdown( options=["32", "64", "128", "256"], value="64", label="Batch size" @@ -136,10 +128,49 @@ form = mo.vstack( mo.hstack([registry_name, collection_name, link_to_registry]), ] ) -form + +mo.vstack( + [ + mo.callout( + mo.md(f"**Device:** `{device}` — {device_note}"), kind=device_kind + ), + mo.md( + "## Configure\n\nSet the hyperparameters and W&B targets, then click " + "**Train model** below. Changing a value here never starts a run on " + "its own — only the button does." + ), + form, + ] +) ``` ```python {.marimo} +# The button must be its own cell: the training cell reads +# `train_button.value`, and a widget is only reactive when a *different* +# cell consumes it. run_button's value is True for exactly the cascade its +# click triggers, then resets to False — so editing the form afterwards +# re-runs the training cell but it stops immediately instead of retraining. +train_button = mo.ui.run_button(label="Train model", kind="success") +mo.vstack( + [ + mo.md( + "## Train\n\nOne click runs the whole pipeline below: start the " + "run, build and train the model, log metrics and example " + "predictions, save the weights as an Artifact, and link it to the " + "Registry. Click again to retrain with new settings — the previous " + "run is finished first." + ), + train_button, + ] +) +``` + +```python {.marimo} +# Everything the Train button triggers, in one cell — no reason to make you +# advance through a chain of output-less code blocks. Each milestone is +# streamed to the cell output with `mo.output.append` as it happens. +mo.stop(not train_button.value, mo.md("Click **Train model** above to begin.")) + config = { "epochs": epochs.value, "batch_size": int(batch_size.value), @@ -149,38 +180,36 @@ config = { "architecture": "CNN", "dataset": "MNIST", } -wandb_project = project.value or None -wandb_entity = entity.value or None -wandb_run_name = run_name.value or None -# Resolve the key but keep it out of `config` so it is never logged to -# W&B. Blank falls back to your ambient login (shell `wandb login`, the -# WANDB_API_KEY env var, or netrc). -wandb_api_key = api_key.value or None registry_name_v = registry_name.value.strip() collection_name_v = collection_name.value.strip() -link_to_registry_v = link_to_registry.value -``` -## Train +# Authenticate. Finish any prior run first (marimo keeps the kernel alive +# across re-clicks). A key pasted into the form wins; otherwise fall back to +# ambient login (shell `wandb login`, WANDB_API_KEY, or netrc). The key is +# never written to the run config. +if wandb.run is not None: + wandb.finish() +if api_key.value: + wandb.login(key=api_key.value) -Click **Train model** to begin. Changing a hyperparameter does not -start a run by itself — the button gates execution. Once a run -completes, clicking the button again starts a new run using the -current form values; the previous run finishes cleanly first. +torch.manual_seed(config["seed"]) -```python {.marimo} -train_button = mo.ui.run_button(label="Train model", kind="success") -train_button -``` +run = wandb.init( + project=project.value or None, + entity=entity.value or None, + name=run_name.value or None, + config=config, + job_type="train", +) +# Use `epoch` as the x-axis for train/test metrics in the W&B UI. +wandb.define_metric("epoch") +wandb.define_metric("train/*", step_metric="epoch") +wandb.define_metric("test/*", step_metric="epoch") +# Surface the run link right away so you can watch metrics stream live. +mo.output.append(mo.md(f"**Run started:** [`{run.name}`]({run.url})")) -```python {.marimo} class Net(nn.Module): - """The same small CNN used in examples/pytorch/pytorch-cnn-mnist. - - Two convolutional layers (10 and 20 filters, 5x5 kernels) feed into two - fully connected layers (50 hidden units, 10 outputs). Roughly 21k - parameters. - """ + """Small CNN: 2 conv layers (10, 20 filters, 5x5) + 2 FC (50, 10).""" def __init__(self): super().__init__() @@ -198,89 +227,31 @@ class Net(nn.Module): x = F.dropout(x, training=self.training) x = self.fc2(x) return F.log_softmax(x, dim=1) -``` -```python {.marimo} +model = Net().to(device) +# `log="gradients"` is the standard choice; `log="all"` also logs parameter +# histograms at extra cost. +wandb.watch(model, log="gradients", log_freq=100) +optimizer = optim.SGD( + model.parameters(), lr=config["lr"], momentum=config["momentum"] +) + transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] ) train_ds = datasets.MNIST("./data", train=True, download=True, transform=transform) test_ds = datasets.MNIST("./data", train=False, download=True, transform=transform) - -# Only `batch_size` and `device` affect the loaders, so we depend on them -# directly rather than the full `config` dict; this avoids re-creating the -# loaders whenever an unrelated hyperparameter changes. -bs = int(batch_size.value) loader_kwargs = ( {"num_workers": 2, "pin_memory": True} if device.type == "cuda" else {} ) train_loader = DataLoader( - train_ds, batch_size=bs, shuffle=True, **loader_kwargs -) -test_loader = DataLoader( - test_ds, batch_size=1000, shuffle=False, **loader_kwargs + train_ds, batch_size=config["batch_size"], shuffle=True, **loader_kwargs ) - -mo.md( - f"**Train:** {len(train_ds):,} examples · " - f"**Test:** {len(test_ds):,} examples · " - f"**Batch size:** {bs}" -) -``` - -```python {.marimo} -mo.stop(not train_button.value, mo.md("Click **Train model** to begin.")) - -# Finish any prior run still attached to this Python process. -# marimo keeps the kernel alive across re-clicks, so a second click — or -# a slider change after the first click — re-executes this cell. Without -# this guard, `wandb.init` warns about a run already being active. -# `wandb.finish` blocks until the prior run's tail logs are uploaded. -if wandb.run is not None: - wandb.finish() - -# Authenticate. A key pasted into the form takes precedence; otherwise -# fall back to your ambient login (shell `wandb login`, the WANDB_API_KEY -# env var, or netrc). The key is never written to the run config. -if wandb_api_key: - wandb.login(key=wandb_api_key) - -torch.manual_seed(config["seed"]) - -run = wandb.init( - project=wandb_project, - entity=wandb_entity, - name=wandb_run_name, - config=config, - job_type="train", -) - -# Use `epoch` as the x-axis for training and test metrics in the W&B UI. -wandb.define_metric("epoch") -wandb.define_metric("train/*", step_metric="epoch") -wandb.define_metric("test/*", step_metric="epoch") - -model = Net().to(device) -# `log="gradients"` is the standard choice for tutorial examples; -# `log="all"` also logs parameter histograms at extra cost. -wandb.watch(model, log="gradients", log_freq=100) -optimizer = optim.SGD( - model.parameters(), lr=config["lr"], momentum=config["momentum"] -) - -mo.md(f"Run started: [`{run.name}`]({run.url})") -``` - -```python {.marimo} -mo.stop(not train_button.value, mo.md("")) +test_loader = DataLoader(test_ds, batch_size=1000, shuffle=False, **loader_kwargs) history = [] best_acc = 0.0 -final_acc = 0.0 -final_loss = 0.0 - for epoch in range(1, config["epochs"] + 1): - # ---- train ---- model.train() for batch_idx, (data, target) in enumerate( tqdm(train_loader, desc=f"epoch {epoch}/{config['epochs']}") @@ -294,7 +265,6 @@ for epoch in range(1, config["epochs"] + 1): if batch_idx % 50 == 0: wandb.log({"train/loss": loss.item(), "epoch": epoch}) - # ---- test ---- model.eval() test_loss = 0.0 correct = 0 @@ -306,24 +276,19 @@ for epoch in range(1, config["epochs"] + 1): test_loss += F.nll_loss(output, target, reduction="sum").item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() - # Pull up to 16 example predictions from the first batch we see. + # Pull up to 16 example predictions from the first batch. while len(example_images) < 16 and len(example_images) < data.size(0): j = len(example_images) example_images.append( wandb.Image( data[j], - caption=( - f"pred={pred[j].item()} " - f"true={target[j].item()}" - ), + caption=f"pred={pred[j].item()} true={target[j].item()}", ) ) test_loss /= len(test_loader.dataset) test_acc = correct / len(test_loader.dataset) best_acc = max(best_acc, test_acc) - final_acc = test_acc - final_loss = test_loss wandb.log( { "test/loss": test_loss, @@ -333,40 +298,23 @@ for epoch in range(1, config["epochs"] + 1): } ) history.append( - {"epoch": epoch, "test_loss": test_loss, "test_acc": test_acc} + {"epoch": epoch, "test_loss": round(test_loss, 4), "test_acc": round(test_acc, 4)} ) -``` -```python {.marimo} -mo.stop(not train_button.value, mo.md("")) -mo.vstack( - [ - mo.md("### Training summary"), - mo.ui.table(history, selection=None), - mo.md( - f"**Final test accuracy:** {final_acc:.2%} · " - f"**Final test loss:** {final_loss:.4f}" - ), - ] +final_acc = history[-1]["test_acc"] +mo.output.append( + mo.vstack( + [ + mo.md("### Training summary"), + mo.ui.table(history, selection=None), + mo.md(f"**Final test accuracy:** {final_acc:.2%}"), + ] + ) ) -``` - -```python {.marimo} -mo.stop(not train_button.value, mo.md("")) +# Save the weights and log them as a model Artifact tagged `latest`. model_path = "mnist_cnn.pt" torch.save(model.state_dict(), model_path) - -mo.md( - f"Saved `{model_path}` ({os.path.getsize(model_path) / 1024:.1f} KB)" -) -``` - -```python {.marimo} -mo.stop(not train_button.value, mo.md("")) - -num_params = sum(p.numel() for p in model.parameters()) - artifact = wandb.Artifact( name=f"mnist-cnn-{run.id}", type="model", @@ -377,7 +325,7 @@ artifact = wandb.Artifact( metadata={ "framework": "pytorch", "architecture": "CNN", - "num_parameters": num_params, + "num_parameters": sum(p.numel() for p in model.parameters()), "dataset": "MNIST", "train_size": len(train_ds), "test_size": len(test_ds), @@ -387,137 +335,90 @@ artifact = wandb.Artifact( }, ) artifact.add_file(model_path) - -# Log a single artifact per run (the final-epoch weights) and tag it -# `latest` unconditionally. Use the Registry UI or the API to promote a -# specific version with aliases like `best` or `production` after -# comparing runs. -aliases = ["latest"] - -logged = run.log_artifact(artifact, aliases=aliases) -# Block until the artifact has fully committed server-side. Without this, -# `link_artifact` below may race on the version reference. +logged = run.log_artifact(artifact, aliases=["latest"]) +# Block until the artifact has committed before linking, to avoid a race. logged.wait() - -mo.md( - f"Artifact logged: `{artifact.name}` with aliases `{aliases}`" -) -``` - -````python {.marimo} -mo.stop(not train_button.value, mo.md("")) -mo.stop( - not link_to_registry_v, - mo.md( - "_Registry linking is disabled (checkbox unchecked). " - "The artifact is logged to the run but not linked to a Registry " - "collection._" - ), -) - -target_path = f"wandb-registry-{registry_name_v}/{collection_name_v}" - -try: - run.link_artifact(artifact=logged, target_path=target_path) - link_result = mo.callout( - mo.md( - f"**Linked to Registry:** `{target_path}`\n\n" - f"Open the Registry at " - f"[https://wandb.ai/registry](https://wandb.ai/registry) to " - f"see the version." - ), - kind="success", - ) -except Exception as exc: # noqa: BLE001 - we want to surface any failure to the reader - link_result = mo.callout( +mo.output.append(mo.md(f"**Artifact logged:** `{artifact.name}` (alias `latest`)")) + +# Link to the Registry, surfacing a remediation note instead of crashing. +if link_to_registry.value: + target_path = f"wandb-registry-{registry_name_v}/{collection_name_v}" + try: + run.link_artifact(artifact=logged, target_path=target_path) + mo.output.append( + mo.callout( + mo.md( + f"**Linked to Registry:** `{target_path}` — see " + f"[wandb.ai/registry](https://wandb.ai/registry)." + ), + kind="success", + ) + ) + except Exception as exc: # noqa: BLE001 - surface any failure to the reader + mo.output.append( + mo.callout( + mo.md( + f"**Registry link failed.** Target `{target_path}` — `{exc}`\n\n" + f"- Linking needs at least the **Member** role on the " + f"Registry. `view-only member cannot write to project` means " + f"your seat is view-only: the run and artifact succeed, but " + f"linking is blocked. An admin can grant access from the " + f"Registry **Members** settings, the Python SDK " + f"(`wandb.Api().registry(...)` then `add_member()` / " + f"`update_member()`), or SCIM (`PATCH /scim/Users/{{id}}` with " + f"`registryRoles`) — see " + f"https://docs.wandb.ai/guides/registry/configure_registry/. " + f"Or set **W&B entity** to an org/team where you have write " + f"access.\n" + f"- The Registry `{registry_name_v}` may not exist; an admin " + f"can create it from the W&B Registry UI.\n" + f"- On the legacy Model Registry, link with " + f"`target_path='model-registry/{collection_name_v}'` instead." + ), + kind="danger", + ) + ) +else: + mo.output.append( mo.md( - f"**Registry link failed.**\n\n" - f"Target path: `{target_path}`\n\n" - f"Error: `{exc}`\n\n" - f"Common causes:\n\n" - f"- Your account lacks **write access** to the Registry. The " - f"error `view-only member cannot write to project` means you " - f"are signed in as a view-only member: the run and artifact " - f"succeed, but linking needs at least the **Member** role on " - f"the Registry. An admin can grant it from the Registry's " - f"**Members** settings, with the Python SDK " - f"(`wandb.Api().registry(...)` then `add_member()` or " - f"`update_member()`), or via SCIM (`PATCH /scim/Users/{{id}}` " - f"with `registryRoles`) — see " - f"https://docs.wandb.ai/guides/registry/configure_registry/. " - f"Alternatively, set the **W&B entity** field to an org or " - f"team where you already have write access.\n" - f"- The Registry `{registry_name_v}` does not exist in your " - f"org. An org admin can create the Model registry from the " - f"W&B Registry UI.\n" - f"- Your org is on the legacy Model Registry. In that case " - f"use the legacy pattern:\n\n" - f" ```python\n" - f" run.link_artifact(\n" - f" logged,\n" - f" target_path='model-registry/{collection_name_v}',\n" - f" )\n" - f" ```" - ), - kind="danger", + "_Registry linking is disabled — the artifact is logged to the run " + "but not linked to a collection._" + ) ) -link_result -```` + +# Close the run so its summary and any Registry link finalize server-side. +wandb.finish() +``` ````python {.marimo hide_code="true"} -mo.stop(not train_button.value, mo.md("")) +# Renders only after a run exists (it consumes `run` from the training +# cell), so it appears once training finishes. mo.md( f""" - ## Verify - - 1. Open the run page: [{run.name}]({run.url}). Confirm the - **Charts**, **System**, and **Examples** panels are populated. - 2. Click **Artifacts** in the run's left nav. Confirm the - `mnist-cnn-{run.id}` model artifact is listed with metadata - (test accuracy, hyperparameters, number of parameters). - 3. Go to [wandb.ai/registry](https://wandb.ai/registry), open the - **{registry_name_v.title()}** registry, then the - **{collection_name_v}** collection. Confirm the linked version - is present. + ## Verify and next steps - ## Consume the registered model + 1. Open the run: [{run.name}]({run.url}) — check the **Charts**, + **System**, and **Examples** panels. + 2. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed + with its metadata (test accuracy, parameter count, hyperparameters). + 3. At [wandb.ai/registry](https://wandb.ai/registry), open the + **{registry_name_v.title()}** registry, then the **{collection_name_v}** + collection, and confirm the linked version. - From any other script or notebook, fetch the latest registered version: + **Consume the registered model** from any script or notebook: ```python import wandb - api = wandb.Api() - art = api.artifact( + art = wandb.Api().artifact( "wandb-registry-{registry_name_v}/{collection_name_v}:latest" ) - art.download() # writes mnist_cnn.pt under ./artifacts/.../ + art.download() # writes mnist_cnn.pt under ./artifacts/ ``` - ## Next steps - - - **Promote a version.** From the Registry UI, add the `production` - alias to the version you want consumers to pick up. The same - collection path with `:production` will then resolve to it. - - **Compare runs.** Re-run with a deeper architecture or a different - learning rate. Group runs in the W&B UI to compare test accuracy - across configurations. - - **Automate on promotion.** Configure a W&B Automation on the - collection to trigger evaluation jobs or webhooks when a new - version is linked. + **Next steps:** promote a version by adding the `production` alias from + the Registry UI; re-run with a deeper architecture or a different + learning rate and compare runs in the W&B UI; or add a W&B Automation to + trigger evaluation when a new version is linked. """ ) -```` - -## Finish - -This cell closes the W&B run so the run summary and the Registry -version finalize on the server. - -```python {.marimo} -mo.stop(not train_button.value, mo.md("")) -# Mirrors the `wandb.finish` guard at the top of the training cell: -# leaves the kernel in a clean state for the next Train click. -if wandb.run is not None: - wandb.finish() -mo.md("Run finished. Click **Train model** again to start a new run.") -``` \ No newline at end of file +```` \ No newline at end of file diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index b051218a..250eac3b 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -15,13 +15,14 @@ uvx marimo edit mnist_registry.py --sandbox -This notebook is interactive: hyperparameters live in a form, and training is -gated by a button so slider changes do not trigger runs. +The notebook has three interactive cells: fill in the form, click **Train +model**, then read the results. Everything between the inputs and the button +runs as a single step, so one click trains, logs, saves, and registers. """ import marimo -__generated_with = "0.9.0" +__generated_with = "0.23.9" app = marimo.App(width="medium", app_title="MNIST -> W&B Registry") @@ -49,11 +50,11 @@ def _(): **W&B API key** field in the form below. Get your key from [wandb.ai/authorize](https://wandb.ai/authorize). - A W&B entity (your user or a team) the run will be written to. - - A **W&B Registry** must exist in your org, and your account needs - **write access** to it. The built-in Model registry is provisioned - automatically in newer orgs. If linking fails (for example, from a - view-only seat), the Registry step shows remediation guidance - inline instead of crashing. + - A **W&B Registry** must exist in your org, and your account needs at + least the **Member** role on it (linking an artifact is a write + action). The built-in Model registry is provisioned automatically in + newer orgs. If linking fails (for example, from a view-only seat), + the run still completes and the Registry step explains how to fix it. - A GPU is optional. The defaults finish in about 2 minutes on CPU. """ ) @@ -61,7 +62,11 @@ def _(): @app.cell -def _(): +def _(mo): + # Imports, device detection, and the input form all live in one cell: this + # is "everything up to collecting your inputs". It defines the form widgets + # but never reads their `.value` — marimo only makes a widget reactive when + # a *different* cell consumes it, which the training cell below does. import os import torch @@ -74,59 +79,22 @@ def _(): import wandb from tqdm.auto import tqdm - return ( - DataLoader, - F, - datasets, - nn, - optim, - os, - torch, - tqdm, - transforms, - wandb, - ) - - -@app.cell -def _(mo, torch): if torch.cuda.is_available(): device = torch.device("cuda") device_note = "CUDA GPU detected. Training will be fast." - callout_kind = "success" + device_kind = "success" elif torch.backends.mps.is_available(): device = torch.device("mps") device_note = "Apple MPS detected. Training will run on the GPU." - callout_kind = "success" + device_kind = "success" else: device = torch.device("cpu") device_note = ( "No GPU detected. Training will run on CPU. With the default " "hyperparameters this takes about 2 minutes." ) - callout_kind = "warn" - - mo.callout(mo.md(f"**Device:** `{device}` — {device_note}"), kind=callout_kind) - return (device,) - - -@app.cell(hide_code=True) -def _(mo): - mo.md( - """ - ## Hyperparameters - - Configure the training run and the Registry target. The defaults reach - roughly 98% test accuracy in about two minutes on CPU. The **Registry** - section controls where the trained model is linked after training - finishes. - """ - ) - return + device_kind = "warn" - -@app.cell -def _(mo): epochs = mo.ui.slider(start=1, stop=10, step=1, value=3, label="Epochs") batch_size = mo.ui.dropdown( options=["32", "64", "128", "256"], value="64", label="Batch size" @@ -163,38 +131,102 @@ def _(mo): mo.hstack([registry_name, collection_name, link_to_registry]), ] ) - form + + mo.vstack( + [ + mo.callout( + mo.md(f"**Device:** `{device}` — {device_note}"), kind=device_kind + ), + mo.md( + "## Configure\n\nSet the hyperparameters and W&B targets, then click " + "**Train model** below. Changing a value here never starts a run on " + "its own — only the button does." + ), + form, + ] + ) return ( + DataLoader, + F, api_key, batch_size, collection_name, + datasets, + device, entity, epochs, link_to_registry, lr, momentum, + nn, + optim, + os, project, registry_name, run_name, seed, + torch, + tqdm, + transforms, + wandb, ) +@app.cell +def _(mo): + # The button must be its own cell: the training cell reads + # `train_button.value`, and a widget is only reactive when a *different* + # cell consumes it. run_button's value is True for exactly the cascade its + # click triggers, then resets to False — so editing the form afterwards + # re-runs the training cell but it stops immediately instead of retraining. + train_button = mo.ui.run_button(label="Train model", kind="success") + mo.vstack( + [ + mo.md( + "## Train\n\nOne click runs the whole pipeline below: start the " + "run, build and train the model, log metrics and example " + "predictions, save the weights as an Artifact, and link it to the " + "Registry. Click again to retrain with new settings — the previous " + "run is finished first." + ), + train_button, + ] + ) + return (train_button,) + + @app.cell def _( + DataLoader, + F, api_key, batch_size, collection_name, + datasets, + device, entity, epochs, link_to_registry, lr, momentum, + mo, + nn, + optim, project, registry_name, run_name, seed, + torch, + tqdm, + train_button, + transforms, + wandb, ): + # Everything the Train button triggers, in one cell — no reason to make you + # advance through a chain of output-less code blocks. Each milestone is + # streamed to the cell output with `mo.output.append` as it happens. + mo.stop(not train_button.value, mo.md("Click **Train model** above to begin.")) + config = { "epochs": epochs.value, "batch_size": int(batch_size.value), @@ -204,59 +236,36 @@ def _( "architecture": "CNN", "dataset": "MNIST", } - wandb_project = project.value or None - wandb_entity = entity.value or None - wandb_run_name = run_name.value or None - # Resolve the key but keep it out of `config` so it is never logged to - # W&B. Blank falls back to your ambient login (shell `wandb login`, the - # WANDB_API_KEY env var, or netrc). - wandb_api_key = api_key.value or None registry_name_v = registry_name.value.strip() collection_name_v = collection_name.value.strip() - link_to_registry_v = link_to_registry.value - return ( - collection_name_v, - config, - link_to_registry_v, - registry_name_v, - wandb_api_key, - wandb_entity, - wandb_project, - wandb_run_name, - ) + # Authenticate. Finish any prior run first (marimo keeps the kernel alive + # across re-clicks). A key pasted into the form wins; otherwise fall back to + # ambient login (shell `wandb login`, WANDB_API_KEY, or netrc). The key is + # never written to the run config. + if wandb.run is not None: + wandb.finish() + if api_key.value: + wandb.login(key=api_key.value) -@app.cell(hide_code=True) -def _(mo): - mo.md( - """ - ## Train + torch.manual_seed(config["seed"]) - Click **Train model** to begin. Changing a hyperparameter does not - start a run by itself — the button gates execution. Once a run - completes, clicking the button again starts a new run using the - current form values; the previous run finishes cleanly first. - """ + run = wandb.init( + project=project.value or None, + entity=entity.value or None, + name=run_name.value or None, + config=config, + job_type="train", ) - return - - -@app.cell -def _(mo): - train_button = mo.ui.run_button(label="Train model", kind="success") - train_button - return (train_button,) - + # Use `epoch` as the x-axis for train/test metrics in the W&B UI. + wandb.define_metric("epoch") + wandb.define_metric("train/*", step_metric="epoch") + wandb.define_metric("test/*", step_metric="epoch") + # Surface the run link right away so you can watch metrics stream live. + mo.output.append(mo.md(f"**Run started:** [`{run.name}`]({run.url})")) -@app.cell -def _(F, nn): class Net(nn.Module): - """The same small CNN used in examples/pytorch/pytorch-cnn-mnist. - - Two convolutional layers (10 and 20 filters, 5x5 kernels) feed into two - fully connected layers (50 hidden units, 10 outputs). Roughly 21k - parameters. - """ + """Small CNN: 2 conv layers (10, 20 filters, 5x5) + 2 FC (50, 10).""" def __init__(self): super().__init__() @@ -275,121 +284,30 @@ def forward(self, x): x = self.fc2(x) return F.log_softmax(x, dim=1) - return (Net,) - + model = Net().to(device) + # `log="gradients"` is the standard choice; `log="all"` also logs parameter + # histograms at extra cost. + wandb.watch(model, log="gradients", log_freq=100) + optimizer = optim.SGD( + model.parameters(), lr=config["lr"], momentum=config["momentum"] + ) -@app.cell -def _(DataLoader, batch_size, datasets, device, mo, transforms): transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] ) train_ds = datasets.MNIST("./data", train=True, download=True, transform=transform) test_ds = datasets.MNIST("./data", train=False, download=True, transform=transform) - - # Only `batch_size` and `device` affect the loaders, so we depend on them - # directly rather than the full `config` dict; this avoids re-creating the - # loaders whenever an unrelated hyperparameter changes. - bs = int(batch_size.value) loader_kwargs = ( {"num_workers": 2, "pin_memory": True} if device.type == "cuda" else {} ) train_loader = DataLoader( - train_ds, batch_size=bs, shuffle=True, **loader_kwargs + train_ds, batch_size=config["batch_size"], shuffle=True, **loader_kwargs ) - test_loader = DataLoader( - test_ds, batch_size=1000, shuffle=False, **loader_kwargs - ) - - mo.md( - f"**Train:** {len(train_ds):,} examples · " - f"**Test:** {len(test_ds):,} examples · " - f"**Batch size:** {bs}" - ) - return test_ds, test_loader, train_ds, train_loader - - -@app.cell -def _( - Net, - config, - device, - mo, - optim, - torch, - train_button, - wandb, - wandb_api_key, - wandb_entity, - wandb_project, - wandb_run_name, -): - mo.stop(not train_button.value, mo.md("Click **Train model** to begin.")) - - # Finish any prior run still attached to this Python process. - # marimo keeps the kernel alive across re-clicks, so a second click — or - # a slider change after the first click — re-executes this cell. Without - # this guard, `wandb.init` warns about a run already being active. - # `wandb.finish` blocks until the prior run's tail logs are uploaded. - if wandb.run is not None: - wandb.finish() - - # Authenticate. A key pasted into the form takes precedence; otherwise - # fall back to your ambient login (shell `wandb login`, the WANDB_API_KEY - # env var, or netrc). The key is never written to the run config. - if wandb_api_key: - wandb.login(key=wandb_api_key) - - torch.manual_seed(config["seed"]) - - run = wandb.init( - project=wandb_project, - entity=wandb_entity, - name=wandb_run_name, - config=config, - job_type="train", - ) - - # Use `epoch` as the x-axis for training and test metrics in the W&B UI. - wandb.define_metric("epoch") - wandb.define_metric("train/*", step_metric="epoch") - wandb.define_metric("test/*", step_metric="epoch") - - model = Net().to(device) - # `log="gradients"` is the standard choice for tutorial examples; - # `log="all"` also logs parameter histograms at extra cost. - wandb.watch(model, log="gradients", log_freq=100) - optimizer = optim.SGD( - model.parameters(), lr=config["lr"], momentum=config["momentum"] - ) - - mo.md(f"Run started: [`{run.name}`]({run.url})") - return model, optimizer, run - - -@app.cell -def _( - F, - config, - device, - mo, - model, - optimizer, - test_loader, - torch, - tqdm, - train_button, - train_loader, - wandb, -): - mo.stop(not train_button.value, mo.md("")) + test_loader = DataLoader(test_ds, batch_size=1000, shuffle=False, **loader_kwargs) history = [] best_acc = 0.0 - final_acc = 0.0 - final_loss = 0.0 - for epoch in range(1, config["epochs"] + 1): - # ---- train ---- model.train() for batch_idx, (data, target) in enumerate( tqdm(train_loader, desc=f"epoch {epoch}/{config['epochs']}") @@ -403,7 +321,6 @@ def _( if batch_idx % 50 == 0: wandb.log({"train/loss": loss.item(), "epoch": epoch}) - # ---- test ---- model.eval() test_loss = 0.0 correct = 0 @@ -415,24 +332,19 @@ def _( test_loss += F.nll_loss(output, target, reduction="sum").item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() - # Pull up to 16 example predictions from the first batch we see. + # Pull up to 16 example predictions from the first batch. while len(example_images) < 16 and len(example_images) < data.size(0): j = len(example_images) example_images.append( wandb.Image( data[j], - caption=( - f"pred={pred[j].item()} " - f"true={target[j].item()}" - ), + caption=f"pred={pred[j].item()} true={target[j].item()}", ) ) test_loss /= len(test_loader.dataset) test_acc = correct / len(test_loader.dataset) best_acc = max(best_acc, test_acc) - final_acc = test_acc - final_loss = test_loss wandb.log( { "test/loss": test_loss, @@ -442,59 +354,23 @@ def _( } ) history.append( - {"epoch": epoch, "test_loss": test_loss, "test_acc": test_acc} + {"epoch": epoch, "test_loss": round(test_loss, 4), "test_acc": round(test_acc, 4)} ) - return best_acc, final_acc, final_loss, history - - -@app.cell -def _(final_acc, final_loss, history, mo, train_button): - mo.stop(not train_button.value, mo.md("")) - mo.vstack( - [ - mo.md("### Training summary"), - mo.ui.table(history, selection=None), - mo.md( - f"**Final test accuracy:** {final_acc:.2%} · " - f"**Final test loss:** {final_loss:.4f}" - ), - ] + final_acc = history[-1]["test_acc"] + mo.output.append( + mo.vstack( + [ + mo.md("### Training summary"), + mo.ui.table(history, selection=None), + mo.md(f"**Final test accuracy:** {final_acc:.2%}"), + ] + ) ) - return - - -@app.cell -def _(mo, model, os, torch, train_button): - mo.stop(not train_button.value, mo.md("")) + # Save the weights and log them as a model Artifact tagged `latest`. model_path = "mnist_cnn.pt" torch.save(model.state_dict(), model_path) - - mo.md( - f"Saved `{model_path}` ({os.path.getsize(model_path) / 1024:.1f} KB)" - ) - return (model_path,) - - -@app.cell -def _( - best_acc, - config, - final_acc, - mo, - model, - model_path, - run, - test_ds, - train_button, - train_ds, - wandb, -): - mo.stop(not train_button.value, mo.md("")) - - num_params = sum(p.numel() for p in model.parameters()) - artifact = wandb.Artifact( name=f"mnist-cnn-{run.id}", type="model", @@ -505,7 +381,7 @@ def _( metadata={ "framework": "pytorch", "architecture": "CNN", - "num_parameters": num_params, + "num_parameters": sum(p.numel() for p in model.parameters()), "dataset": "MNIST", "train_size": len(train_ds), "test_size": len(test_ds), @@ -515,163 +391,96 @@ def _( }, ) artifact.add_file(model_path) - - # Log a single artifact per run (the final-epoch weights) and tag it - # `latest` unconditionally. Use the Registry UI or the API to promote a - # specific version with aliases like `best` or `production` after - # comparing runs. - aliases = ["latest"] - - logged = run.log_artifact(artifact, aliases=aliases) - # Block until the artifact has fully committed server-side. Without this, - # `link_artifact` below may race on the version reference. + logged = run.log_artifact(artifact, aliases=["latest"]) + # Block until the artifact has committed before linking, to avoid a race. logged.wait() - - mo.md( - f"Artifact logged: `{artifact.name}` with aliases `{aliases}`" - ) - return (logged,) - - -@app.cell -def _( - collection_name_v, - link_to_registry_v, - logged, - mo, - registry_name_v, - run, - train_button, -): - mo.stop(not train_button.value, mo.md("")) - mo.stop( - not link_to_registry_v, - mo.md( - "_Registry linking is disabled (checkbox unchecked). " - "The artifact is logged to the run but not linked to a Registry " - "collection._" - ), - ) - - target_path = f"wandb-registry-{registry_name_v}/{collection_name_v}" - - try: - run.link_artifact(artifact=logged, target_path=target_path) - link_result = mo.callout( - mo.md( - f"**Linked to Registry:** `{target_path}`\n\n" - f"Open the Registry at " - f"[https://wandb.ai/registry](https://wandb.ai/registry) to " - f"see the version." - ), - kind="success", - ) - except Exception as exc: # noqa: BLE001 - we want to surface any failure to the reader - link_result = mo.callout( + mo.output.append(mo.md(f"**Artifact logged:** `{artifact.name}` (alias `latest`)")) + + # Link to the Registry, surfacing a remediation note instead of crashing. + if link_to_registry.value: + target_path = f"wandb-registry-{registry_name_v}/{collection_name_v}" + try: + run.link_artifact(artifact=logged, target_path=target_path) + mo.output.append( + mo.callout( + mo.md( + f"**Linked to Registry:** `{target_path}` — see " + f"[wandb.ai/registry](https://wandb.ai/registry)." + ), + kind="success", + ) + ) + except Exception as exc: # noqa: BLE001 - surface any failure to the reader + mo.output.append( + mo.callout( + mo.md( + f"**Registry link failed.** Target `{target_path}` — `{exc}`\n\n" + f"- Linking needs at least the **Member** role on the " + f"Registry. `view-only member cannot write to project` means " + f"your seat is view-only: the run and artifact succeed, but " + f"linking is blocked. An admin can grant access from the " + f"Registry **Members** settings, the Python SDK " + f"(`wandb.Api().registry(...)` then `add_member()` / " + f"`update_member()`), or SCIM (`PATCH /scim/Users/{{id}}` with " + f"`registryRoles`) — see " + f"https://docs.wandb.ai/guides/registry/configure_registry/. " + f"Or set **W&B entity** to an org/team where you have write " + f"access.\n" + f"- The Registry `{registry_name_v}` may not exist; an admin " + f"can create it from the W&B Registry UI.\n" + f"- On the legacy Model Registry, link with " + f"`target_path='model-registry/{collection_name_v}'` instead." + ), + kind="danger", + ) + ) + else: + mo.output.append( mo.md( - f"**Registry link failed.**\n\n" - f"Target path: `{target_path}`\n\n" - f"Error: `{exc}`\n\n" - f"Common causes:\n\n" - f"- Your account lacks **write access** to the Registry. The " - f"error `view-only member cannot write to project` means you " - f"are signed in as a view-only member: the run and artifact " - f"succeed, but linking needs at least the **Member** role on " - f"the Registry. An admin can grant it from the Registry's " - f"**Members** settings, with the Python SDK " - f"(`wandb.Api().registry(...)` then `add_member()` or " - f"`update_member()`), or via SCIM (`PATCH /scim/Users/{{id}}` " - f"with `registryRoles`) — see " - f"https://docs.wandb.ai/guides/registry/configure_registry/. " - f"Alternatively, set the **W&B entity** field to an org or " - f"team where you already have write access.\n" - f"- The Registry `{registry_name_v}` does not exist in your " - f"org. An org admin can create the Model registry from the " - f"W&B Registry UI.\n" - f"- Your org is on the legacy Model Registry. In that case " - f"use the legacy pattern:\n\n" - f" ```python\n" - f" run.link_artifact(\n" - f" logged,\n" - f" target_path='model-registry/{collection_name_v}',\n" - f" )\n" - f" ```" - ), - kind="danger", + "_Registry linking is disabled — the artifact is logged to the run " + "but not linked to a collection._" + ) ) - link_result - return + + # Close the run so its summary and any Registry link finalize server-side. + wandb.finish() + return collection_name_v, registry_name_v, run @app.cell(hide_code=True) -def _(collection_name_v, mo, registry_name_v, run, train_button): - mo.stop(not train_button.value, mo.md("")) +def _(collection_name_v, mo, registry_name_v, run): + # Renders only after a run exists (it consumes `run` from the training + # cell), so it appears once training finishes. mo.md( f""" - ## Verify - - 1. Open the run page: [{run.name}]({run.url}). Confirm the - **Charts**, **System**, and **Examples** panels are populated. - 2. Click **Artifacts** in the run's left nav. Confirm the - `mnist-cnn-{run.id}` model artifact is listed with metadata - (test accuracy, hyperparameters, number of parameters). - 3. Go to [wandb.ai/registry](https://wandb.ai/registry), open the - **{registry_name_v.title()}** registry, then the - **{collection_name_v}** collection. Confirm the linked version - is present. + ## Verify and next steps - ## Consume the registered model + 1. Open the run: [{run.name}]({run.url}) — check the **Charts**, + **System**, and **Examples** panels. + 2. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed + with its metadata (test accuracy, parameter count, hyperparameters). + 3. At [wandb.ai/registry](https://wandb.ai/registry), open the + **{registry_name_v.title()}** registry, then the **{collection_name_v}** + collection, and confirm the linked version. - From any other script or notebook, fetch the latest registered version: + **Consume the registered model** from any script or notebook: ```python import wandb - api = wandb.Api() - art = api.artifact( + art = wandb.Api().artifact( "wandb-registry-{registry_name_v}/{collection_name_v}:latest" ) - art.download() # writes mnist_cnn.pt under ./artifacts/.../ + art.download() # writes mnist_cnn.pt under ./artifacts/ ``` - ## Next steps - - - **Promote a version.** From the Registry UI, add the `production` - alias to the version you want consumers to pick up. The same - collection path with `:production` will then resolve to it. - - **Compare runs.** Re-run with a deeper architecture or a different - learning rate. Group runs in the W&B UI to compare test accuracy - across configurations. - - **Automate on promotion.** Configure a W&B Automation on the - collection to trigger evaluation jobs or webhooks when a new - version is linked. + **Next steps:** promote a version by adding the `production` alias from + the Registry UI; re-run with a deeper architecture or a different + learning rate and compare runs in the W&B UI; or add a W&B Automation to + trigger evaluation when a new version is linked. """ ) return -@app.cell(hide_code=True) -def _(mo): - mo.md( - """ - ## Finish - - This cell closes the W&B run so the run summary and the Registry - version finalize on the server. - """ - ) - return - - -@app.cell -def _(mo, train_button, wandb): - mo.stop(not train_button.value, mo.md("")) - # Mirrors the `wandb.finish` guard at the top of the training cell: - # leaves the kernel in a clean state for the next Train click. - if wandb.run is not None: - wandb.finish() - mo.md("Run finished. Click **Train model** again to start a new run.") - return - - if __name__ == "__main__": app.run() From d1f04a7a763da752447aa05ec3046099f0a41660 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 15:32:33 -0400 Subject: [PATCH 12/22] address Copilot review: unused os, accuracy precision, path glob - Remove the unused `import os` and drop `os` from the setup cell's returns. - Compute final_acc from the full-precision last-epoch test_acc instead of the display-rounded history entry, so artifact `test_accuracy` keeps full precision. - marimo_molab.yml: use `**/*.py` (canonical glob) instead of `**.py` so the path filter reliably matches notebooks in subdirectories. Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/marimo_molab.yml | 2 +- examples/marimo/mnist-registry/mnist_registry.md | 5 ++--- examples/marimo/mnist-registry/mnist_registry.py | 6 ++---- 3 files changed, 5 insertions(+), 8 deletions(-) diff --git a/.github/workflows/marimo_molab.yml b/.github/workflows/marimo_molab.yml index aa2f5020..d9dbb3e9 100644 --- a/.github/workflows/marimo_molab.yml +++ b/.github/workflows/marimo_molab.yml @@ -15,7 +15,7 @@ on: pull_request_target: types: [opened, synchronize, reopened] paths: - - '**.py' + - '**/*.py' permissions: contents: read diff --git a/examples/marimo/mnist-registry/mnist_registry.md b/examples/marimo/mnist-registry/mnist_registry.md index df505c21..c75d4776 100644 --- a/examples/marimo/mnist-registry/mnist_registry.md +++ b/examples/marimo/mnist-registry/mnist_registry.md @@ -64,8 +64,6 @@ mo.md( # is "everything up to collecting your inputs". It defines the form widgets # but never reads their `.value` — marimo only makes a widget reactive when # a *different* cell consumes it, which the training cell below does. -import os - import torch import torch.nn as nn import torch.nn.functional as F @@ -301,7 +299,8 @@ for epoch in range(1, config["epochs"] + 1): {"epoch": epoch, "test_loss": round(test_loss, 4), "test_acc": round(test_acc, 4)} ) -final_acc = history[-1]["test_acc"] +# Full-precision last-epoch accuracy; `history` rounds only for display. +final_acc = test_acc mo.output.append( mo.vstack( [ diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index 250eac3b..2a3df45c 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -67,8 +67,6 @@ def _(mo): # is "everything up to collecting your inputs". It defines the form widgets # but never reads their `.value` — marimo only makes a widget reactive when # a *different* cell consumes it, which the training cell below does. - import os - import torch import torch.nn as nn import torch.nn.functional as F @@ -160,7 +158,6 @@ def _(mo): momentum, nn, optim, - os, project, registry_name, run_name, @@ -357,7 +354,8 @@ def forward(self, x): {"epoch": epoch, "test_loss": round(test_loss, 4), "test_acc": round(test_acc, 4)} ) - final_acc = history[-1]["test_acc"] + # Full-precision last-epoch accuracy; `history` rounds only for display. + final_acc = test_acc mo.output.append( mo.vstack( [ From fcfef8bf5917ebd10ad7c75ef1c68fe6991bea12 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 16:52:10 -0400 Subject: [PATCH 13/22] move Train button below the training cell; clarify it as the run trigger Reorders the notebook to: intro -> setup+form -> training pipeline -> Train button -> verify. The button still gates the pipeline (so it doesn't run on open), but now reads as the explicit "run the code above" trigger. The training cell's pre-run message explains what it does and points to the button below. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../marimo/mnist-registry/mnist_registry.md | 52 +++++++++-------- .../marimo/mnist-registry/mnist_registry.py | 56 +++++++++++-------- 2 files changed, 62 insertions(+), 46 deletions(-) diff --git a/examples/marimo/mnist-registry/mnist_registry.md b/examples/marimo/mnist-registry/mnist_registry.md index c75d4776..541216a8 100644 --- a/examples/marimo/mnist-registry/mnist_registry.md +++ b/examples/marimo/mnist-registry/mnist_registry.md @@ -142,32 +142,18 @@ mo.vstack( ) ``` -```python {.marimo} -# The button must be its own cell: the training cell reads -# `train_button.value`, and a widget is only reactive when a *different* -# cell consumes it. run_button's value is True for exactly the cascade its -# click triggers, then resets to False — so editing the form afterwards -# re-runs the training cell but it stops immediately instead of retraining. -train_button = mo.ui.run_button(label="Train model", kind="success") -mo.vstack( - [ - mo.md( - "## Train\n\nOne click runs the whole pipeline below: start the " - "run, build and train the model, log metrics and example " - "predictions, save the weights as an Artifact, and link it to the " - "Registry. Click again to retrain with new settings — the previous " - "run is finished first." - ), - train_button, - ] -) -``` - ```python {.marimo} # Everything the Train button triggers, in one cell — no reason to make you # advance through a chain of output-less code blocks. Each milestone is # streamed to the cell output with `mo.output.append` as it happens. -mo.stop(not train_button.value, mo.md("Click **Train model** above to begin.")) +mo.stop( + not train_button.value, + mo.md( + "This cell runs the whole pipeline — start the run, train, log " + "metrics and example predictions, save the model Artifact, and link " + "it to the Registry. Click **Train model** below to run it." + ), +) config = { "epochs": epochs.value, @@ -389,6 +375,28 @@ else: wandb.finish() ``` +```python {.marimo} +# Placed after the training cell on purpose: it's the explicit "run" trigger +# for the pipeline above. It must be its own cell because that cell reads +# `train_button.value`, and a widget only drives reactivity when a +# *different* cell consumes it. The gate also stops the pipeline from +# running automatically when the notebook opens; run_button's value is True +# only for the cascade a click triggers (then resets to False), so editing +# the form afterwards re-runs the training cell but it stops immediately. +train_button = mo.ui.run_button(label="Train model", kind="success") +mo.vstack( + [ + train_button, + mo.md( + "Runs the training cell above. It is gated so it does not " + "execute when the notebook opens — click to run, and click " + "again to retrain after editing the form (the previous run is " + "finished first)." + ), + ] +) +``` + ````python {.marimo hide_code="true"} # Renders only after a run exists (it consumes `run` from the training # cell), so it appears once training finishes. diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index 2a3df45c..0e13ecfa 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -169,29 +169,6 @@ def _(mo): ) -@app.cell -def _(mo): - # The button must be its own cell: the training cell reads - # `train_button.value`, and a widget is only reactive when a *different* - # cell consumes it. run_button's value is True for exactly the cascade its - # click triggers, then resets to False — so editing the form afterwards - # re-runs the training cell but it stops immediately instead of retraining. - train_button = mo.ui.run_button(label="Train model", kind="success") - mo.vstack( - [ - mo.md( - "## Train\n\nOne click runs the whole pipeline below: start the " - "run, build and train the model, log metrics and example " - "predictions, save the weights as an Artifact, and link it to the " - "Registry. Click again to retrain with new settings — the previous " - "run is finished first." - ), - train_button, - ] - ) - return (train_button,) - - @app.cell def _( DataLoader, @@ -222,7 +199,14 @@ def _( # Everything the Train button triggers, in one cell — no reason to make you # advance through a chain of output-less code blocks. Each milestone is # streamed to the cell output with `mo.output.append` as it happens. - mo.stop(not train_button.value, mo.md("Click **Train model** above to begin.")) + mo.stop( + not train_button.value, + mo.md( + "This cell runs the whole pipeline — start the run, train, log " + "metrics and example predictions, save the model Artifact, and link " + "it to the Registry. Click **Train model** below to run it." + ), + ) config = { "epochs": epochs.value, @@ -445,6 +429,30 @@ def forward(self, x): return collection_name_v, registry_name_v, run +@app.cell +def _(mo): + # Placed after the training cell on purpose: it's the explicit "run" trigger + # for the pipeline above. It must be its own cell because that cell reads + # `train_button.value`, and a widget only drives reactivity when a + # *different* cell consumes it. The gate also stops the pipeline from + # running automatically when the notebook opens; run_button's value is True + # only for the cascade a click triggers (then resets to False), so editing + # the form afterwards re-runs the training cell but it stops immediately. + train_button = mo.ui.run_button(label="Train model", kind="success") + mo.vstack( + [ + train_button, + mo.md( + "Runs the training cell above. It is gated so it does not " + "execute when the notebook opens — click to run, and click " + "again to retrain after editing the form (the previous run is " + "finished first)." + ), + ] + ) + return (train_button,) + + @app.cell(hide_code=True) def _(collection_name_v, mo, registry_name_v, run): # Renders only after a run exists (it consumes `run` from the training From ce1891498126e0045cf1afe4cff688ad6fc1eb63 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 16:56:07 -0400 Subject: [PATCH 14/22] clarify W&B entity field label (username or team) The entity form field read just "W&B entity", which reads as jargon. Label it "username or team" to match the prerequisites wording. Also tighten the registry-link remediation: an entity is a team or username (not an org), so it now reads "set W&B entity to a team in an org where you have Registry write access." Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/marimo/mnist-registry/mnist_registry.md | 8 +++++--- examples/marimo/mnist-registry/mnist_registry.py | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/examples/marimo/mnist-registry/mnist_registry.md b/examples/marimo/mnist-registry/mnist_registry.md index 541216a8..ae322c40 100644 --- a/examples/marimo/mnist-registry/mnist_registry.md +++ b/examples/marimo/mnist-registry/mnist_registry.md @@ -103,7 +103,9 @@ momentum = mo.ui.slider( seed = mo.ui.number(start=0, stop=99999, value=42, label="Random seed") project = mo.ui.text(value="marimo-mnist-registry", label="W&B project") -entity = mo.ui.text(value="", label="W&B entity (blank uses your default)") +entity = mo.ui.text( + value="", label="W&B entity — username or team (blank uses your default)" +) run_name = mo.ui.text(value="", label="Run name (blank auto-generates)") api_key = mo.ui.text( value="", kind="password", label="W&B API key (blank uses your shell login)" @@ -353,8 +355,8 @@ if link_to_registry.value: f"`update_member()`), or SCIM (`PATCH /scim/Users/{{id}}` with " f"`registryRoles`) — see " f"https://docs.wandb.ai/guides/registry/configure_registry/. " - f"Or set **W&B entity** to an org/team where you have write " - f"access.\n" + f"Or set **W&B entity** to a team in an org where you have " + f"Registry write access.\n" f"- The Registry `{registry_name_v}` may not exist; an admin " f"can create it from the W&B Registry UI.\n" f"- On the legacy Model Registry, link with " diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index 0e13ecfa..5a6b79f8 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -106,7 +106,9 @@ def _(mo): seed = mo.ui.number(start=0, stop=99999, value=42, label="Random seed") project = mo.ui.text(value="marimo-mnist-registry", label="W&B project") - entity = mo.ui.text(value="", label="W&B entity (blank uses your default)") + entity = mo.ui.text( + value="", label="W&B entity — username or team (blank uses your default)" + ) run_name = mo.ui.text(value="", label="Run name (blank auto-generates)") api_key = mo.ui.text( value="", kind="password", label="W&B API key (blank uses your shell login)" @@ -406,8 +408,8 @@ def forward(self, x): f"`update_member()`), or SCIM (`PATCH /scim/Users/{{id}}` with " f"`registryRoles`) — see " f"https://docs.wandb.ai/guides/registry/configure_registry/. " - f"Or set **W&B entity** to an org/team where you have write " - f"access.\n" + f"Or set **W&B entity** to a team in an org where you have " + f"Registry write access.\n" f"- The Registry `{registry_name_v}` may not exist; an admin " f"can create it from the W&B Registry UI.\n" f"- On the legacy Model Registry, link with " From aa342d149205716926ca898a4bda8eb813bb5254 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 17:04:23 -0400 Subject: [PATCH 15/22] correct entity guidance: runs require a team, not a username Personal-username entities were removed for W&B accounts created after 2024-05-21, so a run's entity must be a team. Relabel the entity field accordingly, fix the prerequisite, and wrap wandb.init so an "entity ... not found" failure renders actionable guidance (set the entity to a team) instead of a raw CommError traceback. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../marimo/mnist-registry/mnist_registry.md | 36 ++++++++++++++----- .../marimo/mnist-registry/mnist_registry.py | 36 ++++++++++++++----- 2 files changed, 54 insertions(+), 18 deletions(-) diff --git a/examples/marimo/mnist-registry/mnist_registry.md b/examples/marimo/mnist-registry/mnist_registry.md index ae322c40..78e0869d 100644 --- a/examples/marimo/mnist-registry/mnist_registry.md +++ b/examples/marimo/mnist-registry/mnist_registry.md @@ -48,7 +48,9 @@ mo.md( your shell before starting marimo, or paste your key into the **W&B API key** field in the form below. Get your key from [wandb.ai/authorize](https://wandb.ai/authorize). - - A W&B entity (your user or a team) the run will be written to. + - A W&B **team** to write the run to, set in the **W&B entity** field. + Accounts created after May 2024 have no personal entity, so the run + must go to a team — your username will not work as an entity. - A **W&B Registry** must exist in your org, and your account needs at least the **Member** role on it (linking an artifact is a write action). The built-in Model registry is provisioned automatically in @@ -104,7 +106,7 @@ seed = mo.ui.number(start=0, stop=99999, value=42, label="Random seed") project = mo.ui.text(value="marimo-mnist-registry", label="W&B project") entity = mo.ui.text( - value="", label="W&B entity — username or team (blank uses your default)" + value="", label="W&B entity — a team you belong to (blank uses your default)" ) run_name = mo.ui.text(value="", label="Run name (blank auto-generates)") api_key = mo.ui.text( @@ -180,13 +182,29 @@ if api_key.value: torch.manual_seed(config["seed"]) -run = wandb.init( - project=project.value or None, - entity=entity.value or None, - name=run_name.value or None, - config=config, - job_type="train", -) +try: + run = wandb.init( + project=project.value or None, + entity=entity.value or None, + name=run_name.value or None, + config=config, + job_type="train", + ) +except Exception as exc: # noqa: BLE001 - turn the raw traceback into guidance + mo.stop( + True, + mo.callout( + mo.md( + f"**Could not start the run.** `{exc}`\n\n" + f"An `entity ... not found` error means the **W&B entity** is " + f"not a team you can write to. Personal-username entities were " + f"removed for accounts created after 21 May 2024, so set the " + f"**W&B entity** field to one of your teams (find them in the " + f"left sidebar at [wandb.ai](https://wandb.ai))." + ), + kind="danger", + ), + ) # Use `epoch` as the x-axis for train/test metrics in the W&B UI. wandb.define_metric("epoch") wandb.define_metric("train/*", step_metric="epoch") diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index 5a6b79f8..573303a6 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -49,7 +49,9 @@ def _(): your shell before starting marimo, or paste your key into the **W&B API key** field in the form below. Get your key from [wandb.ai/authorize](https://wandb.ai/authorize). - - A W&B entity (your user or a team) the run will be written to. + - A W&B **team** to write the run to, set in the **W&B entity** field. + Accounts created after May 2024 have no personal entity, so the run + must go to a team — your username will not work as an entity. - A **W&B Registry** must exist in your org, and your account needs at least the **Member** role on it (linking an artifact is a write action). The built-in Model registry is provisioned automatically in @@ -107,7 +109,7 @@ def _(mo): project = mo.ui.text(value="marimo-mnist-registry", label="W&B project") entity = mo.ui.text( - value="", label="W&B entity — username or team (blank uses your default)" + value="", label="W&B entity — a team you belong to (blank uses your default)" ) run_name = mo.ui.text(value="", label="Run name (blank auto-generates)") api_key = mo.ui.text( @@ -233,13 +235,29 @@ def _( torch.manual_seed(config["seed"]) - run = wandb.init( - project=project.value or None, - entity=entity.value or None, - name=run_name.value or None, - config=config, - job_type="train", - ) + try: + run = wandb.init( + project=project.value or None, + entity=entity.value or None, + name=run_name.value or None, + config=config, + job_type="train", + ) + except Exception as exc: # noqa: BLE001 - turn the raw traceback into guidance + mo.stop( + True, + mo.callout( + mo.md( + f"**Could not start the run.** `{exc}`\n\n" + f"An `entity ... not found` error means the **W&B entity** is " + f"not a team you can write to. Personal-username entities were " + f"removed for accounts created after 21 May 2024, so set the " + f"**W&B entity** field to one of your teams (find them in the " + f"left sidebar at [wandb.ai](https://wandb.ai))." + ), + kind="danger", + ), + ) # Use `epoch` as the x-axis for train/test metrics in the W&B UI. wandb.define_metric("epoch") wandb.define_metric("train/*", step_metric="epoch") From 4142351cdb117070f3f707450187894f8c1e9bfd Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 17:27:22 -0400 Subject: [PATCH 16/22] add a consume step that downloads the model and classifies 10 digits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pulls the CNN definition into its own cell so it can be shared, then adds a cell that downloads the model from W&B (preferring the registered version, falling back to the run's artifact so it works even when the registry link is blocked), loads the weights into a fresh network, and classifies 10 held-out MNIST test digits — rendering each as an image with predicted vs. true label and an N/10 correct tally. No new dependency (numpy ships with torch; mo.image renders the arrays). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../marimo/mnist-registry/mnist_registry.md | 99 ++++++++++++---- .../marimo/mnist-registry/mnist_registry.py | 108 ++++++++++++++---- 2 files changed, 165 insertions(+), 42 deletions(-) diff --git a/examples/marimo/mnist-registry/mnist_registry.md b/examples/marimo/mnist-registry/mnist_registry.md index 78e0869d..80dd8bcf 100644 --- a/examples/marimo/mnist-registry/mnist_registry.md +++ b/examples/marimo/mnist-registry/mnist_registry.md @@ -146,6 +146,32 @@ mo.vstack( ) ``` +```python {.marimo} +class Net(nn.Module): + """Small CNN: 2 conv layers (10, 20 filters, 5x5) + 2 FC (50, 10). + + Defined in its own cell so the training cell and the consume cell can + share it (marimo forbids defining the same name in two cells). + """ + + def __init__(self): + super().__init__() + self.conv1 = nn.Conv2d(1, 10, kernel_size=5) + self.conv2 = nn.Conv2d(10, 20, kernel_size=5) + self.conv2_drop = nn.Dropout2d() + self.fc1 = nn.Linear(320, 50) + self.fc2 = nn.Linear(50, 10) + + def forward(self, x): + x = F.relu(F.max_pool2d(self.conv1(x), 2)) + x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) + x = x.view(-1, 320) + x = F.relu(self.fc1(x)) + x = F.dropout(x, training=self.training) + x = self.fc2(x) + return F.log_softmax(x, dim=1) +``` + ```python {.marimo} # Everything the Train button triggers, in one cell — no reason to make you # advance through a chain of output-less code blocks. Each milestone is @@ -212,26 +238,6 @@ wandb.define_metric("test/*", step_metric="epoch") # Surface the run link right away so you can watch metrics stream live. mo.output.append(mo.md(f"**Run started:** [`{run.name}`]({run.url})")) -class Net(nn.Module): - """Small CNN: 2 conv layers (10, 20 filters, 5x5) + 2 FC (50, 10).""" - - def __init__(self): - super().__init__() - self.conv1 = nn.Conv2d(1, 10, kernel_size=5) - self.conv2 = nn.Conv2d(10, 20, kernel_size=5) - self.conv2_drop = nn.Dropout2d() - self.fc1 = nn.Linear(320, 50) - self.fc2 = nn.Linear(50, 10) - - def forward(self, x): - x = F.relu(F.max_pool2d(self.conv1(x), 2)) - x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) - x = x.view(-1, 320) - x = F.relu(self.fc1(x)) - x = F.dropout(x, training=self.training) - x = self.fc2(x) - return F.log_softmax(x, dim=1) - model = Net().to(device) # `log="gradients"` is the standard choice; `log="all"` also logs parameter # histograms at extra cost. @@ -417,6 +423,59 @@ mo.vstack( ) ``` +```python {.marimo} +# Consume the model: download it from W&B (preferring the registered +# version, falling back to the run's own artifact), load the weights into a +# fresh network, and classify 10 held-out test digits. +api = wandb.Api() +try: + consumed = api.artifact( + f"wandb-registry-{registry_name_v}/{collection_name_v}:latest", type="model" + ) + source = f"registry `wandb-registry-{registry_name_v}/{collection_name_v}:latest`" +except Exception: # noqa: BLE001 - registry link may be absent (e.g. a view-only seat) + consumed = api.artifact( + f"{run.entity}/{run.project}/mnist-cnn-{run.id}:latest", type="model" + ) + source = f"run artifact `mnist-cnn-{run.id}:latest`" +weights_dir = consumed.download() + +clf = Net() +clf.load_state_dict(torch.load(f"{weights_dir}/mnist_cnn.pt", map_location="cpu")) +clf.eval() + +cards = [] +correct = 0 +with torch.no_grad(): + for i in range(10): + image, true_label = test_ds[i] + pred = clf(image.unsqueeze(0)).argmax(dim=1).item() + correct += int(pred == true_label) + # Undo the Normalize transform so the digit renders as a clean image. + digit = (image * 0.3081 + 0.1307).clamp(0, 1).squeeze().numpy() + mark = "✅" if pred == true_label else "❌" + cards.append( + mo.vstack( + [ + mo.image(digit, width=64, vmin=0, vmax=1), + mo.md(f"{mark} **{pred}** · true {true_label}"), + ], + align="center", + ) + ) + +mo.vstack( + [ + mo.md( + f"## Classify 10 test digits\n\nConsumed the model from {source}, " + f"loaded the weights into a fresh network, and ran it on 10 held-out " + f"MNIST test images — **{correct}/10 correct**." + ), + mo.hstack(cards, wrap=True, justify="start"), + ] +) +``` + ````python {.marimo hide_code="true"} # Renders only after a run exists (it consumes `run` from the training # cell), so it appears once training finishes. diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index 573303a6..9a1c48cd 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -173,10 +173,40 @@ def _(mo): ) +@app.cell +def _(F, nn): + class Net(nn.Module): + """Small CNN: 2 conv layers (10, 20 filters, 5x5) + 2 FC (50, 10). + + Defined in its own cell so the training cell and the consume cell can + share it (marimo forbids defining the same name in two cells). + """ + + def __init__(self): + super().__init__() + self.conv1 = nn.Conv2d(1, 10, kernel_size=5) + self.conv2 = nn.Conv2d(10, 20, kernel_size=5) + self.conv2_drop = nn.Dropout2d() + self.fc1 = nn.Linear(320, 50) + self.fc2 = nn.Linear(50, 10) + + def forward(self, x): + x = F.relu(F.max_pool2d(self.conv1(x), 2)) + x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) + x = x.view(-1, 320) + x = F.relu(self.fc1(x)) + x = F.dropout(x, training=self.training) + x = self.fc2(x) + return F.log_softmax(x, dim=1) + + return (Net,) + + @app.cell def _( DataLoader, F, + Net, api_key, batch_size, collection_name, @@ -188,7 +218,6 @@ def _( lr, momentum, mo, - nn, optim, project, registry_name, @@ -265,26 +294,6 @@ def _( # Surface the run link right away so you can watch metrics stream live. mo.output.append(mo.md(f"**Run started:** [`{run.name}`]({run.url})")) - class Net(nn.Module): - """Small CNN: 2 conv layers (10, 20 filters, 5x5) + 2 FC (50, 10).""" - - def __init__(self): - super().__init__() - self.conv1 = nn.Conv2d(1, 10, kernel_size=5) - self.conv2 = nn.Conv2d(10, 20, kernel_size=5) - self.conv2_drop = nn.Dropout2d() - self.fc1 = nn.Linear(320, 50) - self.fc2 = nn.Linear(50, 10) - - def forward(self, x): - x = F.relu(F.max_pool2d(self.conv1(x), 2)) - x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) - x = x.view(-1, 320) - x = F.relu(self.fc1(x)) - x = F.dropout(x, training=self.training) - x = self.fc2(x) - return F.log_softmax(x, dim=1) - model = Net().to(device) # `log="gradients"` is the standard choice; `log="all"` also logs parameter # histograms at extra cost. @@ -446,7 +455,7 @@ def forward(self, x): # Close the run so its summary and any Registry link finalize server-side. wandb.finish() - return collection_name_v, registry_name_v, run + return collection_name_v, registry_name_v, run, test_ds @app.cell @@ -473,6 +482,61 @@ def _(mo): return (train_button,) +@app.cell +def _(Net, collection_name_v, mo, registry_name_v, run, test_ds, torch, wandb): + # Consume the model: download it from W&B (preferring the registered + # version, falling back to the run's own artifact), load the weights into a + # fresh network, and classify 10 held-out test digits. + api = wandb.Api() + try: + consumed = api.artifact( + f"wandb-registry-{registry_name_v}/{collection_name_v}:latest", type="model" + ) + source = f"registry `wandb-registry-{registry_name_v}/{collection_name_v}:latest`" + except Exception: # noqa: BLE001 - registry link may be absent (e.g. a view-only seat) + consumed = api.artifact( + f"{run.entity}/{run.project}/mnist-cnn-{run.id}:latest", type="model" + ) + source = f"run artifact `mnist-cnn-{run.id}:latest`" + weights_dir = consumed.download() + + clf = Net() + clf.load_state_dict(torch.load(f"{weights_dir}/mnist_cnn.pt", map_location="cpu")) + clf.eval() + + cards = [] + correct = 0 + with torch.no_grad(): + for i in range(10): + image, true_label = test_ds[i] + pred = clf(image.unsqueeze(0)).argmax(dim=1).item() + correct += int(pred == true_label) + # Undo the Normalize transform so the digit renders as a clean image. + digit = (image * 0.3081 + 0.1307).clamp(0, 1).squeeze().numpy() + mark = "✅" if pred == true_label else "❌" + cards.append( + mo.vstack( + [ + mo.image(digit, width=64, vmin=0, vmax=1), + mo.md(f"{mark} **{pred}** · true {true_label}"), + ], + align="center", + ) + ) + + mo.vstack( + [ + mo.md( + f"## Classify 10 test digits\n\nConsumed the model from {source}, " + f"loaded the weights into a fresh network, and ran it on 10 held-out " + f"MNIST test images — **{correct}/10 correct**." + ), + mo.hstack(cards, wrap=True, justify="start"), + ] + ) + return + + @app.cell(hide_code=True) def _(collection_name_v, mo, registry_name_v, run): # Renders only after a run exists (it consumes `run` from the training From 46e06861db07de0b29cbc784de442c3142b9e086 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Fri, 12 Jun 2026 17:35:00 -0400 Subject: [PATCH 17/22] fix marimo redefinition error in the consume cell MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The consume cell reused `correct` and `pred`, which the training cell already defines — marimo forbids defining a name in more than one cell, so the cell refused to run. Rename them to `n_correct` / `prediction` (both cell-local to the consume step). marimo export does not enforce the single-definition rule, only the kernel does, so this slipped through earlier validation. Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/marimo/mnist-registry/mnist_registry.md | 12 ++++++------ examples/marimo/mnist-registry/mnist_registry.py | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/examples/marimo/mnist-registry/mnist_registry.md b/examples/marimo/mnist-registry/mnist_registry.md index 80dd8bcf..b1c2077c 100644 --- a/examples/marimo/mnist-registry/mnist_registry.md +++ b/examples/marimo/mnist-registry/mnist_registry.md @@ -445,20 +445,20 @@ clf.load_state_dict(torch.load(f"{weights_dir}/mnist_cnn.pt", map_location="cpu" clf.eval() cards = [] -correct = 0 +n_correct = 0 with torch.no_grad(): for i in range(10): image, true_label = test_ds[i] - pred = clf(image.unsqueeze(0)).argmax(dim=1).item() - correct += int(pred == true_label) + prediction = clf(image.unsqueeze(0)).argmax(dim=1).item() + n_correct += int(prediction == true_label) # Undo the Normalize transform so the digit renders as a clean image. digit = (image * 0.3081 + 0.1307).clamp(0, 1).squeeze().numpy() - mark = "✅" if pred == true_label else "❌" + mark = "✅" if prediction == true_label else "❌" cards.append( mo.vstack( [ mo.image(digit, width=64, vmin=0, vmax=1), - mo.md(f"{mark} **{pred}** · true {true_label}"), + mo.md(f"{mark} **{prediction}** · true {true_label}"), ], align="center", ) @@ -469,7 +469,7 @@ mo.vstack( mo.md( f"## Classify 10 test digits\n\nConsumed the model from {source}, " f"loaded the weights into a fresh network, and ran it on 10 held-out " - f"MNIST test images — **{correct}/10 correct**." + f"MNIST test images — **{n_correct}/10 correct**." ), mo.hstack(cards, wrap=True, justify="start"), ] diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index 9a1c48cd..cb78c945 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -505,20 +505,20 @@ def _(Net, collection_name_v, mo, registry_name_v, run, test_ds, torch, wandb): clf.eval() cards = [] - correct = 0 + n_correct = 0 with torch.no_grad(): for i in range(10): image, true_label = test_ds[i] - pred = clf(image.unsqueeze(0)).argmax(dim=1).item() - correct += int(pred == true_label) + prediction = clf(image.unsqueeze(0)).argmax(dim=1).item() + n_correct += int(prediction == true_label) # Undo the Normalize transform so the digit renders as a clean image. digit = (image * 0.3081 + 0.1307).clamp(0, 1).squeeze().numpy() - mark = "✅" if pred == true_label else "❌" + mark = "✅" if prediction == true_label else "❌" cards.append( mo.vstack( [ mo.image(digit, width=64, vmin=0, vmax=1), - mo.md(f"{mark} **{pred}** · true {true_label}"), + mo.md(f"{mark} **{prediction}** · true {true_label}"), ], align="center", ) @@ -529,7 +529,7 @@ def _(Net, collection_name_v, mo, registry_name_v, run, test_ds, torch, wandb): mo.md( f"## Classify 10 test digits\n\nConsumed the model from {source}, " f"loaded the weights into a fresh network, and ran it on 10 held-out " - f"MNIST test images — **{correct}/10 correct**." + f"MNIST test images — **{n_correct}/10 correct**." ), mo.hstack(cards, wrap=True, justify="start"), ] From 91c8161ee921bbd7673822026c384d1e3d59f4af Mon Sep 17 00:00:00 2001 From: Akshay Agrawal Date: Tue, 30 Jun 2026 07:59:12 -0700 Subject: [PATCH 18/22] Refactor the MNIST registry notebook for readability and usability (#626) * Refactor the mnist notebook for readability and usability - Move imports and constants into a setup cell - Wrap the UI elements in a submittable form, removing the ui.run_button() - Refactor the monolithic training and logging cell into multiple helper functions, each declared in their own cell as a reusable function - Separate logic cells from view cells when possible - Clean up the evaluation presentation, using a table instead of markdown+emoji - Add a mo.outline() so users can see an outline of the notebook at a glance Core principles: * Gate execution once on mo.stop(), then let the graph handle the rest. All other cells reference run/config, which don't exist until the form is submitted. This lets us remove lots of confusing conditions. * Separate logic from presentation * Push temporary variables into functions to reduce notebook globals. marimo notebooks work best when there are as few global variables as possible. * docs: export marimo notebook(s) to Markdown [skip ci] --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- .../marimo/mnist-registry/mnist_registry.md | 738 ++++++++------ .../marimo/mnist-registry/mnist_registry.py | 905 ++++++++++-------- 2 files changed, 934 insertions(+), 709 deletions(-) diff --git a/examples/marimo/mnist-registry/mnist_registry.md b/examples/marimo/mnist-registry/mnist_registry.md index b1c2077c..f1c651dd 100644 --- a/examples/marimo/mnist-registry/mnist_registry.md +++ b/examples/marimo/mnist-registry/mnist_registry.md @@ -26,46 +26,8 @@ header: |- """ --- -```python {.marimo hide_code="true"} +```python {.marimo hide_code="true" name="setup"} import marimo as mo - -mo.md( - """ - # MNIST -> W&B Run -> Registry - - ## What you will build - - - A **W&B run** with training and test metrics, gradient histograms, - and example test-set predictions logged as images. - - A **model Artifact** named `mnist-cnn-` of type `model`, - carrying metadata (test accuracy, parameter count, hyperparameters). - - A version of that Artifact **linked into a W&B Registry collection** - so it appears under registered models org-wide. - - ## Prerequisites - - - Authenticate with W&B one of two ways: run **`wandb login`** in - your shell before starting marimo, or paste your key into the - **W&B API key** field in the form below. Get your key from - [wandb.ai/authorize](https://wandb.ai/authorize). - - A W&B **team** to write the run to, set in the **W&B entity** field. - Accounts created after May 2024 have no personal entity, so the run - must go to a team — your username will not work as an entity. - - A **W&B Registry** must exist in your org, and your account needs at - least the **Member** role on it (linking an artifact is a write - action). The built-in Model registry is provisioned automatically in - newer orgs. If linking fails (for example, from a view-only seat), - the run still completes and the Registry step explains how to fix it. - - A GPU is optional. The defaults finish in about 2 minutes on CPU. - """ -) -``` - -```python {.marimo} -# Imports, device detection, and the input form all live in one cell: this -# is "everything up to collecting your inputs". It defines the form widgets -# but never reads their `.value` — marimo only makes a widget reactive when -# a *different* cell consumes it, which the training cell below does. import torch import torch.nn as nn import torch.nn.functional as F @@ -91,62 +53,143 @@ else: "hyperparameters this takes about 2 minutes." ) device_kind = "warn" +``` +# MNIST -> W&B Run -> Registry + +## What you will build + +- A **W&B run** with training and test metrics, gradient histograms, + and example test-set predictions logged as images. +- A **model Artifact** named `mnist-cnn-` of type `model`, + carrying metadata (test accuracy, parameter count, hyperparameters). +- A version of that Artifact **linked into a W&B Registry collection** + so it appears under registered models org-wide. + +## Prerequisites + +- Authenticate with W&B one of two ways: run **`wandb login`** in + your shell before starting marimo, or paste your key into the + **W&B API key** field in the form below. Get your key from + [wandb.ai/authorize](https://wandb.ai/authorize). +- A W&B **team** to write the run to, set in the **W&B entity** field. + Accounts created after May 2024 have no personal entity, so the run + must go to a team — your username will not work as an entity. +- A **W&B Registry** must exist in your org, and your account needs at + least the **Member** role on it (linking an artifact is a write + action). The built-in Model registry is provisioned automatically in + newer orgs. If linking fails (for example, from a view-only seat), + the run still completes and the Registry step explains how to fix it. +- A GPU is optional. The defaults finish in about 2 minutes on CPU. + +```python {.marimo} +mo.outline() +``` + +```python {.marimo hide_code="true"} +mo.callout( + mo.md(f"**Device:** `{device}`. {device_note}"), + kind=device_kind, +) +``` + +## Configuration + +Set the hyperparameters and W&B targets, then click **Train model** below. + +```python {.marimo hide_code="true"} epochs = mo.ui.slider(start=1, stop=10, step=1, value=3, label="Epochs") batch_size = mo.ui.dropdown( options=["32", "64", "128", "256"], value="64", label="Batch size" ) lr = mo.ui.slider( - start=0.001, stop=0.1, step=0.001, value=0.01, label="Learning rate", show_value=True + start=0.001, + stop=0.1, + step=0.001, + value=0.01, + label="Learning rate", + show_value=True, ) momentum = mo.ui.slider( - start=0.0, stop=0.99, step=0.01, value=0.5, label="SGD momentum", show_value=True + start=0.0, + stop=0.99, + step=0.01, + value=0.5, + label="SGD momentum", + show_value=True, ) seed = mo.ui.number(start=0, stop=99999, value=42, label="Random seed") project = mo.ui.text(value="marimo-mnist-registry", label="W&B project") entity = mo.ui.text( - value="", label="W&B entity — a team you belong to (blank uses your default)" + value="", + label="W&B entity \u2014 a team you belong to (blank uses your default)", ) run_name = mo.ui.text(value="", label="Run name (blank auto-generates)") api_key = mo.ui.text( - value="", kind="password", label="W&B API key (blank uses your shell login)" + value="", + kind="password", + label="W&B API key (blank uses your shell login)", ) registry_name = mo.ui.text(value="model", label="W&B Registry name") -collection_name = mo.ui.text(value="MNIST Classifiers", label="Registry collection") -link_to_registry = mo.ui.checkbox(value=True, label="Link artifact to Registry") - -form = mo.vstack( - [ - mo.md("### Training"), - mo.hstack([epochs, batch_size]), - mo.hstack([lr, momentum]), - seed, - mo.md("### W&B run"), - api_key, - mo.hstack([project, entity, run_name]), - mo.md("### Registry"), - mo.hstack([registry_name, collection_name, link_to_registry]), - ] +collection_name = mo.ui.text( + value="MNIST Classifiers", label="Registry collection" +) +link_to_registry = mo.ui.checkbox( + value=True, label="Link artifact to Registry" ) -mo.vstack( - [ - mo.callout( - mo.md(f"**Device:** `{device}` — {device_note}"), kind=device_kind - ), - mo.md( - "## Configure\n\nSet the hyperparameters and W&B targets, then click " - "**Train model** below. Changing a value here never starts a run on " - "its own — only the button does." - ), - form, - ] +# Batch every control into one form so training only kicks off on submit. +# `form.value` is None until the user clicks Train model, then becomes a dict +# keyed by the names below \u2014 the training cell gates on that. +form = ( + mo.md( + """ + **Training.** + + {epochs} {batch_size} + + {lr} {momentum} + + {seed} + + **W&B run.** + + {api_key} + + {project} + + {entity} + + {run_name} + + **Registry.** + + {registry_name} {collection_name} {link_to_registry} + """ + ) + .batch( + epochs=epochs, + batch_size=batch_size, + lr=lr, + momentum=momentum, + seed=seed, + api_key=api_key, + project=project, + entity=entity, + run_name=run_name, + registry_name=registry_name, + collection_name=collection_name, + link_to_registry=link_to_registry, + ) + .form(submit_button_label="Train model", bordered=False) ) + +form ``` -```python {.marimo} +```python {.marimo name="Net"} class Net(nn.Module): """Small CNN: 2 conv layers (10, 20 filters, 5x5) + 2 FC (50, 10). @@ -172,56 +215,56 @@ class Net(nn.Module): return F.log_softmax(x, dim=1) ``` +## Training + ```python {.marimo} -# Everything the Train button triggers, in one cell — no reason to make you -# advance through a chain of output-less code blocks. Each milestone is -# streamed to the cell output with `mo.output.append` as it happens. mo.stop( - not train_button.value, + form.value is None, mo.md( - "This cell runs the whole pipeline — start the run, train, log " - "metrics and example predictions, save the model Artifact, and link " - "it to the Registry. Click **Train model** below to run it." + "Training hasn't started yet. Fill in the form above and click " + "**Train model** to start the run — it trains the model, logs metrics " + "and example predictions, saves the weights as an Artifact, and links " + "them to the Registry." ), ) +cfg = form.value config = { - "epochs": epochs.value, - "batch_size": int(batch_size.value), - "lr": lr.value, - "momentum": momentum.value, - "seed": seed.value, + "epochs": cfg["epochs"], + "batch_size": int(cfg["batch_size"]), + "lr": cfg["lr"], + "momentum": cfg["momentum"], + "seed": cfg["seed"], "architecture": "CNN", "dataset": "MNIST", } -registry_name_v = registry_name.value.strip() -collection_name_v = collection_name.value.strip() +registry_name_v = cfg["registry_name"].strip() +collection_name_v = cfg["collection_name"].strip() -# Authenticate. Finish any prior run first (marimo keeps the kernel alive -# across re-clicks). A key pasted into the form wins; otherwise fall back to -# ambient login (shell `wandb login`, WANDB_API_KEY, or netrc). The key is -# never written to the run config. +# Authenticate and start the run. Finish any prior run first (marimo keeps the +# kernel alive across re-submits). A key pasted into the form wins; otherwise +# fall back to ambient login (shell `wandb login`, WANDB_API_KEY, or netrc). +# The key is never written to the run config. if wandb.run is not None: wandb.finish() -if api_key.value: - wandb.login(key=api_key.value) - +if cfg["api_key"]: + wandb.login(key=cfg["api_key"]) torch.manual_seed(config["seed"]) try: run = wandb.init( - project=project.value or None, - entity=entity.value or None, - name=run_name.value or None, + project=cfg["project"] or None, + entity=cfg["entity"] or None, + name=cfg["run_name"] or None, config=config, job_type="train", ) -except Exception as exc: # noqa: BLE001 - turn the raw traceback into guidance +except Exception as init_exc: # noqa: BLE001 - turn the raw traceback into guidance mo.stop( True, mo.callout( mo.md( - f"**Could not start the run.** `{exc}`\n\n" + f"**Could not start the run.** `{init_exc}`\n\n" f"An `entity ... not found` error means the **W&B entity** is " f"not a team you can write to. Personal-username entities were " f"removed for accounts created after 21 May 2024, so set the " @@ -231,220 +274,139 @@ except Exception as exc: # noqa: BLE001 - turn the raw traceback into guidance kind="danger", ), ) + # Use `epoch` as the x-axis for train/test metrics in the W&B UI. wandb.define_metric("epoch") wandb.define_metric("train/*", step_metric="epoch") wandb.define_metric("test/*", step_metric="epoch") +``` + +```python {.marimo hide_code="true"} # Surface the run link right away so you can watch metrics stream live. -mo.output.append(mo.md(f"**Run started:** [`{run.name}`]({run.url})")) - -model = Net().to(device) -# `log="gradients"` is the standard choice; `log="all"` also logs parameter -# histograms at extra cost. -wandb.watch(model, log="gradients", log_freq=100) -optimizer = optim.SGD( - model.parameters(), lr=config["lr"], momentum=config["momentum"] -) +mo.md(f"**Run started:** [`{run.name}`]({run.url})") +``` -transform = transforms.Compose( - [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] -) -train_ds = datasets.MNIST("./data", train=True, download=True, transform=transform) -test_ds = datasets.MNIST("./data", train=False, download=True, transform=transform) -loader_kwargs = ( - {"num_workers": 2, "pin_memory": True} if device.type == "cuda" else {} -) -train_loader = DataLoader( - train_ds, batch_size=config["batch_size"], shuffle=True, **loader_kwargs +```python {.marimo} +train_ds, test_ds = load_data() +model, history, final_acc, best_acc = run_training( + run, config, train_ds, test_ds ) -test_loader = DataLoader(test_ds, batch_size=1000, shuffle=False, **loader_kwargs) -history = [] -best_acc = 0.0 -for epoch in range(1, config["epochs"] + 1): - model.train() - for batch_idx, (data, target) in enumerate( - tqdm(train_loader, desc=f"epoch {epoch}/{config['epochs']}") - ): - data, target = data.to(device), target.to(device) - optimizer.zero_grad() - output = model(data) - loss = F.nll_loss(output, target) - loss.backward() - optimizer.step() - if batch_idx % 50 == 0: - wandb.log({"train/loss": loss.item(), "epoch": epoch}) +mo.vstack( + [ + mo.md("### Training summary"), + mo.ui.table(history, selection=None), + mo.md(f"**Final test accuracy:** {final_acc:.2%}"), + ] +) +``` - model.eval() - test_loss = 0.0 - correct = 0 - example_images = [] - with torch.no_grad(): - for data, target in test_loader: - data, target = data.to(device), target.to(device) - output = model(data) - test_loss += F.nll_loss(output, target, reduction="sum").item() - pred = output.argmax(dim=1, keepdim=True) - correct += pred.eq(target.view_as(pred)).sum().item() - # Pull up to 16 example predictions from the first batch. - while len(example_images) < 16 and len(example_images) < data.size(0): - j = len(example_images) - example_images.append( - wandb.Image( - data[j], - caption=f"pred={pred[j].item()} true={target[j].item()}", - ) - ) +```python {.marimo} +logged, artifact_name = save_and_log_artifact( + run, + model, + config, + train_size=len(train_ds), + test_size=len(test_ds), + final_acc=final_acc, + best_acc=best_acc, +) - test_loss /= len(test_loader.dataset) - test_acc = correct / len(test_loader.dataset) - best_acc = max(best_acc, test_acc) - wandb.log( - { - "test/loss": test_loss, - "test/accuracy": test_acc, - "epoch": epoch, - "examples": example_images, +# Link to the Registry unless disabled, capturing the outcome for display +# rather than crashing the pipeline. +if not cfg["link_to_registry"]: + registry_status = {"kind": "disabled"} +else: + try: + registry_status = { + "kind": "linked", + "target_path": link_artifact_to_registry( + run, logged, registry_name_v, collection_name_v + ), + } + except Exception as link_exc: # noqa: BLE001 - surface any failure to the reader + registry_status = { + "kind": "failed", + "target_path": f"wandb-registry-{registry_name_v}/{collection_name_v}", + "error": str(link_exc), } - ) - history.append( - {"epoch": epoch, "test_loss": round(test_loss, 4), "test_acc": round(test_acc, 4)} - ) -# Full-precision last-epoch accuracy; `history` rounds only for display. -final_acc = test_acc -mo.output.append( - mo.vstack( - [ - mo.md("### Training summary"), - mo.ui.table(history, selection=None), - mo.md(f"**Final test accuracy:** {final_acc:.2%}"), - ] - ) -) +# Close the run so its summary and any Registry link finalize server-side. +wandb.finish() +``` -# Save the weights and log them as a model Artifact tagged `latest`. -model_path = "mnist_cnn.pt" -torch.save(model.state_dict(), model_path) -artifact = wandb.Artifact( - name=f"mnist-cnn-{run.id}", - type="model", - description=( - "Small CNN trained on MNIST. Architecture: 2 conv layers " - "(10 and 20 filters, 5x5 kernels) + 2 FC layers (50, 10)." - ), - metadata={ - "framework": "pytorch", - "architecture": "CNN", - "num_parameters": sum(p.numel() for p in model.parameters()), - "dataset": "MNIST", - "train_size": len(train_ds), - "test_size": len(test_ds), - "test_accuracy": final_acc, - "best_test_accuracy": best_acc, - "hyperparameters": dict(config), - }, -) -artifact.add_file(model_path) -logged = run.log_artifact(artifact, aliases=["latest"]) -# Block until the artifact has committed before linking, to avoid a race. -logged.wait() -mo.output.append(mo.md(f"**Artifact logged:** `{artifact.name}` (alias `latest`)")) - -# Link to the Registry, surfacing a remediation note instead of crashing. -if link_to_registry.value: - target_path = f"wandb-registry-{registry_name_v}/{collection_name_v}" - try: - run.link_artifact(artifact=logged, target_path=target_path) - mo.output.append( - mo.callout( - mo.md( - f"**Linked to Registry:** `{target_path}` — see " - f"[wandb.ai/registry](https://wandb.ai/registry)." - ), - kind="success", - ) - ) - except Exception as exc: # noqa: BLE001 - surface any failure to the reader - mo.output.append( - mo.callout( - mo.md( - f"**Registry link failed.** Target `{target_path}` — `{exc}`\n\n" - f"- Linking needs at least the **Member** role on the " - f"Registry. `view-only member cannot write to project` means " - f"your seat is view-only: the run and artifact succeed, but " - f"linking is blocked. An admin can grant access from the " - f"Registry **Members** settings, the Python SDK " - f"(`wandb.Api().registry(...)` then `add_member()` / " - f"`update_member()`), or SCIM (`PATCH /scim/Users/{{id}}` with " - f"`registryRoles`) — see " - f"https://docs.wandb.ai/guides/registry/configure_registry/. " - f"Or set **W&B entity** to a team in an org where you have " - f"Registry write access.\n" - f"- The Registry `{registry_name_v}` may not exist; an admin " - f"can create it from the W&B Registry UI.\n" - f"- On the legacy Model Registry, link with " - f"`target_path='model-registry/{collection_name_v}'` instead." - ), - kind="danger", - ) - ) -else: - mo.output.append( - mo.md( +```python {.marimo hide_code="true"} +def _registry_callout(status): + if status["kind"] == "disabled": + return mo.md( "_Registry linking is disabled — the artifact is logged to the run " "but not linked to a collection._" ) + if status["kind"] == "linked": + return mo.callout( + mo.md( + f"**Linked to Registry:** `{status['target_path']}` — see " + f"[wandb.ai/registry](https://wandb.ai/registry)." + ), + kind="success", + ) + return mo.callout( + mo.md( + f"**Registry link failed.** Target `{status['target_path']}` — " + f"`{status['error']}`\n\n" + f"- Linking needs at least the **Member** role on the " + f"Registry. `view-only member cannot write to project` means " + f"your seat is view-only: the run and artifact succeed, but " + f"linking is blocked. An admin can grant access from the " + f"Registry **Members** settings, the Python SDK " + f"(`wandb.Api().registry(...)` then `add_member()` / " + f"`update_member()`), or SCIM (`PATCH /scim/Users/{{id}}` with " + f"`registryRoles`) — see " + f"https://docs.wandb.ai/guides/registry/configure_registry/. " + f"Or set **W&B entity** to a team in an org where you have " + f"Registry write access.\n" + f"- The Registry `{registry_name_v}` may not exist; an admin " + f"can create it from the W&B Registry UI.\n" + f"- On the legacy Model Registry, link with " + f"`target_path='model-registry/{collection_name_v}'` instead." + ), + kind="danger", ) -# Close the run so its summary and any Registry link finalize server-side. -wandb.finish() -``` -```python {.marimo} -# Placed after the training cell on purpose: it's the explicit "run" trigger -# for the pipeline above. It must be its own cell because that cell reads -# `train_button.value`, and a widget only drives reactivity when a -# *different* cell consumes it. The gate also stops the pipeline from -# running automatically when the notebook opens; run_button's value is True -# only for the cascade a click triggers (then resets to False), so editing -# the form afterwards re-runs the training cell but it stops immediately. -train_button = mo.ui.run_button(label="Train model", kind="success") mo.vstack( [ - train_button, - mo.md( - "Runs the training cell above. It is gated so it does not " - "execute when the notebook opens — click to run, and click " - "again to retrain after editing the form (the previous run is " - "finished first)." - ), + mo.md(f"**Artifact logged:** `{artifact_name}` (alias `latest`)"), + _registry_callout(registry_status), ] ) ``` +## Evaluation + ```python {.marimo} -# Consume the model: download it from W&B (preferring the registered -# version, falling back to the run's own artifact), load the weights into a -# fresh network, and classify 10 held-out test digits. api = wandb.Api() try: consumed = api.artifact( - f"wandb-registry-{registry_name_v}/{collection_name_v}:latest", type="model" + f"wandb-registry-{registry_name_v}/{collection_name_v}:latest", + type="model", ) source = f"registry `wandb-registry-{registry_name_v}/{collection_name_v}:latest`" except Exception: # noqa: BLE001 - registry link may be absent (e.g. a view-only seat) consumed = api.artifact( - f"{run.entity}/{run.project}/mnist-cnn-{run.id}:latest", type="model" + f"{run.entity}/{run.project}/{artifact_name}:latest", + type="model", ) - source = f"run artifact `mnist-cnn-{run.id}:latest`" + source = f"run artifact `{artifact_name}:latest`" weights_dir = consumed.download() clf = Net() -clf.load_state_dict(torch.load(f"{weights_dir}/mnist_cnn.pt", map_location="cpu")) +clf.load_state_dict( + torch.load(f"{weights_dir}/mnist_cnn.pt", map_location="cpu") +) clf.eval() -cards = [] +rows = [] n_correct = 0 with torch.no_grad(): for i in range(10): @@ -453,58 +415,208 @@ with torch.no_grad(): n_correct += int(prediction == true_label) # Undo the Normalize transform so the digit renders as a clean image. digit = (image * 0.3081 + 0.1307).clamp(0, 1).squeeze().numpy() - mark = "✅" if prediction == true_label else "❌" - cards.append( - mo.vstack( - [ - mo.image(digit, width=64, vmin=0, vmax=1), - mo.md(f"{mark} **{prediction}** · true {true_label}"), - ], - align="center", - ) + rows.append( + { + "Image": mo.image(digit, width=56, vmin=0, vmax=1), + "Label": true_label, + "Prediction": prediction, + } ) mo.vstack( [ mo.md( - f"## Classify 10 test digits\n\nConsumed the model from {source}, " + f"**Classify 10 test digits.** Consumed the model from {source}, " f"loaded the weights into a fresh network, and ran it on 10 held-out " f"MNIST test images — **{n_correct}/10 correct**." ), - mo.hstack(cards, wrap=True, justify="start"), + mo.ui.table(rows, selection=None), ] ) ``` -````python {.marimo hide_code="true"} -# Renders only after a run exists (it consumes `run` from the training -# cell), so it appears once training finishes. -mo.md( - f""" - ## Verify and next steps - - 1. Open the run: [{run.name}]({run.url}) — check the **Charts**, - **System**, and **Examples** panels. - 2. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed - with its metadata (test accuracy, parameter count, hyperparameters). - 3. At [wandb.ai/registry](https://wandb.ai/registry), open the - **{registry_name_v.title()}** registry, then the **{collection_name_v}** - collection, and confirm the linked version. - - **Consume the registered model** from any script or notebook: - - ```python - import wandb - art = wandb.Api().artifact( - "wandb-registry-{registry_name_v}/{collection_name_v}:latest" - ) - art.download() # writes mnist_cnn.pt under ./artifacts/ - ``` +## Verify and next steps - **Next steps:** promote a version by adding the `production` alias from - the Registry UI; re-run with a deeper architecture or a different - learning rate and compare runs in the W&B UI; or add a W&B Automation to - trigger evaluation when a new version is linked. - """ +1. Open the run: [{run.name}]({run.url}) — check the **Charts**, + **System**, and **Examples** panels. +2. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed + with its metadata (test accuracy, parameter count, hyperparameters). +3. At [wandb.ai/registry](https://wandb.ai/registry), open the + **{registry_name_v.title()}** registry, then the **{collection_name_v}** + collection, and confirm the linked version. + +**Consume the registered model** from any script or notebook: + +```python +import wandb +art = wandb.Api().artifact( + "wandb-registry-{registry_name_v}/{collection_name_v}:latest" ) -```` \ No newline at end of file +art.download() # writes mnist_cnn.pt under ./artifacts/ +``` + +**Next steps:** promote a version by adding the `production` alias from +the Registry UI; re-run with a deeper architecture or a different +learning rate and compare runs in the W&B UI; or add a W&B Automation to +trigger evaluation when a new version is linked. + +## Helper functions + +```python {.marimo name="load_data"} +def load_data(): + """Download (or reuse cached) MNIST with the standard normalization.""" + transform = transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] + ) + train_ds = datasets.MNIST( + "./data", train=True, download=True, transform=transform + ) + test_ds = datasets.MNIST( + "./data", train=False, download=True, transform=transform + ) + return train_ds, test_ds +``` + +```python {.marimo name="make_loaders"} +def make_loaders(train_ds, test_ds, batch_size): + """Wrap the datasets in loaders, enabling CUDA niceties when available.""" + loader_kwargs = ( + {"num_workers": 2, "pin_memory": True} if device.type == "cuda" else {} + ) + train_loader = DataLoader( + train_ds, batch_size=batch_size, shuffle=True, **loader_kwargs + ) + test_loader = DataLoader( + test_ds, batch_size=1000, shuffle=False, **loader_kwargs + ) + return train_loader, test_loader +``` + +```python {.marimo name="train_one_epoch"} +def train_one_epoch(model, loader, optimizer, epoch, epochs): + """Run one training epoch, streaming train loss to W&B every 50 steps.""" + model.train() + for batch_idx, (data, target) in enumerate( + tqdm(loader, desc=f"epoch {epoch}/{epochs}") + ): + data, target = data.to(device), target.to(device) + optimizer.zero_grad() + output = model(data) + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % 50 == 0: + wandb.log({"train/loss": loss.item(), "epoch": epoch}) +``` + +```python {.marimo name="evaluate"} +def evaluate(model, loader, max_examples=16): + """Compute test loss/accuracy and collect a few labelled example images.""" + model.eval() + test_loss = 0.0 + correct = 0 + example_images = [] + with torch.no_grad(): + for data, target in loader: + data, target = data.to(device), target.to(device) + output = model(data) + test_loss += F.nll_loss(output, target, reduction="sum").item() + pred = output.argmax(dim=1, keepdim=True) + correct += pred.eq(target.view_as(pred)).sum().item() + # Pull up to `max_examples` predictions for the W&B Examples panel. + while len(example_images) < max_examples and len( + example_images + ) < data.size(0): + j = len(example_images) + example_images.append( + wandb.Image( + data[j], + caption=f"pred={pred[j].item()} true={target[j].item()}", + ) + ) + n = len(loader.dataset) + return test_loss / n, correct / n, example_images +``` + +```python {.marimo name="run_training"} +def run_training(run, config, train_ds, test_ds): + """Train the CNN, logging metrics each epoch; return the model and history.""" + train_loader, test_loader = make_loaders( + train_ds, test_ds, config["batch_size"] + ) + model = Net().to(device) + # `log="gradients"` is the standard choice; `log="all"` also logs parameter + # histograms at extra cost. + wandb.watch(model, log="gradients", log_freq=100) + optimizer = optim.SGD( + model.parameters(), lr=config["lr"], momentum=config["momentum"] + ) + + history = [] + best_acc = 0.0 + test_acc = 0.0 + for epoch in range(1, config["epochs"] + 1): + train_one_epoch(model, train_loader, optimizer, epoch, config["epochs"]) + test_loss, test_acc, example_images = evaluate(model, test_loader) + best_acc = max(best_acc, test_acc) + wandb.log( + { + "test/loss": test_loss, + "test/accuracy": test_acc, + "epoch": epoch, + "examples": example_images, + } + ) + history.append( + { + "epoch": epoch, + "test_loss": round(test_loss, 4), + "test_acc": round(test_acc, 4), + } + ) + # Full-precision last-epoch accuracy; `history` rounds only for display. + return model, history, test_acc, best_acc +``` + +```python {.marimo name="save_and_log_artifact"} +def save_and_log_artifact( + run, model, config, train_size, test_size, final_acc, best_acc, + model_path="mnist_cnn.pt", +): + """Persist the weights and log them as a `model` Artifact aliased `latest`.""" + torch.save(model.state_dict(), model_path) + name = f"mnist-cnn-{run.id}" + artifact = wandb.Artifact( + name=name, + type="model", + description=( + "Small CNN trained on MNIST. Architecture: 2 conv layers " + "(10 and 20 filters, 5x5 kernels) + 2 FC layers (50, 10)." + ), + metadata={ + "framework": "pytorch", + "architecture": "CNN", + "num_parameters": sum(p.numel() for p in model.parameters()), + "dataset": "MNIST", + "train_size": train_size, + "test_size": test_size, + "test_accuracy": final_acc, + "best_test_accuracy": best_acc, + "hyperparameters": dict(config), + }, + ) + artifact.add_file(model_path) + logged = run.log_artifact(artifact, aliases=["latest"]) + # Block until the artifact has committed before linking, to avoid a race. + logged.wait() + # Return the base name (no version) so callers can build `:latest`. + return logged, name +``` + +```python {.marimo name="link_artifact_to_registry"} +def link_artifact_to_registry(run, logged, registry_name, collection_name): + """Link the logged artifact into a Registry collection; return the target.""" + target_path = f"wandb-registry-{registry_name}/{collection_name}" + run.link_artifact(artifact=logged, target_path=target_path) + return target_path +``` \ No newline at end of file diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index cb78c945..9f3db016 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -22,53 +22,11 @@ import marimo -__generated_with = "0.23.9" +__generated_with = "0.23.11" app = marimo.App(width="medium", app_title="MNIST -> W&B Registry") - -@app.cell(hide_code=True) -def _(): +with app.setup(hide_code=True): import marimo as mo - - mo.md( - """ - # MNIST -> W&B Run -> Registry - - ## What you will build - - - A **W&B run** with training and test metrics, gradient histograms, - and example test-set predictions logged as images. - - A **model Artifact** named `mnist-cnn-` of type `model`, - carrying metadata (test accuracy, parameter count, hyperparameters). - - A version of that Artifact **linked into a W&B Registry collection** - so it appears under registered models org-wide. - - ## Prerequisites - - - Authenticate with W&B one of two ways: run **`wandb login`** in - your shell before starting marimo, or paste your key into the - **W&B API key** field in the form below. Get your key from - [wandb.ai/authorize](https://wandb.ai/authorize). - - A W&B **team** to write the run to, set in the **W&B entity** field. - Accounts created after May 2024 have no personal entity, so the run - must go to a team — your username will not work as an entity. - - A **W&B Registry** must exist in your org, and your account needs at - least the **Member** role on it (linking an artifact is a write - action). The built-in Model registry is provisioned automatically in - newer orgs. If linking fails (for example, from a view-only seat), - the run still completes and the Registry step explains how to fix it. - - A GPU is optional. The defaults finish in about 2 minutes on CPU. - """ - ) - return (mo,) - - -@app.cell -def _(mo): - # Imports, device detection, and the input form all live in one cell: this - # is "everything up to collecting your inputs". It defines the form widgets - # but never reads their `.value` — marimo only makes a widget reactive when - # a *different* cell consumes it, which the training cell below does. import torch import torch.nn as nn import torch.nn.functional as F @@ -95,189 +53,242 @@ def _(mo): ) device_kind = "warn" + +@app.cell(hide_code=True) +def _(): + mo.md(r""" + # MNIST -> W&B Run -> Registry + + ## What you will build + + - A **W&B run** with training and test metrics, gradient histograms, + and example test-set predictions logged as images. + - A **model Artifact** named `mnist-cnn-` of type `model`, + carrying metadata (test accuracy, parameter count, hyperparameters). + - A version of that Artifact **linked into a W&B Registry collection** + so it appears under registered models org-wide. + + ## Prerequisites + + - Authenticate with W&B one of two ways: run **`wandb login`** in + your shell before starting marimo, or paste your key into the + **W&B API key** field in the form below. Get your key from + [wandb.ai/authorize](https://wandb.ai/authorize). + - A W&B **team** to write the run to, set in the **W&B entity** field. + Accounts created after May 2024 have no personal entity, so the run + must go to a team — your username will not work as an entity. + - A **W&B Registry** must exist in your org, and your account needs at + least the **Member** role on it (linking an artifact is a write + action). The built-in Model registry is provisioned automatically in + newer orgs. If linking fails (for example, from a view-only seat), + the run still completes and the Registry step explains how to fix it. + - A GPU is optional. The defaults finish in about 2 minutes on CPU. + """) + return + + +@app.cell +def _(): + mo.outline() + return + + +@app.cell(hide_code=True) +def _(): + mo.callout( + mo.md(f"**Device:** `{device}`. {device_note}"), + kind=device_kind, + ) + return + + +@app.cell(hide_code=True) +def _(): + mo.md(r""" + ## Configuration + + Set the hyperparameters and W&B targets, then click **Train model** below. + """) + return + + +@app.cell(hide_code=True) +def _(): epochs = mo.ui.slider(start=1, stop=10, step=1, value=3, label="Epochs") batch_size = mo.ui.dropdown( options=["32", "64", "128", "256"], value="64", label="Batch size" ) lr = mo.ui.slider( - start=0.001, stop=0.1, step=0.001, value=0.01, label="Learning rate", show_value=True + start=0.001, + stop=0.1, + step=0.001, + value=0.01, + label="Learning rate", + show_value=True, ) momentum = mo.ui.slider( - start=0.0, stop=0.99, step=0.01, value=0.5, label="SGD momentum", show_value=True + start=0.0, + stop=0.99, + step=0.01, + value=0.5, + label="SGD momentum", + show_value=True, ) seed = mo.ui.number(start=0, stop=99999, value=42, label="Random seed") project = mo.ui.text(value="marimo-mnist-registry", label="W&B project") entity = mo.ui.text( - value="", label="W&B entity — a team you belong to (blank uses your default)" + value="", + label="W&B entity \u2014 a team you belong to (blank uses your default)", ) run_name = mo.ui.text(value="", label="Run name (blank auto-generates)") api_key = mo.ui.text( - value="", kind="password", label="W&B API key (blank uses your shell login)" + value="", + kind="password", + label="W&B API key (blank uses your shell login)", ) registry_name = mo.ui.text(value="model", label="W&B Registry name") - collection_name = mo.ui.text(value="MNIST Classifiers", label="Registry collection") - link_to_registry = mo.ui.checkbox(value=True, label="Link artifact to Registry") - - form = mo.vstack( - [ - mo.md("### Training"), - mo.hstack([epochs, batch_size]), - mo.hstack([lr, momentum]), - seed, - mo.md("### W&B run"), - api_key, - mo.hstack([project, entity, run_name]), - mo.md("### Registry"), - mo.hstack([registry_name, collection_name, link_to_registry]), - ] + collection_name = mo.ui.text( + value="MNIST Classifiers", label="Registry collection" ) - - mo.vstack( - [ - mo.callout( - mo.md(f"**Device:** `{device}` — {device_note}"), kind=device_kind - ), - mo.md( - "## Configure\n\nSet the hyperparameters and W&B targets, then click " - "**Train model** below. Changing a value here never starts a run on " - "its own — only the button does." - ), - form, - ] + link_to_registry = mo.ui.checkbox( + value=True, label="Link artifact to Registry" ) - return ( - DataLoader, - F, - api_key, - batch_size, - collection_name, - datasets, - device, - entity, - epochs, - link_to_registry, - lr, - momentum, - nn, - optim, - project, - registry_name, - run_name, - seed, - torch, - tqdm, - transforms, - wandb, + + # Batch every control into one form so training only kicks off on submit. + # `form.value` is None until the user clicks Train model, then becomes a dict + # keyed by the names below \u2014 the training cell gates on that. + form = ( + mo.md( + """ + **Training.** + + {epochs} {batch_size} + + {lr} {momentum} + + {seed} + + **W&B run.** + + {api_key} + + {project} + + {entity} + + {run_name} + + **Registry.** + + {registry_name} {collection_name} {link_to_registry} + """ + ) + .batch( + epochs=epochs, + batch_size=batch_size, + lr=lr, + momentum=momentum, + seed=seed, + api_key=api_key, + project=project, + entity=entity, + run_name=run_name, + registry_name=registry_name, + collection_name=collection_name, + link_to_registry=link_to_registry, + ) + .form(submit_button_label="Train model", bordered=False) ) + form + return (form,) -@app.cell -def _(F, nn): - class Net(nn.Module): - """Small CNN: 2 conv layers (10, 20 filters, 5x5) + 2 FC (50, 10). - - Defined in its own cell so the training cell and the consume cell can - share it (marimo forbids defining the same name in two cells). - """ - - def __init__(self): - super().__init__() - self.conv1 = nn.Conv2d(1, 10, kernel_size=5) - self.conv2 = nn.Conv2d(10, 20, kernel_size=5) - self.conv2_drop = nn.Dropout2d() - self.fc1 = nn.Linear(320, 50) - self.fc2 = nn.Linear(50, 10) - - def forward(self, x): - x = F.relu(F.max_pool2d(self.conv1(x), 2)) - x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) - x = x.view(-1, 320) - x = F.relu(self.fc1(x)) - x = F.dropout(x, training=self.training) - x = self.fc2(x) - return F.log_softmax(x, dim=1) - - return (Net,) + +@app.class_definition +class Net(nn.Module): + """Small CNN: 2 conv layers (10, 20 filters, 5x5) + 2 FC (50, 10). + + Defined in its own cell so the training cell and the consume cell can + share it (marimo forbids defining the same name in two cells). + """ + + def __init__(self): + super().__init__() + self.conv1 = nn.Conv2d(1, 10, kernel_size=5) + self.conv2 = nn.Conv2d(10, 20, kernel_size=5) + self.conv2_drop = nn.Dropout2d() + self.fc1 = nn.Linear(320, 50) + self.fc2 = nn.Linear(50, 10) + + def forward(self, x): + x = F.relu(F.max_pool2d(self.conv1(x), 2)) + x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) + x = x.view(-1, 320) + x = F.relu(self.fc1(x)) + x = F.dropout(x, training=self.training) + x = self.fc2(x) + return F.log_softmax(x, dim=1) + + +@app.cell(hide_code=True) +def _(): + mo.md(""" + ## Training + """) + return @app.cell -def _( - DataLoader, - F, - Net, - api_key, - batch_size, - collection_name, - datasets, - device, - entity, - epochs, - link_to_registry, - lr, - momentum, - mo, - optim, - project, - registry_name, - run_name, - seed, - torch, - tqdm, - train_button, - transforms, - wandb, -): - # Everything the Train button triggers, in one cell — no reason to make you - # advance through a chain of output-less code blocks. Each milestone is - # streamed to the cell output with `mo.output.append` as it happens. +def _(form): mo.stop( - not train_button.value, + form.value is None, mo.md( - "This cell runs the whole pipeline — start the run, train, log " - "metrics and example predictions, save the model Artifact, and link " - "it to the Registry. Click **Train model** below to run it." + "Training hasn't started yet. Fill in the form above and click " + "**Train model** to start the run — it trains the model, logs metrics " + "and example predictions, saves the weights as an Artifact, and links " + "them to the Registry." ), ) + cfg = form.value config = { - "epochs": epochs.value, - "batch_size": int(batch_size.value), - "lr": lr.value, - "momentum": momentum.value, - "seed": seed.value, + "epochs": cfg["epochs"], + "batch_size": int(cfg["batch_size"]), + "lr": cfg["lr"], + "momentum": cfg["momentum"], + "seed": cfg["seed"], "architecture": "CNN", "dataset": "MNIST", } - registry_name_v = registry_name.value.strip() - collection_name_v = collection_name.value.strip() + registry_name_v = cfg["registry_name"].strip() + collection_name_v = cfg["collection_name"].strip() - # Authenticate. Finish any prior run first (marimo keeps the kernel alive - # across re-clicks). A key pasted into the form wins; otherwise fall back to - # ambient login (shell `wandb login`, WANDB_API_KEY, or netrc). The key is - # never written to the run config. + # Authenticate and start the run. Finish any prior run first (marimo keeps the + # kernel alive across re-submits). A key pasted into the form wins; otherwise + # fall back to ambient login (shell `wandb login`, WANDB_API_KEY, or netrc). + # The key is never written to the run config. if wandb.run is not None: wandb.finish() - if api_key.value: - wandb.login(key=api_key.value) - + if cfg["api_key"]: + wandb.login(key=cfg["api_key"]) torch.manual_seed(config["seed"]) try: run = wandb.init( - project=project.value or None, - entity=entity.value or None, - name=run_name.value or None, + project=cfg["project"] or None, + entity=cfg["entity"] or None, + name=cfg["run_name"] or None, config=config, job_type="train", ) - except Exception as exc: # noqa: BLE001 - turn the raw traceback into guidance + except Exception as init_exc: # noqa: BLE001 - turn the raw traceback into guidance mo.stop( True, mo.callout( mo.md( - f"**Could not start the run.** `{exc}`\n\n" + f"**Could not start the run.** `{init_exc}`\n\n" f"An `entity ... not found` error means the **W&B entity** is " f"not a team you can write to. Personal-username entities were " f"removed for accounts created after 21 May 2024, so set the " @@ -287,224 +298,166 @@ def _( kind="danger", ), ) + # Use `epoch` as the x-axis for train/test metrics in the W&B UI. wandb.define_metric("epoch") wandb.define_metric("train/*", step_metric="epoch") wandb.define_metric("test/*", step_metric="epoch") + return cfg, collection_name_v, config, registry_name_v, run + + +@app.cell(hide_code=True) +def _(run): # Surface the run link right away so you can watch metrics stream live. - mo.output.append(mo.md(f"**Run started:** [`{run.name}`]({run.url})")) + mo.md(f"**Run started:** [`{run.name}`]({run.url})") + return - model = Net().to(device) - # `log="gradients"` is the standard choice; `log="all"` also logs parameter - # histograms at extra cost. - wandb.watch(model, log="gradients", log_freq=100) - optimizer = optim.SGD( - model.parameters(), lr=config["lr"], momentum=config["momentum"] - ) - transform = transforms.Compose( - [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] - ) - train_ds = datasets.MNIST("./data", train=True, download=True, transform=transform) - test_ds = datasets.MNIST("./data", train=False, download=True, transform=transform) - loader_kwargs = ( - {"num_workers": 2, "pin_memory": True} if device.type == "cuda" else {} +@app.cell +def _(config, run): + train_ds, test_ds = load_data() + model, history, final_acc, best_acc = run_training( + run, config, train_ds, test_ds ) - train_loader = DataLoader( - train_ds, batch_size=config["batch_size"], shuffle=True, **loader_kwargs + + mo.vstack( + [ + mo.md("### Training summary"), + mo.ui.table(history, selection=None), + mo.md(f"**Final test accuracy:** {final_acc:.2%}"), + ] ) - test_loader = DataLoader(test_ds, batch_size=1000, shuffle=False, **loader_kwargs) + return best_acc, final_acc, model, test_ds, train_ds - history = [] - best_acc = 0.0 - for epoch in range(1, config["epochs"] + 1): - model.train() - for batch_idx, (data, target) in enumerate( - tqdm(train_loader, desc=f"epoch {epoch}/{config['epochs']}") - ): - data, target = data.to(device), target.to(device) - optimizer.zero_grad() - output = model(data) - loss = F.nll_loss(output, target) - loss.backward() - optimizer.step() - if batch_idx % 50 == 0: - wandb.log({"train/loss": loss.item(), "epoch": epoch}) - - model.eval() - test_loss = 0.0 - correct = 0 - example_images = [] - with torch.no_grad(): - for data, target in test_loader: - data, target = data.to(device), target.to(device) - output = model(data) - test_loss += F.nll_loss(output, target, reduction="sum").item() - pred = output.argmax(dim=1, keepdim=True) - correct += pred.eq(target.view_as(pred)).sum().item() - # Pull up to 16 example predictions from the first batch. - while len(example_images) < 16 and len(example_images) < data.size(0): - j = len(example_images) - example_images.append( - wandb.Image( - data[j], - caption=f"pred={pred[j].item()} true={target[j].item()}", - ) - ) - test_loss /= len(test_loader.dataset) - test_acc = correct / len(test_loader.dataset) - best_acc = max(best_acc, test_acc) - wandb.log( - { - "test/loss": test_loss, - "test/accuracy": test_acc, - "epoch": epoch, - "examples": example_images, +@app.cell +def _( + best_acc, + cfg, + collection_name_v, + config, + final_acc, + model, + registry_name_v, + run, + test_ds, + train_ds, +): + logged, artifact_name = save_and_log_artifact( + run, + model, + config, + train_size=len(train_ds), + test_size=len(test_ds), + final_acc=final_acc, + best_acc=best_acc, + ) + + # Link to the Registry unless disabled, capturing the outcome for display + # rather than crashing the pipeline. + if not cfg["link_to_registry"]: + registry_status = {"kind": "disabled"} + else: + try: + registry_status = { + "kind": "linked", + "target_path": link_artifact_to_registry( + run, logged, registry_name_v, collection_name_v + ), + } + except Exception as link_exc: # noqa: BLE001 - surface any failure to the reader + registry_status = { + "kind": "failed", + "target_path": f"wandb-registry-{registry_name_v}/{collection_name_v}", + "error": str(link_exc), } - ) - history.append( - {"epoch": epoch, "test_loss": round(test_loss, 4), "test_acc": round(test_acc, 4)} - ) - # Full-precision last-epoch accuracy; `history` rounds only for display. - final_acc = test_acc - mo.output.append( - mo.vstack( - [ - mo.md("### Training summary"), - mo.ui.table(history, selection=None), - mo.md(f"**Final test accuracy:** {final_acc:.2%}"), - ] - ) - ) + # Close the run so its summary and any Registry link finalize server-side. + wandb.finish() + return artifact_name, registry_status - # Save the weights and log them as a model Artifact tagged `latest`. - model_path = "mnist_cnn.pt" - torch.save(model.state_dict(), model_path) - artifact = wandb.Artifact( - name=f"mnist-cnn-{run.id}", - type="model", - description=( - "Small CNN trained on MNIST. Architecture: 2 conv layers " - "(10 and 20 filters, 5x5 kernels) + 2 FC layers (50, 10)." - ), - metadata={ - "framework": "pytorch", - "architecture": "CNN", - "num_parameters": sum(p.numel() for p in model.parameters()), - "dataset": "MNIST", - "train_size": len(train_ds), - "test_size": len(test_ds), - "test_accuracy": final_acc, - "best_test_accuracy": best_acc, - "hyperparameters": dict(config), - }, - ) - artifact.add_file(model_path) - logged = run.log_artifact(artifact, aliases=["latest"]) - # Block until the artifact has committed before linking, to avoid a race. - logged.wait() - mo.output.append(mo.md(f"**Artifact logged:** `{artifact.name}` (alias `latest`)")) - # Link to the Registry, surfacing a remediation note instead of crashing. - if link_to_registry.value: - target_path = f"wandb-registry-{registry_name_v}/{collection_name_v}" - try: - run.link_artifact(artifact=logged, target_path=target_path) - mo.output.append( - mo.callout( - mo.md( - f"**Linked to Registry:** `{target_path}` — see " - f"[wandb.ai/registry](https://wandb.ai/registry)." - ), - kind="success", - ) - ) - except Exception as exc: # noqa: BLE001 - surface any failure to the reader - mo.output.append( - mo.callout( - mo.md( - f"**Registry link failed.** Target `{target_path}` — `{exc}`\n\n" - f"- Linking needs at least the **Member** role on the " - f"Registry. `view-only member cannot write to project` means " - f"your seat is view-only: the run and artifact succeed, but " - f"linking is blocked. An admin can grant access from the " - f"Registry **Members** settings, the Python SDK " - f"(`wandb.Api().registry(...)` then `add_member()` / " - f"`update_member()`), or SCIM (`PATCH /scim/Users/{{id}}` with " - f"`registryRoles`) — see " - f"https://docs.wandb.ai/guides/registry/configure_registry/. " - f"Or set **W&B entity** to a team in an org where you have " - f"Registry write access.\n" - f"- The Registry `{registry_name_v}` may not exist; an admin " - f"can create it from the W&B Registry UI.\n" - f"- On the legacy Model Registry, link with " - f"`target_path='model-registry/{collection_name_v}'` instead." - ), - kind="danger", - ) - ) - else: - mo.output.append( - mo.md( +@app.cell(hide_code=True) +def _(artifact_name, collection_name_v, registry_name_v, registry_status): + def _registry_callout(status): + if status["kind"] == "disabled": + return mo.md( "_Registry linking is disabled — the artifact is logged to the run " "but not linked to a collection._" ) + if status["kind"] == "linked": + return mo.callout( + mo.md( + f"**Linked to Registry:** `{status['target_path']}` — see " + f"[wandb.ai/registry](https://wandb.ai/registry)." + ), + kind="success", + ) + return mo.callout( + mo.md( + f"**Registry link failed.** Target `{status['target_path']}` — " + f"`{status['error']}`\n\n" + f"- Linking needs at least the **Member** role on the " + f"Registry. `view-only member cannot write to project` means " + f"your seat is view-only: the run and artifact succeed, but " + f"linking is blocked. An admin can grant access from the " + f"Registry **Members** settings, the Python SDK " + f"(`wandb.Api().registry(...)` then `add_member()` / " + f"`update_member()`), or SCIM (`PATCH /scim/Users/{{id}}` with " + f"`registryRoles`) — see " + f"https://docs.wandb.ai/guides/registry/configure_registry/. " + f"Or set **W&B entity** to a team in an org where you have " + f"Registry write access.\n" + f"- The Registry `{registry_name_v}` may not exist; an admin " + f"can create it from the W&B Registry UI.\n" + f"- On the legacy Model Registry, link with " + f"`target_path='model-registry/{collection_name_v}'` instead." + ), + kind="danger", ) - # Close the run so its summary and any Registry link finalize server-side. - wandb.finish() - return collection_name_v, registry_name_v, run, test_ds - -@app.cell -def _(mo): - # Placed after the training cell on purpose: it's the explicit "run" trigger - # for the pipeline above. It must be its own cell because that cell reads - # `train_button.value`, and a widget only drives reactivity when a - # *different* cell consumes it. The gate also stops the pipeline from - # running automatically when the notebook opens; run_button's value is True - # only for the cascade a click triggers (then resets to False), so editing - # the form afterwards re-runs the training cell but it stops immediately. - train_button = mo.ui.run_button(label="Train model", kind="success") mo.vstack( [ - train_button, - mo.md( - "Runs the training cell above. It is gated so it does not " - "execute when the notebook opens — click to run, and click " - "again to retrain after editing the form (the previous run is " - "finished first)." - ), + mo.md(f"**Artifact logged:** `{artifact_name}` (alias `latest`)"), + _registry_callout(registry_status), ] ) - return (train_button,) + return + + +@app.cell(hide_code=True) +def _(): + mo.md(""" + ## Evaluation + """) + return @app.cell -def _(Net, collection_name_v, mo, registry_name_v, run, test_ds, torch, wandb): - # Consume the model: download it from W&B (preferring the registered - # version, falling back to the run's own artifact), load the weights into a - # fresh network, and classify 10 held-out test digits. +def _(artifact_name, collection_name_v, registry_name_v, run, test_ds): api = wandb.Api() try: consumed = api.artifact( - f"wandb-registry-{registry_name_v}/{collection_name_v}:latest", type="model" + f"wandb-registry-{registry_name_v}/{collection_name_v}:latest", + type="model", ) source = f"registry `wandb-registry-{registry_name_v}/{collection_name_v}:latest`" except Exception: # noqa: BLE001 - registry link may be absent (e.g. a view-only seat) consumed = api.artifact( - f"{run.entity}/{run.project}/mnist-cnn-{run.id}:latest", type="model" + f"{run.entity}/{run.project}/{artifact_name}:latest", + type="model", ) - source = f"run artifact `mnist-cnn-{run.id}:latest`" + source = f"run artifact `{artifact_name}:latest`" weights_dir = consumed.download() clf = Net() - clf.load_state_dict(torch.load(f"{weights_dir}/mnist_cnn.pt", map_location="cpu")) + clf.load_state_dict( + torch.load(f"{weights_dir}/mnist_cnn.pt", map_location="cpu") + ) clf.eval() - cards = [] + rows = [] n_correct = 0 with torch.no_grad(): for i in range(10): @@ -513,64 +466,224 @@ def _(Net, collection_name_v, mo, registry_name_v, run, test_ds, torch, wandb): n_correct += int(prediction == true_label) # Undo the Normalize transform so the digit renders as a clean image. digit = (image * 0.3081 + 0.1307).clamp(0, 1).squeeze().numpy() - mark = "✅" if prediction == true_label else "❌" - cards.append( - mo.vstack( - [ - mo.image(digit, width=64, vmin=0, vmax=1), - mo.md(f"{mark} **{prediction}** · true {true_label}"), - ], - align="center", - ) + rows.append( + { + "Image": mo.image(digit, width=56, vmin=0, vmax=1), + "Label": true_label, + "Prediction": prediction, + } ) mo.vstack( [ mo.md( - f"## Classify 10 test digits\n\nConsumed the model from {source}, " + f"**Classify 10 test digits.** Consumed the model from {source}, " f"loaded the weights into a fresh network, and ran it on 10 held-out " f"MNIST test images — **{n_correct}/10 correct**." ), - mo.hstack(cards, wrap=True, justify="start"), + mo.ui.table(rows, selection=None), ] ) return @app.cell(hide_code=True) -def _(collection_name_v, mo, registry_name_v, run): - # Renders only after a run exists (it consumes `run` from the training - # cell), so it appears once training finishes. - mo.md( - f""" - ## Verify and next steps - - 1. Open the run: [{run.name}]({run.url}) — check the **Charts**, - **System**, and **Examples** panels. - 2. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed - with its metadata (test accuracy, parameter count, hyperparameters). - 3. At [wandb.ai/registry](https://wandb.ai/registry), open the - **{registry_name_v.title()}** registry, then the **{collection_name_v}** - collection, and confirm the linked version. - - **Consume the registered model** from any script or notebook: - - ```python - import wandb - art = wandb.Api().artifact( - "wandb-registry-{registry_name_v}/{collection_name_v}:latest" - ) - art.download() # writes mnist_cnn.pt under ./artifacts/ - ``` - - **Next steps:** promote a version by adding the `production` alias from - the Registry UI; re-run with a deeper architecture or a different - learning rate and compare runs in the W&B UI; or add a W&B Automation to - trigger evaluation when a new version is linked. - """ +def _(): + mo.md(r""" + ## Verify and next steps + + 1. Open the run: [{run.name}]({run.url}) — check the **Charts**, + **System**, and **Examples** panels. + 2. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed + with its metadata (test accuracy, parameter count, hyperparameters). + 3. At [wandb.ai/registry](https://wandb.ai/registry), open the + **{registry_name_v.title()}** registry, then the **{collection_name_v}** + collection, and confirm the linked version. + + **Consume the registered model** from any script or notebook: + + ```python + import wandb + art = wandb.Api().artifact( + "wandb-registry-{registry_name_v}/{collection_name_v}:latest" ) + art.download() # writes mnist_cnn.pt under ./artifacts/ + ``` + + **Next steps:** promote a version by adding the `production` alias from + the Registry UI; re-run with a deeper architecture or a different + learning rate and compare runs in the W&B UI; or add a W&B Automation to + trigger evaluation when a new version is linked. + """) return +@app.cell(hide_code=True) +def _(): + mo.md(r""" + ## Helper functions + """) + return + + +@app.function +def load_data(): + """Download (or reuse cached) MNIST with the standard normalization.""" + transform = transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] + ) + train_ds = datasets.MNIST( + "./data", train=True, download=True, transform=transform + ) + test_ds = datasets.MNIST( + "./data", train=False, download=True, transform=transform + ) + return train_ds, test_ds + + +@app.function +def make_loaders(train_ds, test_ds, batch_size): + """Wrap the datasets in loaders, enabling CUDA niceties when available.""" + loader_kwargs = ( + {"num_workers": 2, "pin_memory": True} if device.type == "cuda" else {} + ) + train_loader = DataLoader( + train_ds, batch_size=batch_size, shuffle=True, **loader_kwargs + ) + test_loader = DataLoader( + test_ds, batch_size=1000, shuffle=False, **loader_kwargs + ) + return train_loader, test_loader + + +@app.function +def train_one_epoch(model, loader, optimizer, epoch, epochs): + """Run one training epoch, streaming train loss to W&B every 50 steps.""" + model.train() + for batch_idx, (data, target) in enumerate( + tqdm(loader, desc=f"epoch {epoch}/{epochs}") + ): + data, target = data.to(device), target.to(device) + optimizer.zero_grad() + output = model(data) + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % 50 == 0: + wandb.log({"train/loss": loss.item(), "epoch": epoch}) + + +@app.function +def evaluate(model, loader, max_examples=16): + """Compute test loss/accuracy and collect a few labelled example images.""" + model.eval() + test_loss = 0.0 + correct = 0 + example_images = [] + with torch.no_grad(): + for data, target in loader: + data, target = data.to(device), target.to(device) + output = model(data) + test_loss += F.nll_loss(output, target, reduction="sum").item() + pred = output.argmax(dim=1, keepdim=True) + correct += pred.eq(target.view_as(pred)).sum().item() + # Pull up to `max_examples` predictions for the W&B Examples panel. + while len(example_images) < max_examples and len( + example_images + ) < data.size(0): + j = len(example_images) + example_images.append( + wandb.Image( + data[j], + caption=f"pred={pred[j].item()} true={target[j].item()}", + ) + ) + n = len(loader.dataset) + return test_loss / n, correct / n, example_images + + +@app.function +def run_training(run, config, train_ds, test_ds): + """Train the CNN, logging metrics each epoch; return the model and history.""" + train_loader, test_loader = make_loaders( + train_ds, test_ds, config["batch_size"] + ) + model = Net().to(device) + # `log="gradients"` is the standard choice; `log="all"` also logs parameter + # histograms at extra cost. + wandb.watch(model, log="gradients", log_freq=100) + optimizer = optim.SGD( + model.parameters(), lr=config["lr"], momentum=config["momentum"] + ) + + history = [] + best_acc = 0.0 + test_acc = 0.0 + for epoch in range(1, config["epochs"] + 1): + train_one_epoch(model, train_loader, optimizer, epoch, config["epochs"]) + test_loss, test_acc, example_images = evaluate(model, test_loader) + best_acc = max(best_acc, test_acc) + wandb.log( + { + "test/loss": test_loss, + "test/accuracy": test_acc, + "epoch": epoch, + "examples": example_images, + } + ) + history.append( + { + "epoch": epoch, + "test_loss": round(test_loss, 4), + "test_acc": round(test_acc, 4), + } + ) + # Full-precision last-epoch accuracy; `history` rounds only for display. + return model, history, test_acc, best_acc + + +@app.function +def save_and_log_artifact( + run, model, config, train_size, test_size, final_acc, best_acc, + model_path="mnist_cnn.pt", +): + """Persist the weights and log them as a `model` Artifact aliased `latest`.""" + torch.save(model.state_dict(), model_path) + name = f"mnist-cnn-{run.id}" + artifact = wandb.Artifact( + name=name, + type="model", + description=( + "Small CNN trained on MNIST. Architecture: 2 conv layers " + "(10 and 20 filters, 5x5 kernels) + 2 FC layers (50, 10)." + ), + metadata={ + "framework": "pytorch", + "architecture": "CNN", + "num_parameters": sum(p.numel() for p in model.parameters()), + "dataset": "MNIST", + "train_size": train_size, + "test_size": test_size, + "test_accuracy": final_acc, + "best_test_accuracy": best_acc, + "hyperparameters": dict(config), + }, + ) + artifact.add_file(model_path) + logged = run.log_artifact(artifact, aliases=["latest"]) + # Block until the artifact has committed before linking, to avoid a race. + logged.wait() + # Return the base name (no version) so callers can build `:latest`. + return logged, name + + +@app.function +def link_artifact_to_registry(run, logged, registry_name, collection_name): + """Link the logged artifact into a Registry collection; return the target.""" + target_path = f"wandb-registry-{registry_name}/{collection_name}" + run.link_artifact(artifact=logged, target_path=target_path) + return target_path + + if __name__ == "__main__": app.run() From 6d209433ff02e033cbb7dd0b225364650b4a967f Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Thu, 16 Jul 2026 16:52:15 -0400 Subject: [PATCH 19/22] consolidate W&B logging into a single "Training" section W&B groups panels by the prefix before "/", so train/* and test/* produced separate sections, and the example images were logged from normalized tensors (values outside [0,1]) so they rendered black. Log train loss and test accuracy under a single "Training/" prefix, drop the redundant test/loss, the bare "epoch" metric, the broken "examples" panel, and wandb.watch (which added a separate gradients section). Result: one "Training" section with a loss chart and an accuracy chart. Also restore the verify cell's f-string interpolation (it had regressed to literal {run.name}) and update its panel references. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../marimo/mnist-registry/mnist_registry.md | 61 ++++++------------- .../marimo/mnist-registry/mnist_registry.py | 58 ++++++------------ 2 files changed, 37 insertions(+), 82 deletions(-) diff --git a/examples/marimo/mnist-registry/mnist_registry.md b/examples/marimo/mnist-registry/mnist_registry.md index f1c651dd..d8cebf57 100644 --- a/examples/marimo/mnist-registry/mnist_registry.md +++ b/examples/marimo/mnist-registry/mnist_registry.md @@ -59,8 +59,8 @@ else: ## What you will build -- A **W&B run** with training and test metrics, gradient histograms, - and example test-set predictions logged as images. +- A **W&B run** with a single **Training** section charting loss and + accuracy. - A **model Artifact** named `mnist-cnn-` of type `model`, carrying metadata (test accuracy, parameter count, hyperparameters). - A version of that Artifact **linked into a W&B Registry collection** @@ -222,9 +222,9 @@ mo.stop( form.value is None, mo.md( "Training hasn't started yet. Fill in the form above and click " - "**Train model** to start the run — it trains the model, logs metrics " - "and example predictions, saves the weights as an Artifact, and links " - "them to the Registry." + "**Train model** to start the run — it trains the model, logs loss " + "and accuracy, saves the weights as an Artifact, links them to the " + "Registry, and classifies a few test digits." ), ) @@ -274,11 +274,6 @@ except Exception as init_exc: # noqa: BLE001 - turn the raw traceback into guid kind="danger", ), ) - -# Use `epoch` as the x-axis for train/test metrics in the W&B UI. -wandb.define_metric("epoch") -wandb.define_metric("train/*", step_metric="epoch") -wandb.define_metric("test/*", step_metric="epoch") ``` ```python {.marimo hide_code="true"} @@ -435,10 +430,12 @@ mo.vstack( ) ``` +````python {.marimo hide_code="true"} +mo.md(f""" ## Verify and next steps -1. Open the run: [{run.name}]({run.url}) — check the **Charts**, - **System**, and **Examples** panels. +1. Open the run: [{run.name}]({run.url}) — check the **Training** charts + and the **System** metrics. 2. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed with its metadata (test accuracy, parameter count, hyperparameters). 3. At [wandb.ai/registry](https://wandb.ai/registry), open the @@ -459,7 +456,9 @@ art.download() # writes mnist_cnn.pt under ./artifacts/ the Registry UI; re-run with a deeper architecture or a different learning rate and compare runs in the W&B UI; or add a W&B Automation to trigger evaluation when a new version is linked. - +""") +```` + ## Helper functions ```python {.marimo name="load_data"} @@ -506,16 +505,15 @@ def train_one_epoch(model, loader, optimizer, epoch, epochs): loss.backward() optimizer.step() if batch_idx % 50 == 0: - wandb.log({"train/loss": loss.item(), "epoch": epoch}) + wandb.log({"Training/loss": loss.item()}) ``` ```python {.marimo name="evaluate"} -def evaluate(model, loader, max_examples=16): - """Compute test loss/accuracy and collect a few labelled example images.""" +def evaluate(model, loader): + """Compute test loss and accuracy over a data loader.""" model.eval() test_loss = 0.0 correct = 0 - example_images = [] with torch.no_grad(): for data, target in loader: data, target = data.to(device), target.to(device) @@ -523,19 +521,8 @@ def evaluate(model, loader, max_examples=16): test_loss += F.nll_loss(output, target, reduction="sum").item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() - # Pull up to `max_examples` predictions for the W&B Examples panel. - while len(example_images) < max_examples and len( - example_images - ) < data.size(0): - j = len(example_images) - example_images.append( - wandb.Image( - data[j], - caption=f"pred={pred[j].item()} true={target[j].item()}", - ) - ) n = len(loader.dataset) - return test_loss / n, correct / n, example_images + return test_loss / n, correct / n ``` ```python {.marimo name="run_training"} @@ -545,9 +532,6 @@ def run_training(run, config, train_ds, test_ds): train_ds, test_ds, config["batch_size"] ) model = Net().to(device) - # `log="gradients"` is the standard choice; `log="all"` also logs parameter - # histograms at extra cost. - wandb.watch(model, log="gradients", log_freq=100) optimizer = optim.SGD( model.parameters(), lr=config["lr"], momentum=config["momentum"] ) @@ -557,16 +541,11 @@ def run_training(run, config, train_ds, test_ds): test_acc = 0.0 for epoch in range(1, config["epochs"] + 1): train_one_epoch(model, train_loader, optimizer, epoch, config["epochs"]) - test_loss, test_acc, example_images = evaluate(model, test_loader) + test_loss, test_acc = evaluate(model, test_loader) best_acc = max(best_acc, test_acc) - wandb.log( - { - "test/loss": test_loss, - "test/accuracy": test_acc, - "epoch": epoch, - "examples": example_images, - } - ) + # `train_one_epoch` logs `Training/loss`; logging `Training/accuracy` + # here keeps both charts in a single "Training" section. + wandb.log({"Training/accuracy": test_acc}) history.append( { "epoch": epoch, diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index 9f3db016..dd3e6cd5 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -61,8 +61,8 @@ def _(): ## What you will build - - A **W&B run** with training and test metrics, gradient histograms, - and example test-set predictions logged as images. + - A **W&B run** with a single **Training** section charting loss and + accuracy. - A **model Artifact** named `mnist-cnn-` of type `model`, carrying metadata (test accuracy, parameter count, hyperparameters). - A version of that Artifact **linked into a W&B Registry collection** @@ -246,9 +246,9 @@ def _(form): form.value is None, mo.md( "Training hasn't started yet. Fill in the form above and click " - "**Train model** to start the run — it trains the model, logs metrics " - "and example predictions, saves the weights as an Artifact, and links " - "them to the Registry." + "**Train model** to start the run — it trains the model, logs loss " + "and accuracy, saves the weights as an Artifact, links them to the " + "Registry, and classifies a few test digits." ), ) @@ -299,10 +299,6 @@ def _(form): ), ) - # Use `epoch` as the x-axis for train/test metrics in the W&B UI. - wandb.define_metric("epoch") - wandb.define_metric("train/*", step_metric="epoch") - wandb.define_metric("test/*", step_metric="epoch") return cfg, collection_name_v, config, registry_name_v, run @@ -488,12 +484,12 @@ def _(artifact_name, collection_name_v, registry_name_v, run, test_ds): @app.cell(hide_code=True) -def _(): - mo.md(r""" +def _(collection_name_v, registry_name_v, run): + mo.md(f""" ## Verify and next steps - 1. Open the run: [{run.name}]({run.url}) — check the **Charts**, - **System**, and **Examples** panels. + 1. Open the run: [{run.name}]({run.url}) — check the **Training** charts + and the **System** metrics. 2. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed with its metadata (test accuracy, parameter count, hyperparameters). 3. At [wandb.ai/registry](https://wandb.ai/registry), open the @@ -570,16 +566,15 @@ def train_one_epoch(model, loader, optimizer, epoch, epochs): loss.backward() optimizer.step() if batch_idx % 50 == 0: - wandb.log({"train/loss": loss.item(), "epoch": epoch}) + wandb.log({"Training/loss": loss.item()}) @app.function -def evaluate(model, loader, max_examples=16): - """Compute test loss/accuracy and collect a few labelled example images.""" +def evaluate(model, loader): + """Compute test loss and accuracy over a data loader.""" model.eval() test_loss = 0.0 correct = 0 - example_images = [] with torch.no_grad(): for data, target in loader: data, target = data.to(device), target.to(device) @@ -587,19 +582,8 @@ def evaluate(model, loader, max_examples=16): test_loss += F.nll_loss(output, target, reduction="sum").item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() - # Pull up to `max_examples` predictions for the W&B Examples panel. - while len(example_images) < max_examples and len( - example_images - ) < data.size(0): - j = len(example_images) - example_images.append( - wandb.Image( - data[j], - caption=f"pred={pred[j].item()} true={target[j].item()}", - ) - ) n = len(loader.dataset) - return test_loss / n, correct / n, example_images + return test_loss / n, correct / n @app.function @@ -609,9 +593,6 @@ def run_training(run, config, train_ds, test_ds): train_ds, test_ds, config["batch_size"] ) model = Net().to(device) - # `log="gradients"` is the standard choice; `log="all"` also logs parameter - # histograms at extra cost. - wandb.watch(model, log="gradients", log_freq=100) optimizer = optim.SGD( model.parameters(), lr=config["lr"], momentum=config["momentum"] ) @@ -621,16 +602,11 @@ def run_training(run, config, train_ds, test_ds): test_acc = 0.0 for epoch in range(1, config["epochs"] + 1): train_one_epoch(model, train_loader, optimizer, epoch, config["epochs"]) - test_loss, test_acc, example_images = evaluate(model, test_loader) + test_loss, test_acc = evaluate(model, test_loader) best_acc = max(best_acc, test_acc) - wandb.log( - { - "test/loss": test_loss, - "test/accuracy": test_acc, - "epoch": epoch, - "examples": example_images, - } - ) + # `train_one_epoch` logs `Training/loss`; logging `Training/accuracy` + # here keeps both charts in a single "Training" section. + wandb.log({"Training/accuracy": test_acc}) history.append( { "epoch": epoch, From 85dd33461a7b7e1e57b28daf68a1cdf9d353679b Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Thu, 16 Jul 2026 17:13:27 -0400 Subject: [PATCH 20/22] add confusion matrix and predictions table; restore image previews MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bring back example predictions the right way: a wandb.Table of un-normalized digit images with true/predicted/confidence, logged once (no per-step slider). Add a confusion matrix so "accuracy" is concrete — which digits the model nails and which it confuses. Both log under the Training prefix after training. Update the intro and verify narration to define accuracy and describe the new panels. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../marimo/mnist-registry/mnist_registry.md | 58 +++++++++++++++++-- .../marimo/mnist-registry/mnist_registry.py | 58 +++++++++++++++++-- 2 files changed, 108 insertions(+), 8 deletions(-) diff --git a/examples/marimo/mnist-registry/mnist_registry.md b/examples/marimo/mnist-registry/mnist_registry.md index d8cebf57..b18994b2 100644 --- a/examples/marimo/mnist-registry/mnist_registry.md +++ b/examples/marimo/mnist-registry/mnist_registry.md @@ -59,8 +59,8 @@ else: ## What you will build -- A **W&B run** with a single **Training** section charting loss and - accuracy. +- A **W&B run** whose **Training** section charts loss and accuracy, plus a + confusion matrix and a table of example predictions. - A **model Artifact** named `mnist-cnn-` of type `model`, carrying metadata (test accuracy, parameter count, hyperparameters). - A version of that Artifact **linked into a W&B Registry collection** @@ -434,8 +434,12 @@ mo.vstack( mo.md(f""" ## Verify and next steps -1. Open the run: [{run.name}]({run.url}) — check the **Training** charts - and the **System** metrics. +1. Open the run: [{run.name}]({run.url}). The **Training** section has the + loss and **accuracy** charts (accuracy = the fraction of held-out test + digits whose top prediction matches the true label), a **confusion + matrix** showing which digits get mistaken for which, and a + **predictions** table of example digits. **System** shows hardware + metrics. 2. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed with its metadata (test accuracy, parameter count, hyperparameters). 3. At [wandb.ai/registry](https://wandb.ai/registry), open the @@ -525,6 +529,50 @@ def evaluate(model, loader): return test_loss / n, correct / n ``` +```python {.marimo name="log_eval_report"} +def log_eval_report(model, loader, n_examples=24): + """Log a confusion matrix and a table of example predictions to W&B. + + Both are logged once, after training, so they render as static panels + without a per-step slider. The table stores un-normalized images so the + digits are visible (logging the normalized tensors renders them black). + """ + model.eval() + all_true, all_pred = [], [] + examples = [] + with torch.no_grad(): + for data, target in loader: + probs = model(data.to(device)).exp() # model returns log_softmax + pred = probs.argmax(dim=1).cpu() + all_true.extend(target.tolist()) + all_pred.extend(pred.tolist()) + for k in range(data.size(0)): + if len(examples) >= n_examples: + break + digit = (data[k] * 0.3081 + 0.1307).clamp(0, 1).squeeze().numpy() + examples.append( + [ + wandb.Image(digit), + int(target[k]), + int(pred[k]), + round(float(probs[k].max()), 4), + ] + ) + wandb.log( + { + "Training/predictions": wandb.Table( + columns=["image", "true", "predicted", "confidence"], + data=examples, + ), + "Training/confusion_matrix": wandb.plot.confusion_matrix( + y_true=all_true, + preds=all_pred, + class_names=[str(i) for i in range(10)], + ), + } + ) +``` + ```python {.marimo name="run_training"} def run_training(run, config, train_ds, test_ds): """Train the CNN, logging metrics each epoch; return the model and history.""" @@ -553,6 +601,8 @@ def run_training(run, config, train_ds, test_ds): "test_acc": round(test_acc, 4), } ) + # Log a confusion matrix and example predictions once, after training. + log_eval_report(model, test_loader) # Full-precision last-epoch accuracy; `history` rounds only for display. return model, history, test_acc, best_acc ``` diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index dd3e6cd5..0795e152 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -61,8 +61,8 @@ def _(): ## What you will build - - A **W&B run** with a single **Training** section charting loss and - accuracy. + - A **W&B run** whose **Training** section charts loss and accuracy, plus a + confusion matrix and a table of example predictions. - A **model Artifact** named `mnist-cnn-` of type `model`, carrying metadata (test accuracy, parameter count, hyperparameters). - A version of that Artifact **linked into a W&B Registry collection** @@ -488,8 +488,12 @@ def _(collection_name_v, registry_name_v, run): mo.md(f""" ## Verify and next steps - 1. Open the run: [{run.name}]({run.url}) — check the **Training** charts - and the **System** metrics. + 1. Open the run: [{run.name}]({run.url}). The **Training** section has the + loss and **accuracy** charts (accuracy = the fraction of held-out test + digits whose top prediction matches the true label), a **confusion + matrix** showing which digits get mistaken for which, and a + **predictions** table of example digits. **System** shows hardware + metrics. 2. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed with its metadata (test accuracy, parameter count, hyperparameters). 3. At [wandb.ai/registry](https://wandb.ai/registry), open the @@ -586,6 +590,50 @@ def evaluate(model, loader): return test_loss / n, correct / n +@app.function +def log_eval_report(model, loader, n_examples=24): + """Log a confusion matrix and a table of example predictions to W&B. + + Both are logged once, after training, so they render as static panels + without a per-step slider. The table stores un-normalized images so the + digits are visible (logging the normalized tensors renders them black). + """ + model.eval() + all_true, all_pred = [], [] + examples = [] + with torch.no_grad(): + for data, target in loader: + probs = model(data.to(device)).exp() # model returns log_softmax + pred = probs.argmax(dim=1).cpu() + all_true.extend(target.tolist()) + all_pred.extend(pred.tolist()) + for k in range(data.size(0)): + if len(examples) >= n_examples: + break + digit = (data[k] * 0.3081 + 0.1307).clamp(0, 1).squeeze().numpy() + examples.append( + [ + wandb.Image(digit), + int(target[k]), + int(pred[k]), + round(float(probs[k].max()), 4), + ] + ) + wandb.log( + { + "Training/predictions": wandb.Table( + columns=["image", "true", "predicted", "confidence"], + data=examples, + ), + "Training/confusion_matrix": wandb.plot.confusion_matrix( + y_true=all_true, + preds=all_pred, + class_names=[str(i) for i in range(10)], + ), + } + ) + + @app.function def run_training(run, config, train_ds, test_ds): """Train the CNN, logging metrics each epoch; return the model and history.""" @@ -614,6 +662,8 @@ def run_training(run, config, train_ds, test_ds): "test_acc": round(test_acc, 4), } ) + # Log a confusion matrix and example predictions once, after training. + log_eval_report(model, test_loader) # Full-precision last-epoch accuracy; `history` rounds only for display. return model, history, test_acc, best_acc From 2983bf68e7de93bb14d25ac73b6843b801865705 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Thu, 16 Jul 2026 17:27:06 -0400 Subject: [PATCH 21/22] add PR curve, embedding projector, and run summary to the W&B report MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Expand log_eval_report to also log a per-digit PR curve and penultimate-layer embeddings (Net.features) for W&B's 2D projector — one cluster per digit, colored by label. Set run.summary with final/best accuracy and parameter count so they surface on the run overview. Document how to add the 2D Projection panel in the verify steps. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../marimo/mnist-registry/mnist_registry.md | 76 +++++++++++++------ .../marimo/mnist-registry/mnist_registry.py | 76 +++++++++++++------ 2 files changed, 106 insertions(+), 46 deletions(-) diff --git a/examples/marimo/mnist-registry/mnist_registry.md b/examples/marimo/mnist-registry/mnist_registry.md index b18994b2..38322d43 100644 --- a/examples/marimo/mnist-registry/mnist_registry.md +++ b/examples/marimo/mnist-registry/mnist_registry.md @@ -59,8 +59,9 @@ else: ## What you will build -- A **W&B run** whose **Training** section charts loss and accuracy, plus a - confusion matrix and a table of example predictions. +- A **W&B run** with a **Training** section (loss, accuracy), a confusion + matrix, a PR curve, a predictions table, and penultimate-layer embeddings + for the 2D projector — plus headline metrics on the run overview. - A **model Artifact** named `mnist-cnn-` of type `model`, carrying metadata (test accuracy, parameter count, hyperparameters). - A version of that Artifact **linked into a W&B Registry collection** @@ -205,12 +206,15 @@ class Net(nn.Module): self.fc1 = nn.Linear(320, 50) self.fc2 = nn.Linear(50, 10) - def forward(self, x): + def features(self, x): + """Penultimate 50-dim representation, used for the embedding projector.""" x = F.relu(F.max_pool2d(self.conv1(x), 2)) x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) x = x.view(-1, 320) - x = F.relu(self.fc1(x)) - x = F.dropout(x, training=self.training) + return F.relu(self.fc1(x)) + + def forward(self, x): + x = F.dropout(self.features(x), training=self.training) x = self.fc2(x) return F.log_softmax(x, dim=1) ``` @@ -437,12 +441,17 @@ mo.md(f""" 1. Open the run: [{run.name}]({run.url}). The **Training** section has the loss and **accuracy** charts (accuracy = the fraction of held-out test digits whose top prediction matches the true label), a **confusion - matrix** showing which digits get mistaken for which, and a - **predictions** table of example digits. **System** shows hardware - metrics. -2. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed + matrix** showing which digits get mistaken for which, a per-digit + **PR curve**, and a **predictions** table of example digits. The run + **Overview** shows headline numbers (final/best accuracy, parameter + count); **System** shows hardware metrics. +2. **See the digit clusters.** Add a panel → **2D Projection**, choose the + `Training/embeddings` table, set the vector to `embedding` and color by + `digit`. W&B projects the penultimate-layer features to 2D — one cluster + per digit, tighter as the model improves. +3. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed with its metadata (test accuracy, parameter count, hyperparameters). -3. At [wandb.ai/registry](https://wandb.ai/registry), open the +4. At [wandb.ai/registry](https://wandb.ai/registry), open the **{registry_name_v.title()}** registry, then the **{collection_name_v}** collection, and confirm the linked version. @@ -530,22 +539,34 @@ def evaluate(model, loader): ``` ```python {.marimo name="log_eval_report"} -def log_eval_report(model, loader, n_examples=24): - """Log a confusion matrix and a table of example predictions to W&B. - - Both are logged once, after training, so they render as static panels - without a per-step slider. The table stores un-normalized images so the - digits are visible (logging the normalized tensors renders them black). +def log_eval_report(model, loader, n_examples=24, n_embeddings=500): + """Log evaluation visuals to W&B once, after training: + + - a **predictions** table of un-normalized example digits (true / + predicted / confidence) — logged once, so there is no per-step slider; + - a **confusion matrix** over the whole test set; + - a per-digit **PR curve** from the predicted probabilities; + - **penultimate-layer embeddings** for the W&B 2D projector (add a "2D + Projection" panel and color by `digit` to see one cluster per class). """ model.eval() - all_true, all_pred = [], [] - examples = [] + all_true, all_pred, all_probs = [], [], [] + examples, embeddings = [], [] with torch.no_grad(): for data, target in loader: - probs = model(data.to(device)).exp() # model returns log_softmax + data_dev = data.to(device) + probs = model(data_dev).exp() # model returns log_softmax pred = probs.argmax(dim=1).cpu() + probs = probs.cpu() all_true.extend(target.tolist()) all_pred.extend(pred.tolist()) + all_probs.extend(probs.tolist()) + if len(embeddings) < n_embeddings: + feats = model.features(data_dev).cpu() + for k in range(data.size(0)): + if len(embeddings) >= n_embeddings: + break + embeddings.append([int(target[k]), feats[k].tolist()]) for k in range(data.size(0)): if len(examples) >= n_examples: break @@ -558,6 +579,7 @@ def log_eval_report(model, loader, n_examples=24): round(float(probs[k].max()), 4), ] ) + class_names = [str(i) for i in range(10)] wandb.log( { "Training/predictions": wandb.Table( @@ -565,9 +587,13 @@ def log_eval_report(model, loader, n_examples=24): data=examples, ), "Training/confusion_matrix": wandb.plot.confusion_matrix( - y_true=all_true, - preds=all_pred, - class_names=[str(i) for i in range(10)], + y_true=all_true, preds=all_pred, class_names=class_names + ), + "Training/pr_curve": wandb.plot.pr_curve( + all_true, all_probs, labels=class_names + ), + "Training/embeddings": wandb.Table( + columns=["digit", "embedding"], data=embeddings ), } ) @@ -601,8 +627,12 @@ def run_training(run, config, train_ds, test_ds): "test_acc": round(test_acc, 4), } ) - # Log a confusion matrix and example predictions once, after training. + # Log a confusion matrix, PR curve, predictions, and embeddings once. log_eval_report(model, test_loader) + # Headline numbers on the run overview page. + run.summary["final_test_accuracy"] = test_acc + run.summary["best_test_accuracy"] = best_acc + run.summary["num_parameters"] = sum(p.numel() for p in model.parameters()) # Full-precision last-epoch accuracy; `history` rounds only for display. return model, history, test_acc, best_acc ``` diff --git a/examples/marimo/mnist-registry/mnist_registry.py b/examples/marimo/mnist-registry/mnist_registry.py index 0795e152..064ba503 100644 --- a/examples/marimo/mnist-registry/mnist_registry.py +++ b/examples/marimo/mnist-registry/mnist_registry.py @@ -61,8 +61,9 @@ def _(): ## What you will build - - A **W&B run** whose **Training** section charts loss and accuracy, plus a - confusion matrix and a table of example predictions. + - A **W&B run** with a **Training** section (loss, accuracy), a confusion + matrix, a PR curve, a predictions table, and penultimate-layer embeddings + for the 2D projector — plus headline metrics on the run overview. - A **model Artifact** named `mnist-cnn-` of type `model`, carrying metadata (test accuracy, parameter count, hyperparameters). - A version of that Artifact **linked into a W&B Registry collection** @@ -222,12 +223,15 @@ def __init__(self): self.fc1 = nn.Linear(320, 50) self.fc2 = nn.Linear(50, 10) - def forward(self, x): + def features(self, x): + """Penultimate 50-dim representation, used for the embedding projector.""" x = F.relu(F.max_pool2d(self.conv1(x), 2)) x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) x = x.view(-1, 320) - x = F.relu(self.fc1(x)) - x = F.dropout(x, training=self.training) + return F.relu(self.fc1(x)) + + def forward(self, x): + x = F.dropout(self.features(x), training=self.training) x = self.fc2(x) return F.log_softmax(x, dim=1) @@ -491,12 +495,17 @@ def _(collection_name_v, registry_name_v, run): 1. Open the run: [{run.name}]({run.url}). The **Training** section has the loss and **accuracy** charts (accuracy = the fraction of held-out test digits whose top prediction matches the true label), a **confusion - matrix** showing which digits get mistaken for which, and a - **predictions** table of example digits. **System** shows hardware - metrics. - 2. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed + matrix** showing which digits get mistaken for which, a per-digit + **PR curve**, and a **predictions** table of example digits. The run + **Overview** shows headline numbers (final/best accuracy, parameter + count); **System** shows hardware metrics. + 2. **See the digit clusters.** Add a panel → **2D Projection**, choose the + `Training/embeddings` table, set the vector to `embedding` and color by + `digit`. W&B projects the penultimate-layer features to 2D — one cluster + per digit, tighter as the model improves. + 3. In the run's **Artifacts** tab, confirm `mnist-cnn-{run.id}` is listed with its metadata (test accuracy, parameter count, hyperparameters). - 3. At [wandb.ai/registry](https://wandb.ai/registry), open the + 4. At [wandb.ai/registry](https://wandb.ai/registry), open the **{registry_name_v.title()}** registry, then the **{collection_name_v}** collection, and confirm the linked version. @@ -591,22 +600,34 @@ def evaluate(model, loader): @app.function -def log_eval_report(model, loader, n_examples=24): - """Log a confusion matrix and a table of example predictions to W&B. - - Both are logged once, after training, so they render as static panels - without a per-step slider. The table stores un-normalized images so the - digits are visible (logging the normalized tensors renders them black). +def log_eval_report(model, loader, n_examples=24, n_embeddings=500): + """Log evaluation visuals to W&B once, after training: + + - a **predictions** table of un-normalized example digits (true / + predicted / confidence) — logged once, so there is no per-step slider; + - a **confusion matrix** over the whole test set; + - a per-digit **PR curve** from the predicted probabilities; + - **penultimate-layer embeddings** for the W&B 2D projector (add a "2D + Projection" panel and color by `digit` to see one cluster per class). """ model.eval() - all_true, all_pred = [], [] - examples = [] + all_true, all_pred, all_probs = [], [], [] + examples, embeddings = [], [] with torch.no_grad(): for data, target in loader: - probs = model(data.to(device)).exp() # model returns log_softmax + data_dev = data.to(device) + probs = model(data_dev).exp() # model returns log_softmax pred = probs.argmax(dim=1).cpu() + probs = probs.cpu() all_true.extend(target.tolist()) all_pred.extend(pred.tolist()) + all_probs.extend(probs.tolist()) + if len(embeddings) < n_embeddings: + feats = model.features(data_dev).cpu() + for k in range(data.size(0)): + if len(embeddings) >= n_embeddings: + break + embeddings.append([int(target[k]), feats[k].tolist()]) for k in range(data.size(0)): if len(examples) >= n_examples: break @@ -619,6 +640,7 @@ def log_eval_report(model, loader, n_examples=24): round(float(probs[k].max()), 4), ] ) + class_names = [str(i) for i in range(10)] wandb.log( { "Training/predictions": wandb.Table( @@ -626,9 +648,13 @@ def log_eval_report(model, loader, n_examples=24): data=examples, ), "Training/confusion_matrix": wandb.plot.confusion_matrix( - y_true=all_true, - preds=all_pred, - class_names=[str(i) for i in range(10)], + y_true=all_true, preds=all_pred, class_names=class_names + ), + "Training/pr_curve": wandb.plot.pr_curve( + all_true, all_probs, labels=class_names + ), + "Training/embeddings": wandb.Table( + columns=["digit", "embedding"], data=embeddings ), } ) @@ -662,8 +688,12 @@ def run_training(run, config, train_ds, test_ds): "test_acc": round(test_acc, 4), } ) - # Log a confusion matrix and example predictions once, after training. + # Log a confusion matrix, PR curve, predictions, and embeddings once. log_eval_report(model, test_loader) + # Headline numbers on the run overview page. + run.summary["final_test_accuracy"] = test_acc + run.summary["best_test_accuracy"] = best_acc + run.summary["num_parameters"] = sum(p.numel() for p in model.parameters()) # Full-precision last-epoch accuracy; `history` rounds only for display. return model, history, test_acc, best_acc From 9c104aa135c1980a697406ae1d4b4be61b3cd163 Mon Sep 17 00:00:00 2001 From: John Mulhausen Date: Thu, 30 Jul 2026 16:53:39 -0400 Subject: [PATCH 22/22] sync README with the current notebook The README had drifted: it still described wandb.watch gradient histograms and a 16-image prediction gallery, both removed. Update "What you get" to the actual W&B report (Training loss/accuracy, confusion matrix, PR curve, predictions table, 2D-projector embeddings, run-overview headline metrics), describe the in-notebook Evaluation step, add the team-entity prerequisite, and reword the gating note for the marimo form. An agent audit confirms README<->notebook consistency. Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/marimo/mnist-registry/README.md | 37 +++++++++++++++--------- 1 file changed, 23 insertions(+), 14 deletions(-) diff --git a/examples/marimo/mnist-registry/README.md b/examples/marimo/mnist-registry/README.md index cb3b89db..7df84374 100644 --- a/examples/marimo/mnist-registry/README.md +++ b/examples/marimo/mnist-registry/README.md @@ -16,6 +16,9 @@ can resolve them automatically. shell before launching the notebook, or paste your key into the **W&B API key** field in the form. Get your key from [wandb.ai/authorize](https://wandb.ai/authorize). +- A W&B **team** to write the run to, set in the **W&B entity** field. Accounts + created after May 2024 have no personal entity, so the run must go to a team + — your username will not work as an entity. - A W&B **Registry** must exist in your org, and your account needs at least the **Member** role on it for the final linking step (linking an artifact is a write action). The built-in Model registry is provisioned automatically in @@ -45,23 +48,29 @@ pip install -r requirements.txt marimo edit mnist_registry.py ``` -The notebook is interactive-only by design: training is gated by a button -click, so `marimo run` renders the form but never starts training without -an explicit click. +The notebook is interactive-only by design: training is gated by submitting +the form, so `marimo run` renders the form but never starts training until you +click **Train model**. ## What you get After a successful run: -- A W&B run with training and test metrics, gradient histograms (`wandb.watch`), - and up to 16 example test-set predictions logged as images. +- A W&B run whose **Training** section charts loss and accuracy, alongside a + confusion matrix, a per-digit PR curve, a table of example predictions (with + images), and penultimate-layer embeddings you can open in a 2D projector — + one cluster per digit. Final/best accuracy and the parameter count are + surfaced on the run overview. - A model Artifact named `mnist-cnn-` of type `model` with metadata for test accuracy, parameter count, dataset sizes, and the full hyperparameter dict. Tagged with the `latest` alias. - A version of that Artifact linked into the configured Registry collection (default: `wandb-registry-model/MNIST Classifiers`). -To consume the registered model from another script or notebook: +The notebook then **consumes the model in place**: its Evaluation cell +downloads the artifact (preferring the registered version, falling back to the +run's own artifact) and classifies ten held-out test digits, showing predicted +versus true labels. To consume it from another script or notebook: ```python import wandb @@ -72,20 +81,20 @@ art.download() # writes mnist_cnn.pt under ./artifacts/ ## Design notes -- **Training is gated by a button.** marimo cells re-run reactively when their - inputs change. Before the first click of **Train model**, slider changes do - not start a run. After a run completes, clicking **Train model** again - starts a new run with the current form values; the previous run finishes - cleanly first. +- **Training is gated by a form.** Hyperparameters live in a marimo form, so + changing a field does nothing until you submit it with **Train model**. + Submitting again after a run starts a new run with the current values; the + previous run is finished cleanly first. - **`wandb.run` finishes defensively** at the top of the training cell so a second click of **Train model** does not nest runs in the same marimo kernel. - **`logged.wait()` runs** after `log_artifact` and before `link_artifact` to avoid a race where the link tries to resolve a version that has not finished committing server-side. -- **Registry failures soft-fail.** If `link_artifact` raises — usually - because the Registry does not exist in your org — the notebook - surfaces remediation guidance through `mo.callout` rather than aborting. +- **Registry failures soft-fail.** If linking raises — usually a + view-only seat or a Registry that does not exist in your org — the + notebook surfaces remediation guidance through `mo.callout` rather than + aborting; the run and artifact still succeed. ## Reference