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

Autopilot Standard

このページでは、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 を実行して最新のバージョンを取得する。
注: gcloud CLI がすでにインストールされている場合には、必ず compute/region と compute/zone のプロパティを設定してください。デフォルトのロケーションを設定することで、gcloud CLI のエラー（One of [--zone, --region] must be supplied: Please specify location など）を防止できます。

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

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

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

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

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
selector:
  matchLabels:
    app: maxengine-server
spec:
  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
```
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 連携を使用している場合は、アダプタが実行されているサービスアカウントにもモニタリング閲覧者のロールを付与する必要があります。
```
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 つのサービスアカウントをバインドします。
```
gcloud iam service-accounts add-iam-policy-binding \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:tpu-vm-gke-testing.svc.id.goog[default/prom-frontend-sa]" \
  jetstream-iam-sa@tpu-vm-gke-testing.iam.gserviceaccount.com \
&&
kubectl annotate serviceaccount \
  --namespace default \
  prom-frontend-sa \
  iam.gke.io/gcp-service-account=jetstream-iam-sa@tpu-vm-gke-testing.iam.gserviceaccount.com
```
3. 次に、Google Cloud サービスアカウントに monitoring.viewer ロールを付与します。
```
gcloud projects add-iam-policy-binding tpu-vm-gke-testing \
  --member=serviceAccount:jetstream-iam-sa@tpu-vm-gke-testing.iam.gserviceaccount.com \
  --role=roles/monitoring.viewer
```
4. 最後に、フロントエンドデプロイのサービスアカウントを新しいクラスタ内サービスアカウントに設定します。
```
kubectl set serviceaccount deployment frontend prom-frontend-sa
```
指標ベースの 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 自動スケーリングの詳細を確認する。