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 フラグは、サンプリングする組み合わせの数と RNG シードを制御します。

このリファレンスはコンピュートプロバイダーとして Hugging Face Jobs を使用していますが、これは 1 つのオプションに過ぎません。ユーザーは Modal、RunPod、または uv スクリプトを実行できる他のプロバイダーに置き換えることができます。