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 を実行して最新のバージョンを取得する。

• GKE の TPU で JetStream を使用して Gemma をサービングするでワークフローを理解する。JetStream デプロイ マニフェストで PROMETHEUS_PORT 引数が設定されていることを確認する。

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

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 自動スケーリングをご覧ください。

   • JetStream を個々の指標でスケーリングする場合は、カスタム指標 Stackdriver アダプタを使用します。
   • 複数の異なる指標で構成される式の値で JetStream をスケーリングする場合は、サードパーティの Prometheus アダプタを使用します。

   カスタム指標 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 クラスタで GKE 用 Workload Identity 連携が有効になっていて、GKE 用 Workload Identity 連携を使用している場合は、アダプタが実行されているサービス アカウントにもモニタリング閲覧者のロールを付与する必要があります。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 }'

次のステップ

• Cloud Monitoring の指標に基づいて Pod の自動スケーリングを最適化する方法を確認する。
• オープンソースの Kubernetes ドキュメントで、水平 Pod 自動スケーリングの詳細を確認する。