GKE またはセルフマネージド Kubernetes を構成する

ML Diagnostics プラットフォームは、Google Kubernetes Engine（GKE）またはセルフマネージド Kubernetes で実行される ML ワークロード用に構成できます。

ML 診断には、次の一般提供機能が含まれています。

  • ワークロードの自動モニタリングとシステム指標の収集: GKE でワークロードの自動モニタリングとシステム指標の収集を行う場合、コードを計測したり、追加の構成を行ったりする必要はありません。ワークロード モニタリングとシステム指標の収集はデフォルトで有効になっているため、このガイドの手順を行う必要はありません。ワークロード モニタリングは、JobSet と Job の GKE サービスの種類をサポートしており、GKE バージョン 1.36.4-gke.1247000 以降と互換性があります。
  • ML Diagnostics SDK の指標とプログラムによるプロファイリング: ML Diagnostics SDK を使用して指標を収集し、プログラムによるプロファイリングを実行する場合は、injection-webhook Helm チャートをインストールし、ワークロードにラベルを付け、ML Diagnostics SDK を使用します。GKE では、サポートされている GKE バージョンを使用している場合は、1 つの Google Cloud CLI コマンドを使用できます。または、手動インストールを実行することもできます。セルフマネージド Kubernetes では、Webhook を手動でインストールする必要があります。
  • オンデマンド プロファイリング: オンデマンド プロファイリングを使用する場合は、connection-operator Helm チャートをインストールします。GKE では、サポートされている GKE バージョンを使用している場合は、1 つの gcloud CLI コマンドを使用できます。または、手動インストールを実行することもできます。セルフマネージド Kubernetes では、オペレーターを手動でインストールする必要があります。

ML 診断には、次のプレビュー機能も含まれています。

  • Diagon DaemonSet プロファイリング: GKE ワークロードでオンデマンド プロファイリング、継続的なプロファイリング、ワークロード モニタリングからの自動プロファイル キャプチャを使用する場合は、GKE TPU クラスタに Diagon DaemonSet をデプロイします。この DaemonSet は、セルフマネージド Kubernetes クラスタでも動作します。このクラスタでは、オンデマンド プロファイリングと継続的なプロファイリングのみがサポートされています。セルフマネージド Kubernetes では、ワークロード モニタリングはサポートされていません。

Diagon DaemonSet をデプロイする

GKE またはセルフマネージド Kubernetes クラスタでコンテナ化されたワークロードを実行する場合は、Diagon デーモンを DaemonSet としてデプロイできます。これにより、デーモンのコピーがクラスタ内のすべての（または選択した）GPU ノードと TPU ノードで実行されます。

コンテナ イメージ リポジトリ

Diagon デーモンは Docker イメージとして公開されており、次の Artifact Registry パスから入手できます。

us-central1-docker.pkg.dev/hypercomputediagon-release/daemon-docker/mldiagnostics-daemon:latest

デプロイ手順を実行する

Diagon DaemonSet をクラスタにデプロイする手順は次のとおりです。

  1. DaemonSet 構成を作成する: 次の YAML 構成を diagon-daemonset.yaml として保存します。この構成では、Diagon デーモンが特権モードで実行され、基盤となるホスト GPU デバイスと TPU デバイスから診断指標とプロファイリング指標を収集できます。

    # diagon-daemonset.yaml
apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: diagon-daemon
  namespace: diagon
  labels:
    app.kubernetes.io/name: diagon-daemon
    app.kubernetes.io/part-of: diagon
spec:
  selector:
    matchLabels:
      name: diagon-daemon
  template:
    metadata:
      labels:
        name: diagon-daemon
    spec:
      hostNetwork: true
      dnsPolicy: ClusterFirstWithHostNet
      serviceAccountName: SERVICE_ACCOUNT_NAME
      terminationGracePeriodSeconds: 30
      tolerations:
      - operator: Exists
        effect: NoSchedule
      - operator: Exists
        effect: NoExecute
      containers:
      - name: diagon-daemon
        image: us-central1-docker.pkg.dev/hypercomputediagon-release/daemon-docker/mldiagnostics-daemon:latest
        imagePullPolicy: Always
        command:
        - /mldiagnostics_extension_files/bin/mldiagnostics_extension
        - -standalone=true
        - -channel=hypercomputecluster.googleapis.com/diagnostics
        - -alsologtostderr=true
        securityContext:
          privileged: true
        resources:
          requests:
            cpu: 10m
            memory: 32Mi
          limits:
            memory: 256Mi  # No CPU limit.

  2. DaemonSet をクラスタにデプロイする: 次のコマンドを実行して、アクティブな Kubernetes コンテキストにデーモンをデプロイします。

    kubectl apply -f diagon-daemonset.yaml

  3. デプロイとデーモンの健全性を確認する: 管理ワークステーションまたは Cloud Shell で、すべてのデーモン Pod がワーカーノードにスケジュールされ、Running 状態に達していることを確認します。

    kubectl get pods -n diagon -l name=diagon-daemon -o wide

    ストリーミング ログを調べて、 Google Cloud ML 診断コントロール プレーンへの接続とアクティブなハートビートを確認します。

    kubectl logs -n diagon -l name=diagon-daemon --tail=50

    出力には、次の想定されるログ署名が含まれます。

    Stream established.
