mcpField
M
marimo-batch

marimo-batch

github.com/marimo-team数据与数据库1773.7k 次安装10 周前更新
代码仓库

关于 marimo-batch

marimo-batch 是一项将 marimo 笔记本改造成可配置批处理作业的技能，面向运行机器学习或其他参数化任务的笔记本用户。它涵盖 UI 与命令行参数、环境配置及可选的网格搜索。

主要功能

  1. 01用 Pydantic 定义批处理参数
  2. 02通过 marimo UI 和命令行传参
  3. 03使用 EnvConfig 配置环境变量
  4. 04设置可选的网格搜索任务

适用场景

  1. 01按指定命令行参数运行笔记本
  2. 02用较小参数快速检查笔记本
  3. 03抽样并启动超参数组合任务

安装

复制下面的命令，在终端中运行。

$npx skills add https://github.com/marimo-team/skills --skill marimo-batch

安装后在 Claude Code 中输入/marimo-batch确认已加载。

附带文件references/grid.pyreferences/starting-point.py

详细介绍

Pydantic 是声明批处理任务事实来源的绝佳方式，尤其适用于 ML。你可以声明如下内容：

from pydantic import BaseModel, Field

class ModelParams(BaseModel):
    sample_size: int = Field(
        default=1024 * 4, description="Number of training samples per epoch."
    )
    learning_rate: float = Field(default=0.01, description="Learning rate for the optimizer.")

你也可以通过两种方式填写这些模型参数，还可以想象一下 UI 中的表单。

el = mo.md("""
{sample_size} 
{learning_rate}
""").batch(
    sample_size=mo.ui.slider(1024, 1024 * 10, value=1024 * 4, step=1024, label="Sample size"),
    learning_rate=mo.ui.slider(0.001, 0.1, value=0.01, step=0.001, label="Learning rate"),
).form()
el

不过，你也可以使用 marimo 的 CLI。

if mo.app_meta().mode == "script":
    if "help" in mo.cli_args() or len(cli_args) == 0:
        print("Usage: uv run git_archaeology.py --repo <url> [--samples <n>]")
        print()
        for name, field in ModelParams.model_fields.items():
            default = f" (default: {field.default})" if field.default is not None else " (required)"
            print(f"  --{name:12s} {field.description}{default}")
        exit()
    model_params = ModelParams(
        **{k.replace("-", "_"): v for k, v in mo.cli_args().items()
    })
else: 
    model_params = ModelParams(**el.value)

现在，用户可以通过命令行运行：

uv run notebook.py --sample-size 4096 --learning-rate 0.005

这样就能兼顾两全：你可以使用 UI 进行测试和迭代，然后使用 CLI 运行批处理任务。另一个好处是，你可以使用相关设置运行笔记本，以便快速检查笔记本中是否存在错误。

用户希望使用这种模式运行笔记本，因此请务必询问用户希望通过 CLI 配置哪些参数，然后再对笔记本进行相应修改。修改前，请务必先让用户确认这些更改。

Weights and Biases

用户可能有兴趣添加对 weights and biases 的支持。请务必确认用户是否有此需求（是/否）。如果有，请确保记录这些 ModelParams。还要确保将 wandb_project 和 wandb_run_name 纳入 ModelParams（如果用户希望采用这种方案）。

如果用户希望启动 ML 训练任务，请务必使用这个起点。请确保保留此笔记本中的列！

Environment Variables

你可能需要读取任务的环境变量。如果 .env 文件存在，请使用 python-dotenv 读取；同时还要添加一个 EnvConfig，以便用户可以在 UI 中手动添加键。

from wigglystuff import EnvConfig

# With validators
config = EnvConfig({
    "OPENAI_API_KEY": lambda k: openai.Client(api_key=k).models.list(),
    "WANDB_API_KEY": lambda k: wandb.login(key=k, verify=True)
})

# Block until valid, useful in cell that needs the key
config.require_valid()

# Access values
config["OPENAI_API_KEY"]
config.get("OPENAI_API_KEY", "some default")

请务必在笔记本顶部添加此 EnvConfig。

Columns

较大的 marimo 笔记本通常会使用 columns 功能，以便于导航。如果是这种情况，必须保留这些列！

@app.cell(column=0, hide_code=True)
def _(mo):
    mo.md(r"""demo""")

Compute platform

当任务准备好使用大量算力时，请务必遵循良好的实践。考虑数据集的批次大小，并确保提供充足的日志，以便用户发现可能出现的问题。

Grid search

当用户希望运行超参数扫描时，请引导他们使用这个网格启动器。它可以直接与 references/starting-point.py 中的笔记本配合使用：它会从与笔记本 ModelParams 字段匹配的搜索空间中随机抽取组合，并将每个组合作为独立任务启动。

默认情况下，脚本会进行试运行（uv run grid.py），以便用户在消耗算力前检查这些组合。传入 --launch 可实际提交任务。--count 和 --seed 标志分别用于控制抽取的组合数量和 RNG 种子。

此参考方案使用 Hugging Face Jobs 作为计算提供方，但这只是其中一种选择。用户可以将其替换为 Modal、RunPod，或任何能够运行 uv 脚本的其他提供方。

评论

支持 Markdown