TPU 上の LLM ワークロードの自動スケーリングを構成する

このページでは、GKE HorizontalPodAutoscaler（HPA）を使用して自動スケーリング インフラストラクチャを設定し、単一ホストの JetStream を使用して Gemma 大規模言語モデル（LLM）をデプロイする方法について説明します。

自動スケーリングの指標の選択の詳細については、GKE で TPU を使用して LLM ワークロードを自動スケーリングするためのベスト プラクティスをご覧ください。

始める前に

始める前に、次の作業が完了していることを確認してください。

  • Google Kubernetes Engine API を有効にする。
    • Google Kubernetes Engine API の有効化
  • このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。すでに gcloud CLI をインストールしている場合は、gcloud components update を実行して最新のバージョンを取得する。

指標を使用して自動スケーリングする

JetStream 推論サーバーまたは TPU パフォーマンス指標から出力されるワークロード固有のパフォーマンス指標を使用して、Pod の自動スケーリングを直接制御できます。

指標を使用して自動スケーリングを設定する手順は次のとおりです。

  1. JetStream サーバーから Cloud Monitoring に指標をエクスポートします。Google Cloud Managed Service for Prometheus を使用すると、Prometheus コレクタのデプロイと構成が簡素化されます。Google Cloud Managed Service for Prometheus は、GKE クラスタでデフォルトで有効になっています。手動で有効にすることもできます。

    次のマニフェストの例は、Google Cloud Managed Service for Prometheus に Pod から 15 秒間隔で指標をスクレイピングするよう指示する PodMonitoring リソース定義を設定する方法を示しています。

    サーバー指標をスクレイピングする必要がある場合は、次のマニフェストを使用します。サーバー指標では、5 秒間隔の取得間隔がサポートされています。

    apiVersion: monitoring.googleapis.com/v1
kind: PodMonitoring
metadata:
  name: jetstream-podmonitoring
spec:
  selector:
    matchLabels:
      app: maxengine-server
  endpoints:
  - interval: 15s
    path: "/"
    port: PROMETHEUS_PORT
  targetLabels:
    metadata:
    - pod
    - container
    - node

    TPU 指標をスクレイピングする必要がある場合は、次のマニフェストを使用します。システム指標では、15 秒という短い取得間隔がサポートされています。

    apiVersion: monitoring.googleapis.com/v1
kind: PodMonitoring
metadata:
  name: tpu-metrics-exporter
  namespace: kube-system
  labels:
    k8s-app: tpu-device-plugin
spec:
  endpoints:
    - port: 2112
      interval: 15s
  selector:
    matchLabels:
      k8s-app: tpu-device-plugin

  2. Metrics アダプタをインストールします。このアダプタにより、Monitoring にエクスポートされたサーバー指標が HPA コントローラに認識されるようになります。詳細については、Google Cloud Managed Service for Prometheus ドキュメントの水平 Pod 自動スケーリングをご覧ください。

    カスタム指標 Stackdriver アダプタ

    v0.13.1 のアダプタから、カスタム指標 Stackdriver アダプタは Google Cloud Managed Service for Prometheus からの指標のクエリをサポートしています。

    カスタム指標 Stackdriver アダプタをインストールするには、次の操作を行います。

    1. クラスタのマネージド コレクションを設定します。

    2. カスタム指標 Stackdriver アダプタをクラスタにインストールします。

      kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/k8s-stackdriver/master/custom-metrics-stackdriver-adapter/deploy/production/adapter_new_resource_model.yaml

    3. Kubernetes クラスタで Workload Identity Federation for GKE が有効になっていて、Workload Identity Federation for GKE を使用している場合は、アダプタが実行されているサービス アカウントにもモニタリング閲覧者のロールを付与する必要があります。 PROJECT_ID は実際のプロジェクト ID に置き換えます。

    export PROJECT_NUMBER=$(gcloud projects describe PROJECT_ID --format 'get(projectNumber)')