Received successful DaemonRegistrationResponse from backend
Sending DaemonHeartbeatRequest

  4. DaemonSet イメージを更新する: 新しいバージョンのデーモンに更新するか、ロールバックするには、YAML ファイルのイメージタグを更新するか、次のコマンドを直接実行します。

    kubectl set image daemonset/diagon-daemon diagon-daemon=us-docker.pkg.dev/hypercomputediagon-release/daemon/diagon-daemon:v2.1.0 --namespace=kube-system

    デプロイまたは更新の進行状況を確認するには、ロールアウト ステータスを確認します。

    kubectl rollout status daemonset/diagon-daemon --namespace=kube-system

オペレーターと Webhook を設定する

オペレーターと Webhook の設定には、次のオプションがあります。

  • 1.35.0-gke.3065000 より後のバージョンの GKE を使用している場合は、1 つの gcloud CLI コマンド、 Google Cloud コンソール、または Terraform を使用して、ML 診断用に GKE クラスタを設定できます。

  • 1.35.0-gke.3065000 より前の GKE バージョンとセルフマネージド Kubernetes の場合は、cert-manager、injection-webhook、connection-operator アーティファクトをインストールするようにクラスタを手動で構成する必要があります。詳細については、手動インストールをご覧ください。

Diagon DaemonSet とともに、GKE クラスタにオペレーターと Webhook をデプロイします。オペレーターと Webhook をすでにデプロイしている場合は、削除する必要はありません。代わりに、オペレーターはオンデマンド プロファイル キャプチャのフォールバック パスとして機能し、SDK を使用するには引き続き Webhook が必要です。

GKE ではなくセルフマネージド Kubernetes 環境を使用する場合は、SDK の Webhook をインストールし、Diagon DaemonSet をデプロイします。セルフマネージド Kubernetes 環境では、ワークロード モニタリングは機能しません。SDK 指標、プログラムによるプロファイリング、オンデマンド プロファイリング、継続的なプロファイリングのみがサポートされています。

1.35.0-gke.3065000 より後の GKE バージョンの場合は、次のいずれかの方法で、必要な ML 診断コンポーネント（connection-operator と injection-webhook の両方）を GKE クラスタにデプロイします。

gcloud

新しい GKE クラスタの場合:

gcloud beta container clusters create CLUSTER_NAME --enable-managed-mldiagnostics

既存の GKE クラスタの場合:

gcloud beta container clusters update CLUSTER_NAME --enable-managed-mldiagnostics

ML 診断を無効にするには、次のコマンドを使用します。

gcloud beta container clusters update CLUSTER_NAME --no-enable-managed-mldiagnostics

ML 診断用に GKE クラスタを設定する gcloud CLI コマンドの詳細については、次の API リファレンス ページの enable-managed-mldiagnostics フラグをご覧ください。

コンソール

  • 新しい GKE クラスタの場合は、[Feature Manager] > [Managed Machine Learning Diagnostics] に移動します。

    GKE マネージド ML 診断に移動

  • 既存の GKE クラスタの場合は、[クラスタ] に移動し、クラスタの名前を選択して [編集] に移動し、[機能] の [マネージド機械学習診断] を編集します。

Terraform

Terraform を使用して GKE インフラストラクチャをプロビジョニングする場合は、terraform-provider-google-beta プロバイダ（バージョン v7.28.0 以降）または terraform-google-modules/kubernetes-engine/google モジュール（バージョン v45.0.0 以降）を使用できます。

Google Cloud Terraform プロバイダの使用

google_container_cluster リソースを使用して新しいクラスタでマネージド ML 診断を有効にするには、google-beta プロバイダ（バージョン v7.28.0 以降）を使用して、enabled = true を含む managed_machine_learning_diagnostics_config ブロックを追加します。

terraform {
  required_providers {
    google-beta = {
      source  = "hashicorp/google-beta"
      version = ">= 7.28.0"
    }
  }
}

resource "google_container_cluster" "primary" {
  provider = google-beta
  name     = "CLUSTER_NAME"
  location = "LOCATION"

  # ... other cluster configuration ...

  managed_machine_learning_diagnostics_config {
    enabled = true
  }
}

Terraform リソースを使用してクラスタでマネージド ML 診断を無効にするには、managed_machine_learning_diagnostics_config ブロックで enabled = false を設定します。

resource "google_container_cluster" "primary" {
  provider = google-beta
  name     = "CLUSTER_NAME"
  location = "LOCATION"

  # ... other cluster configuration ...

  managed_machine_learning_diagnostics_config {
    enabled = false
  }
}

GKE Terraform モジュールの使用

