チュートリアル: Vertex AI で W&B Launch を設定する

このチュートリアルでは、W&B Launch を設定して、Vertex AI のトレーニングジョブとして実行するジョブを送信する方法を説明します。これにより、トレーニングワークロードを Google Cloud のマネージドインフラストラクチャーにオフロードできます。Vertex AI のトレーニングジョブでは、Vertex AI プラットフォーム上で、提供されるアルゴリズムまたはカスタムアルゴリズムを使用して機械学習モデルをトレーニングできます。Launch ジョブを開始すると、Vertex AI が基盤となるインフラストラクチャー、スケーリング、オーケストレーションを管理します。このガイドは、すでに W&B Launch を使用していて、Google Cloud Vertex AI 上でジョブを実行したい ML エンジニアおよびプラットフォーム管理者を対象としています。 W&B Launch は、google-cloud-aiplatform SDK の CustomJob クラスを介して Vertex AI と連携します。CustomJob のパラメーターは、Launch キュー設定で制御できます。Google Cloud 外のプライベート Registry からイメージを pull するように Vertex AI を設定することはできません。つまり、Vertex AI を W&B Launch とともに使用する場合は、コンテナーイメージを Google Cloud 内またはパブリック Registry に保存する必要があります。Vertex ジョブからコンテナーイメージにアクセスできるようにする方法の詳細については、Vertex AI documentation を参照してください。

前提条件

Launch queue を設定する前に、次の Google Cloud リソースと権限が用意されていることを確認してください。

Vertex AI API が有効な Google Cloud project を作成するか、その project にアクセスできるようにします。 API の有効化の詳細については、Google Cloud API Console docsを参照してください。
Google Cloud Artifact Registry リポジトリを作成します。これは、Vertex で実行するイメージの保存先として使用します。詳細については、Google Cloud Artifact Registry documentationを参照してください。
ステージング用の GCS バケットを作成します。これは、Vertex AI がメタデータを保存するために使用します。ステージングバケットとして使用するには、このバケットを Vertex AI ワークロードと同じリージョンに配置する必要があります。ステージングとビルドコンテキストには同じバケットを使用できます。
サービスアカウントを作成します。このサービスアカウントには、Vertex AI ジョブを起動するために必要な権限が必要です。サービスアカウントへの権限の割り当ての詳細については、Google Cloud IAM documentationを参照してください。
次の表に示すように、サービスアカウントに Vertex ジョブを管理する権限を付与します。

Permission	Resource Scope	説明
`aiplatform.customJobs.create`	Specified Google Cloud Project	project 内で新しい機械学習ジョブを作成できます。
`aiplatform.customJobs.list`	Specified Google Cloud Project	project 内の機械学習ジョブを一覧表示できます。
`aiplatform.customJobs.get`	Specified Google Cloud Project	project 内の特定の機械学習ジョブに関する情報を取得できます。

Vertex AI ワークロードで標準以外のサービスアカウントのアイデンティティを使用する場合は、サービスアカウントの作成方法と必要な権限について、Vertex AI documentation を参照してください。W&B runs にカスタムサービスアカウントを選択するには、Launch queue 設定の spec.service_account フィールドを使用します。

Vertex AI 用のキューを設定する

Google Cloud の前提条件を満たしたら、次のステップは W&B Launch が Vertex AI ジョブの送信に使用するキュー設定を計画することです。Vertex AI リソースのキュー設定では、Vertex AI Python SDK の CustomJob コンストラクタと CustomJob の run method に渡す入力を指定します。リソースの設定は、spec キーと run キーの下に保存されます。

spec キーには、Vertex AI Python SDK の CustomJob コンストラクタの名前付き引数に対応する値が含まれます。
run キーには、Vertex AI Python SDK の CustomJob クラスの run method の名前付き引数に対応する値が含まれます。

実行環境のカスタマイズは、spec.worker_pool_specs リストで行います。worker pool spec は、ジョブを実行する worker のグループを定義します。デフォルトの設定に含まれる worker spec では、アクセラレータなしの n1-standard-4 マシン 1 台が指定されています。必要に応じて、マシンタイプ、アクセラレータタイプ、アクセラレータ数を変更できます。利用可能なマシンタイプとアクセラレータタイプの詳細については、Vertex AI documentation を参照してください。

キューを作成する

キュー設定の計画が済んだら、Vertex AI をコンピュートリソースとして使用するキューを W&B App で作成します。

Launch page にアクセスします。
Create Queue ボタンをクリックします。
キューを作成する Entity を選択します。
Name フィールドにキューの名前を入力します。
Resource として Google Cloud Vertex AI を選択します。

Configuration フィールドで、Configure a queue for Vertex AI で定義した Vertex AI CustomJob に関する情報を入力します。デフォルトでは、W&B によって次のような YAML および JSON のリクエストボディが自動入力されます。

spec:
  worker_pool_specs:
    - machine_spec:
        machine_type: n1-standard-4
        accelerator_type: ACCELERATOR_TYPE_UNSPECIFIED
        accelerator_count: 0
      replica_count: 1
      container_spec:
        image_uri: ${image_uri}
  staging_bucket: [STAGING-BUCKET]
run:
  restart_job_on_worker_restart: false

キューの設定が完了したら、Create Queue ボタンをクリックします。

最低限、次のフィールドを指定する必要があります。

spec.worker_pool_specs: 空ではないワーカープール仕様のリスト。
spec.staging_bucket: Vertex AI のアセットとメタデータをステージングするための GCS バケット。

Vertex AI のドキュメントの一部では、ワーカープール仕様のキーがすべてキャメルケースで記載されています。たとえば、workerPoolSpecs です。Vertex AI Python SDK では、これらのキーにスネークケースを使用します。たとえば、worker_pool_specs です。Launch queue 設定内のすべてのキーでは、スネークケースを使用してください。

Launch エージェントを設定する

キューを作成したら、Launch エージェントを設定してキューをポーリングし、ジョブを Vertex AI に送信します。Launch エージェントは、デフォルトでは ~/.config/wandb/launch-config.yaml にある設定ファイルで設定できます。

max_jobs: [N-CONCURRENT-JOBS]
queues:
  - [QUEUE-NAME]

Vertex AI で実行するイメージを Launch エージェントにビルドさせる場合は、エージェントの高度な設定を参照してください。

エージェントの権限を設定する

最後に、前提条件で作成したサービスアカウントとして動作するために必要な認証情報を Launch エージェントに付与します。このサービスアカウントとして認証するには、複数の方法があります。Workload Identity、ダウンロードしたサービスアカウント JSON、環境変数、Google Cloud Platform のコマンドラインツール、またはこれらの方法を組み合わせて認証できます。

Documentation Index

​前提条件

​Vertex AI 用のキューを設定する

​キューを作成する

​Launch エージェントを設定する

​エージェントの権限を設定する

前提条件

Vertex AI 用のキューを設定する

キューを作成する

Launch エージェントを設定する

エージェントの権限を設定する