gcloud projects add-iam-policy-binding projects/PROJECT_ID \
  --role roles/monitoring.viewer \
  --member=principal://iam.googleapis.com/projects/$PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/custom-metrics/sa/custom-metrics-stackdriver-adapter

    Prometheus アダプタ

    prometheus-adapter を使用して Google Cloud Managed Service for Prometheus をスケーリングする場合は、次の点を考慮してください。

    • Prometheus API または UI を使用して Google Cloud Managed Service for Prometheus のクエリを実行する場合と同様に、クエリを Prometheus フロントエンド UI プロキシ経由でルーティングします。このフロントエンドは、後のステップでインストールします。
    • デフォルトでは、prometheus-adapter Deployment の prometheus-url 引数は --prometheus-url=http://frontend.default.svc:9090/ に設定されます。ここで、default はフロントエンドをデプロイした Namespace です。別の Namespace にフロントエンド サービスをデプロイした場合は、それに合わせてこの引数を構成します。
    • ルール構成ファイルの .seriesQuery フィールドでは、指標名に正規表現マッチャーを使用できません。代わりに、完全な指標名を指定します。

    アップストリームの Prometheus と比較して Google Cloud Managed Service for Prometheus ではデータが利用可能になるまでに少し時間がかかる場合があります。過剰な自動スケーリング ロジックを構成すると、望ましくない動作が発生する可能性があります。データの鮮度は保証されませんが、通常はネットワーク レイテンシを除いて Google Cloud Managed Service for Prometheus に送信した 3～7 秒後にデータはクエリの実行に使用できます。

    prometheus-adapter によって発行されるすべてのクエリは、グローバル スコープです。つまり、同じ名前の指標を出力する 2 つの Namespace にアプリケーションがある場合、その指標を使用する HPA 構成は両方のアプリケーションのデータを使用してスケーリングされます。間違ったデータを使用してスケーリングしないように、PromQL では常に namespace または cluster フィルタを使用します。

    prometheus-adapter とマネージド コレクションを使用して HPA 構成の例を設定するには、次の操作を行います。

    1. クラスタのマネージド コレクションを設定します。

    2. クラスタに Prometheus フロントエンド UI プロキシをデプロイします。次の内容のマニフェストを prometheus-frontend.yaml という名前で作成します。

        apiVersion: apps/v1
  kind: Deployment
  metadata:
    name: frontend
  spec:
    replicas: 2
    selector:
      matchLabels:
        app: frontend
    template:
      metadata:
        labels:
          app: frontend
      spec:
        automountServiceAccountToken: true
        affinity:
          nodeAffinity:
            requiredDuringSchedulingIgnoredDuringExecution:
              nodeSelectorTerms:
              - matchExpressions:
                - key: kubernetes.io/arch
                  operator: In
                  values:
                  - arm64
                  - amd64
                - key: kubernetes.io/os
                  operator: In
                  values:
                  - linux
        containers:
        - name: frontend
          image: gke.gcr.io/prometheus-engine/frontend:v0.8.0-gke.4
          args:
          - "--web.listen-address=:9090"
          - "--query.project-id=PROJECT_ID"
          ports:
          - name: web
            containerPort: 9090
          readinessProbe:
            httpGet:
              path: /-/ready
              port: web
          securityContext:
            allowPrivilegeEscalation: false
            capabilities:
              drop:
              - all
            privileged: false
            runAsGroup: 1000
            runAsNonRoot: true
            runAsUser: 1000
          livenessProbe:
            httpGet:
              path: /-/healthy
              port: web
  ---
  apiVersion: v1
  kind: Service
  metadata:
    name: prometheus
  spec:
    clusterIP: None
    selector:
      app: frontend
    ports:
    - name: web
      port: 9090

      マニフェストを適用します。

      kubectl apply -f prometheus-frontend.yaml

    3. prometheus-community/prometheus-adapter Helm チャートをインストールして、クラスタに prometheus-adapter がインストールされていることを確認します。次の values.yaml ファイルを作成します。

      rules:
  default: false
  external:
  - seriesQuery: 'jetstream_prefill_backlog_size'
    resources:
      template: <<.Resource>>
    name:
      matches: ""
      as: "jetstream_prefill_backlog_size"
    metricsQuery: avg(<<.Series>>{<<.LabelMatchers>>,cluster="CLUSTER_NAME"})
  - seriesQuery: 'jetstream_slots_used_percentage'
    resources:
      template: <<.Resource>>
    name:
      matches: ""
      as: "jetstream_slots_used_percentage"
    metricsQuery: avg(<<.Series>>{<<.LabelMatchers>>,cluster="CLUSTER_NAME"})
  - seriesQuery: 'memory_used'
    resources:
      template: <<.Resource>>
    name:
      matches: ""
      as: "memory_used_percentage"
    metricsQuery: avg(memory_used{cluster="CLUSTER_NAME",exported_namespace="default",container="jetstream-http"}) / avg(memory_total{cluster="CLUSTER_NAME",exported_namespace="default",container="jetstream-http"})

      このファイルを Helm チャートのデプロイ用の値ファイルとして使用します。

      helm repo add prometheus-community https://prometheus-community.github.io/helm-charts && helm repo update && helm install example-release prometheus-community/prometheus-adapter -f values.yaml

    GKE 用 Workload Identity 連携を使用する場合は、次のコマンドを実行し、サービス アカウントを構成して認可すること必要があります。

    1. まず、クラスタ内サービス アカウントと Google Cloud サービス アカウントを作成します。

      gcloud iam service-accounts create prom-frontend-sa && kubectl create sa prom-frontend-sa

    2. この 2 つのサービス アカウントをバインドします。PROJECT_ID は、プロジェクト ID に置き換えてください。

      gcloud iam service-accounts add-iam-policy-binding \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:PROJECT_ID.svc.id.goog[default/prom-frontend-sa]" \
  jetstream-iam-sa@PROJECT_ID.iam.gserviceaccount.com \
&&
kubectl annotate serviceaccount \
  --namespace default \
  prom-frontend-sa \
  iam.gke.io/gcp-service-account=jetstream-iam-sa@PROJECT_ID.iam.gserviceaccount.com

    3. 次に、Google Cloud サービス アカウントに monitoring.viewer ロールを付与します。

      gcloud projects add-iam-policy-binding PROJECT_ID \
  --member=serviceAccount:jetstream-iam-sa@PROJECT_ID.iam.gserviceaccount.com \
  --role=roles/monitoring.viewer

    4. 最後に、フロントエンド デプロイのサービス アカウントを新しいクラスタ内サービス アカウントに設定します。

      kubectl set serviceaccount deployment frontend prom-frontend-sa

  3. 指標ベースの HPA リソースを設定します。優先するサーバー指標に基づく HPA リソースをデプロイします。詳細については、Google Cloud Managed Service for Prometheus ドキュメントの水平 Pod 自動スケーリングをご覧ください。 具体的な HPA 構成は、指標のタイプ（サーバーまたは TPU）とインストールされている指標アダプタによって異なります。

    すべての HPA 構成で必要な値がいくつかあります。HPA リソースを作成するには、これらの値を設定する必要があります。

    • MIN_REPLICAS: 許可される JetStream Pod レプリカの最小数。JetStream をデプロイする手順で JetStream のデプロイ マニフェストを変更しない場合は、これを 1 に設定することをおすすめします。
    • MAX_REPLICAS: 許容される JetStream Pod レプリカの最大数。サンプルの JetStream デプロイでは、レプリカごとに 8 個のチップが必要です。ノードプールには 16 個のチップが含まれています。スケールアップのレイテンシを低く抑えるには、これを 2 に設定します。値を大きくすると、クラスタ オートスケーラーがトリガーされ、ノードプールに新しいノードが作成されるため、スケールアップのレイテンシが増加します。

    • TARGET: JetStream インスタンス全体でのこの指標のターゲット平均。この値からレプリカ数がどのように決定されるかについては、Kubernetes のドキュメントで自動スケーリングの説明をご覧ください。

    カスタム指標 Stackdriver アダプタ

    カスタム指標 Stackdriver アダプタを使用すると、Pod 全体での Google Cloud Managed Service for Prometheus からの指標クエリの平均値でワークロードをスケーリングできます。カスタム指標 Stackdriver アダプタを使用する場合は、jetstream_prefill_backlog_size サーバー指標、jetstream_slots_used_percentage サーバー指標、memory_used TPU 指標を使用してスケーリングすることをおすすめします。

    サーバー指標でスケーリングする HPA マニフェストを作成するには、次の hpa.yaml ファイルを作成します。

    apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: jetstream-hpa
  namespace: default
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: maxengine-server
  minReplicas: MIN_REPLICAS
  maxReplicas: MAX_REPLICAS
  metrics:
  - type: Pods
    pods:
      metric:
        name: prometheus.googleapis.com|jetstream_METRIC|gauge
      target:
        type: AverageValue
        averageValue: TARGET

    カスタム指標 Stackdriver アダプタを TPU 指標で使用する場合は、スケーリングに kubernetes.io|node|accelerator|memory_used 指標のみを使用することをおすすめします。この指標でスケーリングする HPA マニフェストを作成するには、次の hpa.yaml ファイルを作成します。

    apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: jetstream-hpa
  namespace: default
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: maxengine-server
  minReplicas: MIN_REPLICAS
  maxReplicas: MAX_REPLICAS
  metrics:
  - type: External
    external:
      metric:
        name: prometheus.googleapis.com|memory_used|gauge
        selector:
          matchLabels:
            metric.labels.container: jetstream-http
            metric.labels.exported_namespace: default
      target:
        type: AverageValue
        averageValue: TARGET

    Prometheus アダプタ

    Prometheus アダプタを使用すると、Google Cloud Managed Service for Prometheus の PromQL クエリの値を使用してワークロードをスケーリングできます。Pod 全体の平均値を表す jetstream_prefill_backlog_size と jetstream_slots_used_percentage のサーバー指標をすでに定義しています。

    サーバー指標でスケーリングする HPA マニフェストを作成するには、次の hpa.yaml ファイルを作成します。

    apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: jetstream-hpa
  namespace: default
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: maxengine-server
  minReplicas: MIN_REPLICAS
  maxReplicas: MAX_REPLICAS
  metrics:
  - type: External
    external:
      metric:
        name: jetstream_METRIC
      target:
        type: AverageValue
        averageValue: TARGET

    TPU 指標を使用してスケーリングする HPA のマニフェストを作成するには、prometheus-adapter helm 値ファイルで定義されている memory_used_percentage のみを使用することをおすすめします。memory_used_percentage は次の PromQL クエリに付けられた名前で、これは、すべてのアクセラレータで使用されている現在の平均メモリを反映します。

    avg(kubernetes_io:node_accelerator_memory_used{cluster_name="CLUSTER_NAME"}) / avg(kubernetes_io:node_accelerator_memory_total{cluster_name="CLUSTER_NAME"})

    memory_used_percentage でスケーリングする HPA マニフェストを作成するには、次の hpa.yaml ファイルを作成します。

    apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: jetstream-hpa
  namespace: default
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: maxengine-server
  minReplicas: MIN_REPLICAS
  maxReplicas: MAX_REPLICAS
  metrics:
  - type: External
    external:
      metric:
        name: memory_used_percentage
      target:
        type: AverageValue
        averageValue: TARGET

