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

このチュートリアルでは、クラスター管理者向けに、Kubernetes クラスターに W&B Launch を設定する手順を説明します。設定が完了すると、ML エンジニアは W&B から直接トレーニングワークロードを送信、管理できるようになります。W&B Launch を使用すると、ML ワークロードを Kubernetes クラスターにプッシュできます。これにより、Kubernetes ですでに管理しているリソースを、ML エンジニアが W&B のインターフェースから直接使用できるようになります。 W&B は公式の Launch エージェントイメージを提供しています。このイメージは、W&B が管理する Helm チャートを使用してクラスターにデプロイできます。 W&B では、Kaniko ビルダーを使用して、Launch エージェントが Kubernetes クラスター内で Docker イメージをビルドできるようにしています。Launch エージェント用に Kaniko を設定する方法や、ジョブのビルドを無効にしてビルド済みの Docker イメージのみを使用する方法については、エージェントの高度なセットアップを参照してください。

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

Kubernetes 用のキューを設定する

Launch キューは、エージェントが各ジョブの実行に使用する Kubernetes ワークロードの spec を定義します。Kubernetes のターゲットリソース向けの Launch キュー設定は、Kubernetes のジョブ spec または Kubernetes カスタムリソース spec に似た形式になります。 Launch キューを作成する際に、Kubernetes ワークロードリソース spec のあらゆる項目を制御できます。

Kubernetes job spec
カスタムリソース 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 定義を使用すると便利な場合があります。たとえば、マルチノード分散トレーニングを実行する場合です。アプリケーションの例については、Volcano を使用してマルチノードジョブで Launch を使用するチュートリアルを参照してください。また、Kubeflow で Launch を使用する場合にも利用できます。次の 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"

キューを作成する

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

Launch ページにアクセスします。
Create Queue ボタンをクリックします。
キューを作成する Entity を選択します。
Name フィールドにキューの名前を入力します。
Resource に Kubernetes を選択します。
Configuration フィールドに、Kubernetes 用のキューを設定するで設定した Kubernetes ジョブのワークフロー spec またはカスタムリソース spec を入力します。

Helm を使用して Launch エージェントを設定する

キューを用意したら、次に、キューからジョブをプルしてクラスター上で実行する Launch エージェントをデプロイします。W&B が提供する Helm チャートを使用して、Kubernetes クラスターに Launch エージェントをデプロイします。Launch エージェントの動作は values.yaml ファイルで制御します。 values.yaml ファイルの launchConfig キーには、通常は Launch エージェントの設定ファイル (~/.config/wandb/launch-config.yaml) に定義する内容を指定します。たとえば、Kaniko Docker イメージビルダーを使用する Launch エージェントを EKS で実行するためのエージェント設定があるとします。[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
  # エージェント spec のリソースブロック。
  resources:
    limits:
      cpu: 1000m
      memory: 1Gi

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

# W&B API URL（使用する URL をここに設定します）
baseUrl: https://api.wandb.ai

# Launch エージェントがデプロイできる追加のターゲット namespace
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 エージェントを設定する

Kubernetes 用のキューを設定する

キューを作成する

Helm を使用して Launch エージェントを設定する