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

このチュートリアルでは、ML エンジニアが W&B から直接トレーニングのワークロードを送信して管理できるように、クラスター管理者向けに Kubernetes クラスター上で W&B Launch を設定する方法を説明します。W&B Launch を使用すると、ML ワークロードを Kubernetes クラスターに送信でき、ML エンジニアは W&B 上から、すでに Kubernetes で管理しているリソースを利用できます。 W&B は、W&B が提供する Helm chart を使ってクラスターにデプロイできる、公式の Launch agent image を提供しています。 W&B は Kaniko builder を使用して、Launch agent が Kubernetes クラスター内で Docker イメージをビルドできるようにしています。Launch agent 用に Kaniko を設定する方法や、ジョブのビルドを無効にしてビルド済みの Docker イメージのみを使用する方法について詳しくは、Advanced agent setup を参照してください。

Helm をインストールし、W&B Launch agent の Helm chart を適用またはアップグレードするには、Kubernetes リソースの作成、更新、削除に必要な十分な権限を持つ kubectl でクラスターにアクセスできる必要があります。通常は、cluster-admin または同等の権限を持つカスタムロールが割り当てられたユーザーが必要です。

Kubernetes向けのキューを設定する

Launchキューは、各ジョブの実行時にエージェントが使用するKubernetesワークロードspecを定義します。Kubernetesのターゲットリソースに対するLaunchのキュー設定は、Kubernetes job spec または Kubernetes custom resource spec のいずれかに似たものになります。 Launchキューの作成時には、Kubernetesワークロードリソースのspecのあらゆる要素を制御できます。

Kubernetes job spec
Custom resource spec

spec:
  template:
    spec:
      containers:
        - env:
            - name: MY_ENV_VAR
              value: some-value
          resources:
            requests:
              cpu: 1000m
              memory: 1Gi
metadata:
  labels:
    queue: k8s-test
namespace: wandb

ユースケースによっては、CustomResource 定義を使用したい場合があります。たとえば、CustomResource 定義はマルチノードの分散トレーニングを実行する場合に便利です。アプリケーション例については、Volcanoを使用したマルチノードジョブでLaunchを使用するチュートリアルを参照してください。別のユースケースとして、LaunchをKubeflowと組み合わせて使用する場合があります。次のYAMLスニペットは、Kubeflowを使用するLaunchキュー設定の例を示しています。

kubernetes:
  kind: PyTorchJob
  spec:
    pytorchReplicaSpecs:
      Master:
        replicas: 1
        template:
          spec:
            containers:
              - name: pytorch
                image: '${image_uri}'
                imagePullPolicy: Always
        restartPolicy: Never
      Worker:
        replicas: 2
        template:
          spec:
            containers:
              - name: pytorch
                image: '${image_uri}'
                imagePullPolicy: Always
        restartPolicy: Never
    ttlSecondsAfterFinished: 600
  metadata:
    name: '${run_id}-pytorch-job'
  apiVersion: kubeflow.org/v1

セキュリティ上の理由から、これらが指定されていない場合、W&Bは次のリソースをLaunchキューに追加します。

securityContext
backOffLimit
ttlSecondsAfterFinished

次のYAMLスニペットは、これらの値がLaunchキュー内でどのように表示されるかを示しています。

example-spec.yaml

spec:
  template:
    backOffLimit: 0
    ttlSecondsAfterFinished: 60
    securityContext:
      allowPrivilegeEscalation: false
      capabilities:
        drop:
          - ALL
      seccompProfile:
        type: "RuntimeDefault"

キューを作成する

Kubernetes をコンピュートリソースとして使用するキューを W&B App で作成します。

Launch ページにアクセスします。
Create Queue ボタンをクリックします。
キューを作成する Entity を選択します。
Name フィールドにキューの名前を入力します。
Resource として Kubernetes を選択します。
Configuration フィールドに、Configure a queue for Kubernetes で設定した Kubernetes job workflow spec または custom resource spec を指定します。

Helm を使用して Launch agent を設定する

キューを用意したら、次にそのキューからジョブを取得してクラスター上で実行する Launch agent をデプロイします。W&B が提供する Helm chart を使用して、Launch agent を Kubernetes クラスターにデプロイします。Launch agent の動作は、values.yaml file で制御できます。 values.yaml file の launchConfig キーには、通常は Launch agent の設定ファイル (~/.config/wandb/launch-config.yaml) で定義する内容を指定します。たとえば、Kaniko Docker image builder を使用して EKS 上で Launch agent を実行できるようにする Launch agent 設定があるとします。[QUEUE-NAME]、[MAX-CONCURRENT-JOBS]、[MY-REGISTRY-URI]、[S3-BUCKET-URI] をご自身の値に置き換えてください。

launch-config.yaml

queues:
  - [QUEUE-NAME]
max_jobs: [MAX-CONCURRENT-JOBS]
environment:
  type: aws
  region: us-east-1
registry:
  type: ecr
  uri: [MY-REGISTRY-URI]
builder:
  type: kaniko
  build-context-store: [S3-BUCKET-URI]

values.yaml ファイルでは、次のようになります。[QUEUE-NAME]、[MAX-CONCURRENT-JOBS]、[AWS-REGION]、[MY-REGISTRY-URI]、および [S3-BUCKET-URI] は、ご自身の値に置き換えてください。

values.yaml

agent:
  labels: {}
  # W&B APIキー。
  apiKey: ''
  # エージェントに使用するコンテナーイメージ。
  image: wandb/launch-agent:latest
  # エージェントイメージのイメージプルポリシー。
  imagePullPolicy: Always
  # エージェント仕様のリソースブロック。
  resources:
    limits:
      cpu: 1000m
      memory: 1Gi

# Launch エージェントのデプロイ先の名前空間
namespace: wandb

# W&B API URL（ご自身の URL を設定してください）
baseUrl: https://api.wandb.ai

# Launch エージェントがデプロイできる追加のターゲット名前空間
additionalTargetNamespaces:
  - default
  - wandb

# Launch エージェントの設定ファイルの内容をそのまま指定してください。
launchConfig: |
  queues:
    - [QUEUE-NAME]
  max_jobs: [MAX-CONCURRENT-JOBS]
  environment:
    type: aws
    region: [AWS-REGION]
  registry:
    type: ecr
    uri: [MY-REGISTRY-URI]
  builder:
    type: kaniko
    build-context-store: [S3-BUCKET-URI]

# git 認証情報ファイルの内容。k8s シークレットに保存され、
# エージェントコンテナーにマウントされます。プライベートリポジトリを
# クローンする場合に設定してください。
gitCreds: |

# wandb サービスアカウントのアノテーション。GCP でワークロード アイデンティティを設定する際に役立ちます。
serviceAccount:
  annotations:
    iam.gke.io/gcp-service-account:
    azure.workload.identity/client-id:

# Azure で Kaniko を使用する場合は、Azure ストレージのアクセスキーを設定してください。
azureStorageAccessKey: ''

レジストリ、環境、必須のエージェント権限の詳細については、エージェントの高度な設定を参照してください。

​Kubernetes向けのキューを設定する

​キューを作成する

​Helm を使用して Launch agent を設定する

Kubernetes向けのキューを設定する

キューを作成する

Helm を使用して Launch agent を設定する