marimo-batch

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.")

このモデルのパラメーターには、2つの方法で値を設定できます。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 のトレーニングジョブを開始したい場合は、こちらの開始用ファイルを使用してください。このノートブックの列はそのまま維持してください。

環境変数

ジョブで環境変数を読み取る必要がある場合があります。.env ファイルが存在すれば python-dotenv で読み取るようにし、ユーザーが UI でキーを手動で追加できるように EnvConfig も追加してください。

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 はノートブックの先頭に追加してください。

列

大きな marimo ノートブックでは、ナビゲーションしやすくするために列機能を使うのが一般的です。その場合は、列を必ず維持してください。

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

コンピューティングプラットフォーム

ジョブに本格的なコンピューティングリソースを割り当てる準備ができたら、適切なプラクティスを心がけることが重要です。データセットのバッチサイズを検討し、問題が発生した際にユーザーが気づけるよう、十分なログを出力してください。

グリッドサーチ

ユーザーがハイパーパラメーターのスイープを実行したい場合は、こちらのグリッドランチャーを案内してください。references/starting-point.py のノートブックでそのまま動作します。ノートブックの ModelParams フィールドに合致する探索空間からランダムな組み合わせをサンプリングし、それぞれを個別のジョブとして起動します。

デフォルトではスクリプトはドライラン（uv run grid.py）を実行するため、コンピューティングリソースを使う前に組み合わせを確認できます。ジョブを実際に送信するには --launch を指定します。--count フラグと --seed フラグで、サンプリングする組み合わせの数と乱数シードを設定します。

このリファレンスではコンピューティングプロバイダーとして Hugging Face Jobs を使用していますが、これは選択肢の一つにすぎません。uv スクリプトを実行できる Modal、RunPod、その他のプロバイダーに置き換えることができます。