複数の指標を使用してスケーリングする

複数の指標に基づいてスケーリングを構成することもできます。複数の指標を使用してレプリカ数を決定する方法については、Kubernetes のドキュメントで自動スケーリングに関する説明をご覧ください。このタイプの HPA マニフェストを作成するには、各 HPA リソースの spec.metrics フィールドのすべてのエントリを 1 つの HPA リソースに収集します。次のスニペットは、HPA リソースをバンドルする方法の例を示しています。

apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: jetstream-hpa-multiple-metrics
  namespace: default
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: maxengine-server
  minReplicas: MIN_REPLICAS
  maxReplicas: MAX_REPLICAS
  metrics:
  - type: Pods
    pods:
      metric:
        name: jetstream_METRIC
      target:
        type: AverageValue
      averageValue: JETSTREAM_METRIC_TARGET
  - type: External
    external:
      metric:
        name: memory_used_percentage
      target:
        type: AverageValue
      averageValue: EXTERNAL_METRIC_TARGET

自動スケーリングをモニタリングしてテストする

HPA 構成に基づいて JetStream ワークロードがどのようにスケーリングされるかを確認できます。

レプリカ数をリアルタイムで確認するには、次のコマンドを実行します。

kubectl get hpa --watch

このコマンドの出力は次のようになります。

NAME            REFERENCE                     TARGETS      MINPODS   MAXPODS   REPLICAS   AGE
jetstream-hpa   Deployment/maxengine-server   0/10 (avg)   1         2         1          1m

HPA のスケーリング機能をテストするには、次のコマンドを使用して、短時間に 100 件のリクエストをモデル エンドポイントに送信します。これにより、使用可能なデコード スロットが使い果たされ、プリフィル キューでリクエストのバックログが発生し、HPA がトリガーされてモデルのデプロイサイズが増加します。

seq 100 | xargs -P 100 -n 1 curl --request POST --header "Content-type: application/json" -s localhost:8000/generate --data '{ "prompt": "Can you provide a comprehensive and detailed overview of the history and development of artificial intelligence.", "max_tokens": 200 }'

次のステップ