terraform-google-modules/kubernetes-engine/google モジュール（beta-public-cluster サブモジュールなど）を使用してクラスタでマネージド ML 診断を有効にするには、enable_managed_machine_learning_diagnostics = true を設定します。

module "gke" {
  source  = "terraform-google-modules/kubernetes-engine/google//modules/beta-public-cluster"
  version = ">= 45.0"

  project_id = "PROJECT_ID"
  name       = "CLUSTER_NAME"
  region     = "LOCATION"

  # ... other cluster configuration ...

  enable_managed_machine_learning_diagnostics = true
}

Terraform モジュールを使用してクラスタでマネージド ML 診断を無効にするには、enable_managed_machine_learning_diagnostics = false を設定します。

module "gke" {
  source  = "terraform-google-modules/kubernetes-engine/google//modules/beta-public-cluster"
  version = ">= 45.0"

  # ... other cluster configuration ...

  enable_managed_machine_learning_diagnostics = false
}

構成を更新したら、変更を適用します。

terraform apply

GKE で Terraform を使用する方法については、GKE での Terraform のサポートをご覧ください。

手動インストール

1.35.0-gke.3065000 より前の GKE バージョンの場合は、GKE クラスタを手動で構成して、次のものをインストールする必要があります。

  • cert-manager: injection-webhook の前提条件。
  • injection-webhook: 必要なメタデータを SDK に提供します。JobSet、RayJob、LeaderWorkerSet などの一般的な ML Kubernetes ワークロードをサポートしています。
  • connection-operator: GKE でのオンデマンド プロファイリング。connection-operator と injection-webhook を GKE クラスタにデプロイすると、オンデマンド キャプチャをトリガーしたときに、プロファイリング サーバーが実行されているターゲット Pod へのプロファイリング リクエストが初期化されます。

GKE クラスタの設定の詳細については、GKE クラスタを構成するをご覧ください。

Cert-manager

cert-manager アドオンは、クラスタの証明書コントローラとして機能し、アプリケーションの安全性を確保し、証明書が意図せず期限切れにならないようにします。

Helm を使用して、次のものをインストールします。

helm repo add jetstack https://charts.jetstack.io
helm repo update

helm install \
  cert-manager jetstack/cert-manager \
  --namespace cert-manager \
  --create-namespace \
  --version v1.13.0 \
  --set installCRDs=true \
  --set global.leaderElection.namespace=cert-manager \
  --timeout 10m

injection-webhook

injection-webhook Helm チャートは、メタデータを SDK に渡します。helm upgrade --install を使用して、初めてインストールするか、既存のインストールをアップグレードします。

helm upgrade --install mldiagnostics-injection-webhook \
  --namespace=gke-mldiagnostics \
  --create-namespace \
  --version 0.25.0 \
  oci://us-docker.pkg.dev/ai-on-gke/mldiagnostics-webhook-and-operator-helm/mldiagnostics-injection-webhook

connection-operator

connection-operator Helm チャートを使用すると、GKE でオンデマンド プロファイリングを有効にできます。次の表を使用して、適切な mldiagnostics-connection-operator バージョンを確認してください。

JAX のバージョン Helm チャートのバージョン
0.8.x 0.24.0
0.9.x+ 0.24.0+

Helm を使用して、必要なバージョンをインストールします。

JAX 0.8.x の場合:

helm upgrade --install mldiagnostics-connection-operator \
  --namespace=gke-mldiagnostics \
  --create-namespace \
  --version 0.24.0 \
  oci://us-docker.pkg.dev/ai-on-gke/mldiagnostics-webhook-and-operator-helm/mldiagnostics-connection-operator \
  --set 'mldiagnosticsConnectionOperator.controller.args={--metrics-bind-address=:8443,--health-probe-bind-address=:8081,--sidecar-timeout=65m,--disable-hostname-override}'

JAX 0.9.x+ の場合:

helm upgrade --install mldiagnostics-connection-operator \
  --namespace=gke-mldiagnostics \
  --create-namespace \
  --version 0.24.0 \
  oci://us-docker.pkg.dev/ai-on-gke/mldiagnostics-webhook-and-operator-helm/mldiagnostics-connection-operator

ワークロードにラベルを付ける

プログラムによるプロファイリングの場合は、injection-webhook をトリガーしてメタデータを Pod に挿入する必要があります。ワークロードをデプロイする前に、ワークロードまたはその Namespace に managed-mldiagnostics-gke=true のラベルを付けます。

  • ワークロードにラベルを付ける: JobSet、LeaderWorkerSet、RayJob ワークロードにラベルを付けます。これにより、その特定のワークロードの Webhook が有効になります。JobSet ワークロードの例を次に示します。

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: single-host-tpu-v3-jobset2
  namespace: default
  labels:
    managed-mldiagnostics-gke: "true"

  • Namespace にラベルを付ける: これにより、その Namespace 内のすべての JobSet、LeaderWorkerSet、RayJob ワークロードで Webhook が有効になります。

    kubectl create namespace ai-workloads
kubectl label namespace ai-workloads managed-mldiagnostics-gke=true