ジョブ入力を管理する

Launch のコア体験は、ハイパーパラメーターやデータセットのような異なるジョブの入力を簡単に実験し、それらのジョブを適切なハードウェアにルーティングすることです。一度ジョブが作成されると、元の作成者以外のユーザーも W&B GUI または CLI を介してこれらの入力を調整できます。CLI または UI からジョブ入力を設定する方法については、Enqueue jobs ガイドを参照してください。

このセクションでは、プログラム的にジョブの調整可能な入力を制御する方法を説明します。

デフォルトでは、W&B ジョブは Run.config 全体をジョブの入力としてキャプチャしますが、Launch SDK は run config の選択したキーを制御したり、JSON または YAML ファイルを入力として指定するための機能を提供します。

Run オブジェクトの再設定

ジョブ内の wandb.init によって返される Run オブジェクトは、デフォルトで再設定可能です。Launch SDK は、ジョブの Launch 時に Run.config オブジェクトのどの部分が再設定可能かをカスタマイズする方法を提供します。

import wandb
from wandb.sdk import launch

# Launch SDK の使用に必要
wandb.require("core")

config = {
    "trainer": {
        "learning_rate": 0.01,
        "batch_size": 32,
        "model": "resnet",
        "dataset": "cifar10",
        "private": {
            "key": "value",
        },
    },
    "seed": 42,
}

with wandb.init(config=config):
    launch.manage_wandb_config(
        include=["trainer"], 
        exclude=["trainer.private"],
    )
    # 等

関数 launch.manage_wandb_config は、Run.config オブジェクトに対する入力値をジョブに受け入れるように設定します。オプションの include および exclude オプションは、ネストされた config オブジェクト内のパスプレフィクスを取ります。例えば、ジョブがエンドユーザーに公開したくないオプションを持つライブラリを使用している場合に役立ちます。

include プレフィクスが提供されている場合、config 内で include プレフィクスと一致するパスのみが入力値を受け入れます。exclude プレフィクスが提供されている場合、exclude リストと一致するパスは入力値からフィルタリングされません。あるパスが include および exclude プレフィクスの両方に一致する場合、exclude プレフィクスが優先されます。

前述の例では、パス ["trainer.private"] は trainer オブジェクトから private キーをフィルタリングし、パス ["trainer"] は trainer オブジェクト下のすべてのキー以外をフィルタリングします。

上記のコードがパッケージ化され、ジョブとして実行される場合、ジョブの入力タイプは次のようになります：

{
    "trainer": {
        "learning_rate": "float",
        "batch_size": "int",
        "model": "str",
        "dataset": "str",
    },
}

W&B CLI または UI からジョブをローンチする際、ユーザーは 4 つの trainer パラメーターのみを上書きできます。

run config の入力へのアクセス

run config の入力でローンチされたジョブは、Run.config を通じて入力値にアクセスできます。ジョブ コード内で wandb.init によって返された Run は、入力値を自動的に設定します。ジョブ コードのどこでも run config の入力値をロードするには、次を使用します：

from wandb.sdk import launch

run_config_overrides = launch.load_wandb_config()

ファイルの再設定

Launch SDK はまた、ジョブ コードで設定ファイルに格納された入力値を管理する方法も提供します。これは、多くのディープラーニングや大規模な言語モデルのユースケースで一般的なパターンであり、例えばこの torchtune の例やこの Axolotl config などがあります。

launch.manage_config_file 関数を使用して、設定ファイルを Launch ジョブの入力として追加し、ジョブをローンチする際に設定ファイル内の値の編集アクセスを提供します。

デフォルトでは、launch.manage_config_file が使用されると run config 入力はキャプチャされません。launch.manage_wandb_config を呼び出すことで、この振る舞いを上書きします。

以下の例を考えてみましょう：

import yaml
import wandb
from wandb.sdk import launch

# Launch SDK の使用に必要
wandb.require("core")

launch.manage_config_file("config.yaml")

with open("config.yaml", "r") as f:
    config = yaml.safe_load(f)

with wandb.init(config=config):
    # 等
    pass

コードが隣接するファイル config.yaml と一緒に実行されていると想像してください：

learning_rate: 0.01
batch_size: 32
model: resnet
dataset: cifar10

launch.manage_config_file の呼び出しは、config.yaml ファイルをジョブの入力として追加し、W&B CLI または UI からローンチする際に再設定可能にします。

launch.manage_wandb_config と同じようにして、設定ファイルに対する許容入力キーをフィルタリングするために include と exclude のキーワード引数を使用できます。

設定ファイルの入力へのアクセス

Launch によって作成された run で launch.manage_config_file が呼び出されると、launch は設定ファイルの内容を入力値でパッチします。修正された設定ファイルはジョブ 環境で利用可能です。

ジョブの launch ドロワー UI のカスタマイズ

ジョブの入力スキーマを定義することで、ジョブをローンチするためのカスタム UI を作成できます。ジョブのスキーマを定義するには、launch.manage_wandb_config または launch.manage_config_file の呼び出しにそれを含めます。スキーマは、JSON Schema の形での Python 辞書または Pydantic モデル クラスのいずれかです。

• JSON schema
• Pydantic model

schema = {
    "type": "object",
    "properties": {
        "seed": {
          "type": "integer"
        },
        "trainer": {
            "type": "object",
            "properties": {
                "learning_rate": {
                    "type": "number",
                    "description": "モデルの学習率",
                    "exclusiveMinimum": 0,
                },
                "batch_size": {
                    "type": "integer",
                    "description": "バッチごとのサンプル数",
                    "enum": [16, 64, 256]
                },
                "dataset": {
                    "type": "string",
                    "description": "使用するデータセットの名前",
                    "enum": ["cifar10", "cifar100"]
                }
            }
        }
    }
}

launch.manage_wandb_config(
    include=["seed", "trainer"], 
    exclude=["trainer.private"],
    schema=schema,
)

class DatasetEnum(str, Enum):
    cifar10 = "cifar10"
    cifar100 = "cifar100"

class Trainer(BaseModel):
    learning_rate: float = Field(gt=0, description="モデルの学習率")
    batch_size: int = Field(ge=1, le=256, description="バッチごとのサンプル数")
    dataset: DatasetEnum = Field(title="Dataset", description="使用するデータセットの名前")

class Schema(BaseModel):
    seed: int
    trainer: Trainer

launch.manage_wandb_config(
    include=["seed", "trainer"],
    exclude=["trainer.private"],
    schema=Schema,
)

t = Trainer(learning_rate=0.01, batch_size=32, dataset=DatasetEnum.cifar10)
s = Schema(seed=42, trainer=t)
launch.manage_wandb_config(
    include=["seed", "trainer"],
    exclude=["trainer.private"],
    input_schema=s,
)

ジョブ入力スキーマを追加すると、launch ドロワーに構造化されたフォームが作成され、ジョブのローンチが容易になります。