Tutorial: Set up W&B Launch on Kubernetes

W&B Launch を使用して、ML ワークロードを Kubernetes クラスターにプッシュできます。これにより、ML エンジニアは、Kubernetes で既に管理しているリソースを使用するために、W&B 内でシンプルなインターフェースを利用できます。

W&B は、W&B が管理する 公式 Launch agent イメージ を保持しており、Helm chart を使用してクラスターにデプロイできます。

W&B は Kaniko ビルダーを使用して、Launch エージェントが Kubernetes クラスターで Docker イメージを構築できるようにします。Launch エージェント用に Kaniko をセットアップする方法、またはジョブの構築をオフにして、事前構築済みの Docker イメージのみを使用する方法の詳細については、エージェントの詳細設定 を参照してください。

Kubernetes のキューを設定する

Kubernetes ターゲットリソースの Launch キュー設定は、Kubernetes Job spec または Kubernetes Custom Resource spec のいずれかに類似します。

Launch キューを作成するときに、Kubernetes ワークロードリソース 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 を使用するチュートリアルで、アプリケーションの例を参照してください。別のユースケースとして、W&B 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 キューにどのように表示されるかを示しています。

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

キューを作成する

Kubernetes をコンピューティングリソースとして使用するキューを W&B アプリケーションで作成します。

  1. Launch ページ に移動します。
  2. [キューを作成] ボタンをクリックします。
  3. キューを作成する Entity を選択します。
  4. [名前] フィールドにキューの名前を入力します。
  5. [リソース] として Kubernetes を選択します。
  6. [設定] フィールドで、前のセクションで設定した Kubernetes Job ワークフロー spec または Custom Resource spec を指定します。

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

W&B が提供する Helm chart を使用して、Launch エージェントを Kubernetes クラスターにデプロイします。values.yaml ファイル で Launch エージェントの振る舞いを制御します。

通常は Launch エージェント設定ファイル (~/.config/wandb/launch-config.yaml) で定義される内容を、values.yaml ファイルの launchConfig キー内に指定します。

たとえば、Kaniko Docker イメージビルダーを使用する EKS で Launch エージェントを実行できるようにする Launch エージェント設定があるとします。

queues:
	- <queue name>
max_jobs: <n 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 ファイル内では、次のようになります。

agent:
  labels: {}
  # W&B API key.
  apiKey: ''
  # Container image to use for the agent.
  image: wandb/launch-agent:latest
  # Image pull policy for agent image.
  imagePullPolicy: Always
  # Resources block for the agent spec.
  resources:
    limits:
      cpu: 1000m
      memory: 1Gi

# Namespace to deploy launch agent into
namespace: wandb

# W&B api url (Set yours here)
baseUrl: https://api.wandb.ai

# Additional target namespaces that the launch agent can deploy into
additionalTargetNamespaces:
  - default
  - wandb

# This should be set to the literal contents of your launch agent config.
launchConfig: |
  queues:
    - <queue name>
  max_jobs: <n concurrent jobs>
  environment:
    type: aws
    region: <aws-region>
  registry:
    type: ecr
    uri: <my-registry-uri>
  builder:
    type: kaniko
    build-context-store: <s3-bucket-uri>  

# The contents of a git credentials file. This will be stored in a k8s secret
# and mounted into the agent container. Set this if you want to clone private
# repos.
gitCreds: |

# Annotations for the wandb service account. Useful when setting up workload identity on gcp.
serviceAccount:
  annotations:
    iam.gke.io/gcp-service-account:
    azure.workload.identity/client-id:

# Set to access key for azure storage if using kaniko with azure.
azureStorageAccessKey: ''

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