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

​
前提条件

  1. Vertex AI API を有効にした Google Cloud プロジェクトを作成するか、アクセスできるようにします。 API の有効化について詳しくは、Google Cloud API Console のドキュメントを参照してください。
  2. Google Cloud Artifact Registry リポジトリを作成します。Vertex で実行するイメージの保存先として使用します。詳しくは、Google Cloud Artifact Registry のドキュメントを参照してください。
  3. ステージング用の GCS バケットを作成します。これは Vertex AI がメタデータを保存するために使用します。このバケットをステージング用バケットとして使用するには、Vertex AI ワークロードと同じリージョンにある必要があります。ステージングとビルドコンテキストの両方に同じバケットを使用できます。
  4. サービスアカウントを作成します。Vertex AI ジョブを起動するために必要な権限を付与してください。サービスアカウントへの権限の割り当てについて詳しくは、Google Cloud IAM のドキュメントを参照してください。
  5. サービスアカウントに Vertex ジョブを管理する権限を付与します
PermissionResource ScopeDescription
aiplatform.customJobs.create指定した Google Cloud プロジェクトプロジェクト内で新しい機械学習ジョブを作成できます。
aiplatform.customJobs.list指定した Google Cloud プロジェクトプロジェクト内の機械学習ジョブを一覧表示できます。
aiplatform.customJobs.get指定した Google Cloud プロジェクトプロジェクト内の特定の機械学習ジョブに関する情報を取得できます。
Vertex AI ワークロードで標準以外のサービスアカウントのアイデンティティを使用する場合は、サービスアカウントの作成方法と必要な権限について Vertex AI のドキュメントを参照してください。Launch キュー設定の spec.service_account フィールドを使用すると、W&B run にカスタムサービスアカウントを選択できます。

​
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 で作成します。
  1. Launch ページにアクセスします。
  2. Create Queue ボタンをクリックします。
  3. キューを作成する Entity を選択します。
  4. 名 フィールドにキューの名前を入力します。
  5. Resource として Google Cloud Vertex AI を選択します。
  6. 設定 フィールドで、前のセクションで定義した 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: <REQUIRED>
run:
  restart_job_on_worker_restart: false
  1. キューを設定したら、Create Queue ボタンをクリックします。
最低限、以下を指定する必要があります。
  • spec.worker_pool_specs : 空でない worker pool specification のリスト。
  • spec.staging_bucket : Vertex AI のアセットとメタデータのステージングに使用する GCS bucket。
Vertex AI のドキュメントには、worker pool specification のキーがすべてキャメルケースで記載されているものがあります。たとえば  workerPoolSpecs です。Vertex AI Python SDK では、これらのキーにスネークケースを使用します。たとえば worker_pool_specs です。Launch キュー設定内のすべてのキーは、スネークケースを使用する必要があります。

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

Launch エージェントは、デフォルトでは ~/.config/wandb/launch-config.yaml にある設定ファイルで設定できます。 
max_jobs: <n-concurrent-jobs>
queues:
  - <queue-name>
Vertex AI で実行するイメージを Launch エージェントにビルドさせる場合は、エージェントの高度な設定を参照してください。

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

このサービスアカウントとして認証するには、複数の方法があります。Workload Identity、ダウンロードしたサービスアカウント JSON、環境変数、Google Cloud Platform のコマンドライン ツール、またはこれらの方法を組み合わせて実現できます。