ジョブの入力を管理する

Launch の中心となるのは、ハイパーパラメーターやデータセットなどのジョブ入力をさまざまに変えて実験し、それらのジョブを適切なハードウェアにルーティングすることです。ジョブを作成すると、元の作成者以外のユーザーも W&B UI または CLI からこれらの入力を調整できます。CLI または UI から Launch する際にジョブ入力を設定する方法については、ジョブをキューに追加するガイドを参照してください。このガイドでは、ジョブで調整可能な入力をプログラムで制御し、エンドユーザーに変更を許可したいパラメーターのみを公開する方法について説明します。デフォルトでは、W&B のジョブは Run.config 全体をジョブの入力として取得します。ただし、Launch SDK には、run 設定内の特定のキーのみを制御対象にしたり、JSON または YAML ファイルを入力として指定したりするための関数が用意されています。

Launch SDK の関数を使用するには wandb-core が必要です。詳細については、wandb-core README を参照してください。

`Run` オブジェクトを再設定する

デフォルトでは、ジョブ内で wandb.init() が返す Run オブジェクトを再設定できます。Launch SDK を使用すると、ジョブの起動時に 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 には、ネストされた設定オブジェクト内のパス接頭辞を指定します。これは、たとえばジョブが使用するライブラリのオプションをエンドユーザーに公開したくない場合に便利です。 include 接頭辞を指定すると、設定内のパスのうち include 接頭辞に一致するものだけが入力値を受け付けます。exclude 接頭辞を指定すると、exclude リストに一致するパスが入力値から除外されます。パスが include 接頭辞と exclude 接頭辞の両方に一致する場合は、exclude 接頭辞が優先されます。上記の例では、パス ["trainer.private"] は trainer オブジェクトから private キーを除外し、パス ["trainer"] は trainer オブジェクト配下以外のすべてのキーを除外します。

名前に . を含むキーを除外するには、\ でエスケープした . を使用します。たとえば、r"trainer\.private" は、trainer オブジェクト配下の private キーではなく、trainer.private キーを除外します。上記の例の r 接頭辞は raw string を表します。

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

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

W&B CLI または UI からジョブを起動する場合、上書きできるのは 4 つの trainer パラメーターだけです。

run 設定の入力にアクセスする

run 設定の入力を指定して起動したジョブでは、Run.config から入力値にアクセスできます。ジョブのコード内で wandb.init() が返す Run には、入力値が自動的に設定されています。ジョブのコード内の任意の場所で run 設定の入力値を読み込むには、launch.load_wandb_config() を使用します。

from wandb.sdk import launch

run_config_overrides = launch.load_wandb_config()

ファイルを再設定する

Launch SDK では、ジョブのコード内の設定ファイルに保存された入力値も管理できます。これは、多くのディープラーニングや大規模言語モデルのユースケースでよく見られるパターンです。たとえば、この torchtune の例や、この Axolotl の設定などが挙げられます。

Launch で Sweeps では、設定ファイルの入力を sweep のパラメーターとして使用することはできません。sweep のパラメーターは Run.config オブジェクトで制御する必要があります。

launch.manage_config_file() 関数を使用すると、設定ファイルを Launch job の入力として追加できます。これにより、ジョブの起動時に設定ファイル内の値を編集できるようになります。デフォルトでは、launch.manage_config_file() を使用すると run 設定の入力は取得されません。この動作は、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.manage_config_file() を呼び出してください。

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

公開する入力をフィルターするだけでなく、ジョブの入力に対してスキーマを定義し、ジョブを起動するためのカスタム UI を作成することもできます。スキーマを定義すると、launch ドロワーには自由形式のテキスト入力欄の代わりに、構造化されたフィールド、検証のヒント、ドロップダウンが表示されます。ジョブのスキーマを定義するには、launch.manage_wandb_config() または launch.manage_config_file() の呼び出しにスキーマを含めます。スキーマには、JSON スキーマ形式の Python dict、または Pydantic モデルクラスを使用できます。

ジョブ入力スキーマは入力を検証しません。launch ドロワーの UI を定義するだけです。

JSON スキーマ
Pydantic モデル

次の例は、以下のプロパティを持つスキーマです。

seed: 整数。
trainer: いくつかのキーが指定された辞書。
- trainer.learning_rate: 0 より大きい浮動小数点数。
- trainer.batch_size: 16、64、256 のいずれかの整数。
- trainer.dataset: cifar10 または cifar100 のいずれかの string。

schema = {
    "type": "object",
    "properties": {
        "seed": {
          "type": "integer"
        }
        "trainer": {
            "type": "object",
            "properties": {
                "learning_rate": {
                    "type": "number",
                    "description": "Learning rate of the model",
                    "exclusiveMinimum": 0,
                },
                "batch_size": {
                    "type": "integer",
                    "description": "Number of samples per batch",
                    "enum": [16, 64, 256]
                },
                "dataset": {
                    "type": "string",
                    "description": "Name of the dataset to use",
                    "enum": ["cifar10", "cifar100"]
                }
            }
        }
    }
}

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

一般的に、次の JSON スキーマ属性がサポートされています。

属性	必須	メモ
`type`	はい	`number`、`integer`、`string`、`object` のいずれかである必要があります
`title`	いいえ	プロパティの表示名を上書きします
`description`	いいえ	プロパティにヘルパーテキストを追加します
`enum`	いいえ	自由形式のテキスト入力欄の代わりにドロップダウンの選択肢を作成します
`minimum`	いいえ	`type` が `number` または `integer` の場合にのみ使用できます
`maximum`	いいえ	`type` が `number` または `integer` の場合にのみ使用できます
`exclusiveMinimum`	いいえ	`type` が `number` または `integer` の場合にのみ使用できます
`exclusiveMaximum`	いいえ	`type` が `number` または `integer` の場合にのみ使用できます
`properties`	いいえ	`type` が `object` の場合に、ネストされた設定を定義します

次の例は、以下のプロパティを持つスキーマです。

seed: 整数。
trainer: いくつかのサブ属性が指定されたスキーマ。
- trainer.learning_rate: 0 より大きい浮動小数点数。
- trainer.batch_size: 1 以上 256 以下の整数。
- trainer.dataset: cifar10 または cifar100 のいずれかの string。

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

class Trainer(BaseModel):
    learning_rate: float = Field(gt=0, description="Learning rate of the model")
    batch_size: int = Field(ge=1, le=256, description="Number of samples per batch")
    dataset: DatasetEnum = Field(title="Dataset", description="Name of the dataset to use")

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 ドロワーに構造化されたフォームが表示されます。

​Run オブジェクトを再設定する

​run 設定の入力にアクセスする

​ファイルを再設定する

​設定ファイルの入力にアクセスする

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

`Run` オブジェクトを再設定する

run 設定の入力にアクセスする

ファイルを再設定する

設定ファイルの入力にアクセスする

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