marimo-batch-mlflow¶
Opinionated fork of marimo-team/skills/marimo-batch that swaps three things:
| Concern | upstream marimo-batch |
this skill |
|---|---|---|
| CLI parsing | mo.cli_args() + hand-rolled rich.Table for --help |
tyro.cli(ModelParams) — auto --help, type coercion, validation |
| Params model | Pydantic BaseModel |
dataclass (primary) or pydantic.BaseModel (alternative) |
| Tracking | Weights and Biases | MLflow (mlflow + optional mlflow-widgets for live charts) |
Pick this skill when the user has (or is happy with) MLflow and prefers strongly-typed CLIs. Pick upstream when the user is already on W&B.
What ships¶
- The full SKILL.md (skills/local/marimo-batch-mlflow/SKILL.md) with a decision matrix vs upstream, a dual-mode pattern recipe, and cell-level templates for params / EnvConfig / training loop / live chart.
- A reference notebook
(starting-point.py)
— full PyTorch training notebook with dataclass params, Tyro CLI, MLflow
logging, and a live
MlflowChartcell that's gated off in script mode. - A Pydantic-variant snippet (params-pydantic.py) — only the params cell differs; rest of the notebook is identical.
- A grid-search launcher
(grid.py)
— randomized sweeps via Hugging Face Jobs with MLflow secrets injected.
Dry-run by default;
--launchactually submits.
Dual-mode pattern (the core idiom)¶
import marimo as mo
import tyro
is_script_mode = mo.app_meta().mode == "script"
if is_script_mode:
params = tyro.cli(ModelParams) # CLI flags
else:
mo.stop(form.value is None, mo.md("*Submit form to start.*"))
params = ModelParams(**form.value) # UI form
# Every cell below uses `params.epochs`, `params.batch_size`, ...
# unaware of which branch produced it.
Same notebook serves as the UI for fast iteration and the entry point
for uv run notebook.py --epochs 50 batch jobs. No code duplication.
Why a fork instead of vendoring¶
The upstream marimo-batch skill is opinionated towards W&B and
mo.cli_args(). Vendoring + patching would be clobbered every sync; a
local fork in skills/local/ keeps the divergence safe and explicit. The
SKILL.md cross-references upstream so users can pick the right variant.
Cross-references¶
- Upstream
marimo-batch— the W&B variant; pick this if the user is already on W&B. marimo-notebook— general marimo authoring patterns (vendored from marimo-team).anywidget-generator— for building custom live-chart widgets ifmlflow-widgetsdoesn't cover the case (vendored from marimo-team).mlflow-widgets— the anywidget-based MLflow chart/table/parallel-coordinates components used instarting-point.py's live-chart cell.- Tyro docs — CLI generation reference; supports dataclass, Pydantic, attrs.