为 TPU 上的 LLM 工作负载配置自动扩缩

Autopilot Standard

本页面介绍如何使用 GKE Pod 横向自动扩缩器 (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。

熟悉通过 JetStream 使用 GKE 中的 TPU 应用 Gemma 中的工作流并完成相应工作。确保已在 JetStream 部署清单中设置 PROMETHEUS_PORT 参数。

使用指标进行自动扩缩

您可以使用 JetStream 推理服务器发出的特定于工作负载的性能指标或 TPU 性能指标来引导 Pod 进行自动扩缩。

如需使用指标设置自动扩缩，请按照以下步骤操作：

将指标从 JetStream 服务器导出到 Cloud Monitoring。您将使用 Google Cloud Managed Service for Prometheus，该功能可简化 Prometheus 收集器的部署和配置。默认情况下，GKE 集群中会启用 Google Cloud Managed Service for Prometheus；您也可以手动启用该功能。

以下示例清单展示了如何设置 PodMonitoring 资源定义，以引导 Google Cloud Managed Service for Prometheus 以 15 秒的周期性间隔从 Pod 爬取指标：

如果您需要爬取服务器指标，请使用以下清单。使用服务器指标时，支持爬取间隔为 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
```

安装指标适配器。此适配器可使您导出到 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 适配器，请执行以下操作：

在集群中设置代管式收集。

在集群中安装自定义指标 Stackdriver 适配器。

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

如果您在 Kubernetes 集群上启用了适用于 GKE 的工作负载身份联合，并且使用适用于 GKE 的工作负载身份联合，则还必须向用于运行适配器的服务账号授予“Monitoring Viewer”角色。请将 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 或界面查询 Google Cloud Managed Service for Prometheus 时一样，请通过 Prometheus 前端界面代理路由查询。此前端将在后续步骤中安装。
默认情况下，prometheus-adapter 部署的 prometheus-url 参数设置为 --prometheus-url=http://frontend.default.svc:9090/，其中 default 是您部署前端的命名空间。如果您在其他命名空间中部署了前端，请相应地配置此参数。
在规则配置的 .seriesQuery 字段中，您不能对指标名称使用正则表达式 (regex) 匹配器。请改为完全指定指标名称。

由于与上游 Prometheus 相比，在 Google Cloud Managed Service for Prometheus 中使数据可供使用所需的时间略长，因此配置过于激增的自动扩缩逻辑可能会导致不希望的行为。虽然无法保证数据新鲜度，但数据在发送到 Google Cloud Managed Service for Prometheus 后，通常在 3-7 秒内可用于查询（不包括任何网络延迟时间）。

prometheus-adapter 发出的所有查询都是全局性的。这意味着，如果两个命名空间中的应用发出具有相同名称的指标，则使用该指标的 HPA 配置会使用这两个应用中的数据进行扩缩。如需避免使用错误的数据进行扩缩，请始终在 PromQL 中使用 namespace 或 cluster 过滤条件。

如需使用 prometheus-adapter 和托管式集合设置示例 HPA 配置，请按以下步骤操作：

在集群中设置代管式收集。

在集群中部署 Prometheus 前端界面代理。创建以下名为 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

通过安装 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 的工作负载身份联合，则还需要通过运行以下命令来配置服务账号并向其授权：

首先，创建集群内服务账号和 Google Cloud 服务账号：

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

然后，绑定这两个服务账号，请务必将 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

接下来，向 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

最后，将前端部署服务账号设置为新的集群内服务账号：
```
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 适配器支持根据 Google Cloud Managed Service for Prometheus 在所有 Pod 间进行的各个指标查询的平均值来扩缩工作负载。使用自定义指标 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 查询的值来扩缩工作负载。您之前定义了 jetstream_prefill_backlog_size 和 jetstream_slots_used_percentage 服务器指标，它们表示所有 Pod 间的平均值。

如需创建 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
```
如需创建 HPA 清单以便根据 TPU 指标进行扩缩，我们建议仅使用在 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"})
```
如需创建 HPA 清单以便根据 memory_used_percentage 进行扩缩，请创建以下 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 字段中的所有条目收集到单个 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

监控和测试自动扩缩

您可以观察 JetStream 工作负载如何基于 HPA 配置进行扩缩。

如需实时观察副本数量，请运行以下命令：

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 横向自动扩缩。