Autoscaling für LLM-Arbeitslasten auf TPUs konfigurieren

Konfigurieren Sie das Autoscaling für Ihre LLM-Arbeitslasten (Large Language Model) in Google Kubernetes Engine (GKE), die auf TPUs ausgeführt werden, mit dem horizontalen Pod-Autoscaler (HPA) und JetStream mit einem einzelnen Host. So können Sie schwankende LLM-Inferenzarbeitslasten effizient verwalten, die Ressourcennutzung optimieren und eine konsistente Leistung für Ihre Anwendungen gewährleisten.

HPA passt die Anzahl der Pod-Replikate automatisch basierend auf der beobachteten CPU-Auslastung oder anderen benutzerdefinierten Messwerten an. JetStream erleichtert die Bereitstellung von LLMs und ermöglicht eine effektive Skalierung. Best Practices für die Auswahl von Messwerten für das Autoscaling finden Sie unter Best Practices für das Autoscaling von LLM-Arbeitslasten mit TPUs in GKE.

Hinweis

Führen Sie die folgenden Aufgaben aus, bevor Sie beginnen:

  • Aktivieren Sie die Google Kubernetes Engine API.
    • Google Kubernetes Engine API aktivieren
  • Wenn Sie die Google Cloud CLI für diesen Task verwenden möchten, installieren und dann initialisieren Sie die gcloud CLI. Wenn Sie die gcloud CLI bereits installiert haben, rufen Sie die neueste Version mit dem gcloud components update Befehl ab. Ältere gcloud CLI-Versionen unterstützen möglicherweise nicht die Ausführung der Befehle in diesem Dokument.

Autoscaling mit Messwerten

Sie können die arbeitslastspezifischen Leistungsmesswerte, die vom JetStream-Inferenzserver ausgegeben werden, oder TPU-Leistungsmesswerte verwenden, um das Autoscaling für Ihre Pods auszurichten.

So richten Sie das Autoscaling mit Messwerten ein:

  1. Exportieren Sie die Messwerte vom JetStream-Server nach Cloud Monitoring. Verwenden Sie Google Cloud Managed Service for Prometheus, um die Bereitstellung und Konfiguration Ihres Prometheus-Collectors zu vereinfachen. Google Cloud Managed Service for Prometheus ist in Ihrem GKE-Cluster standardmäßig aktiviert. Sie können es auch manuell aktivieren.

    Das folgende Beispielmanifest zeigt, wie Sie Ihre PodMonitoring-Ressourcendefinitionen einrichten, damit Google Cloud Managed Service for Prometheus Messwerte in regelmäßigen Abständen von 15 Sekunden aus Ihren Pods extrahiert:

    Wenn Sie Servermesswerte extrahieren müssen, verwenden Sie das folgende Manifest. Bei Servermesswerten werden Extraktionsintervalle von bis zu 5 Sekunden unterstützt.

    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

    Wenn Sie TPU-Messwerte extrahieren müssen, verwenden Sie das folgende Manifest. Bei Systemmesswerten werden Extraktionsintervalle von bis zu 15 Sekunden unterstützt.

    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. Installieren Sie einen Messwertadapter. Dieser Adapter macht die Servermesswerte, die Sie nach Monitoring exportiert haben, für den HPA-Controller sichtbar. Weitere Informationen finden Sie unter Horizontales Pod-Autoscaling in der Dokumentation zu Google Cloud Managed Service for Prometheus.

    Stackdriver-Adapter für benutzerdefinierte Messwerte

    Der Stackdriver-Adapter für benutzerdefinierte Messwerte unterstützt die Abfrage von Messwerten aus Google Cloud Managed Service for Prometheus ab Version 0.13.1 des Adapters.

    So installieren Sie den Stackdriver-Adapter für benutzerdefinierte Messwerte:

    1. Verwaltete Erfassung für Ihren Cluster einrichten

    2. Installieren Sie den Stackdriver-Adapter für benutzerdefinierte Messwerte in Ihrem Cluster.

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

    3. Wenn Sie die Workload Identity Federation for GKE in Ihrem Kubernetes-Cluster aktiviert haben und sie verwenden, müssen Sie dem Dienstkonto, unter dem der Adapter ausgeführt wird, auch die Rolle „Monitoring-Betrachter“ zuweisen. Ersetzen Sie PROJECT_ID durch Ihre Projekt-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-Adapter

    Beachten Sie diese Überlegungen, wenn Sie prometheus-adapter verwenden, um mit Google Cloud Managed Service for Prometheus zu skalieren:

    • Sie können Abfragen über den Frontend-UI-Proxy von Prometheus weiterleiten, genauso wie bei der Abfrage von Google Cloud Managed Service for Prometheus mithilfe der Prometheus API oder UI. Dieses Frontend wird in einem späteren Schritt installiert.
    • Standardmäßig ist das prometheus-url-Argument der prometheus-adapter-Bereitstellung auf --prometheus-url=http://frontend.default.svc:9090/ festgelegt, wobei default der Namespace ist, in dem Sie das Frontend bereitgestellt haben. Wenn Sie das Frontend in einem anderen Namespace bereitgestellt haben, konfigurieren Sie dieses Argument entsprechend.
    • Im Feld .seriesQuery der Regelkonfiguration können Sie keinen Regex-Matcher (regulärer Ausdruck) für einen Messwertnamen verwenden. Geben Sie stattdessen Messwertnamen vollständig an.

    Da die Bereitstellung von Daten in Google Cloud Managed Service for Prometheus im Vergleich zu Upstream-Prometheus etwas länger dauern kann, kann das Konfigurieren einer übermäßig flexiblen Autoscaling-Logik zu unerwünschtem Verhalten führen. Obwohl es keine Garantie für die Datenaktualität gibt, sind Daten in der Regel 3–7 Sekunden nach dem Senden an Google Cloud Managed Service for Prometheus verfügbar, und zwar ohne Netzwerklatenz.

    Alle von prometheus-adapter ausgegebenen Abfragen sind global. Wenn Sie also Anwendungen in zwei Namespaces haben, die identisch benannte Messwerte ausgeben, wird eine HPA-Konfiguration, die diese Messwerte verwendet, anhand von Daten aus beiden Anwendungen skaliert. Verwenden Sie in Ihrer PromQL immer die Filter namespace oder cluster, um die Skalierung mit falschen Daten zu vermeiden.

    Führen Sie die folgenden Schritte aus, um eine HPA-Beispielkonfiguration mit prometheus-adapter und der verwalteten Sammlung einzurichten:

    1. Verwaltete Erfassung für Ihren Cluster einrichten

    2. Stellen Sie den Prometheus-Frontend-UI-Proxy in Ihrem Cluster bereit. Erstellen Sie das folgende Manifest mit dem Namen 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

      Wenden Sie dann das Manifest an:

      kubectl apply -f prometheus-frontend.yaml

    3. Prüfen Sie, ob prometheus-adapter in Ihrem Cluster installiert ist, indem Sie das Helm-Diagramm prometheus-community/prometheus-adapter installieren. Erstellen Sie die folgende Datei 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"})

      Verwenden Sie diese Datei dann als Werte-Datei für die Bereitstellung Ihres Helm-Diagramms:

      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

    Wenn Sie die Workload Identity Federation for GKE verwenden, müssen Sie auch ein Dienstkonto konfigurieren und autorisieren. Führen Sie dazu die folgenden Befehle aus:

    1. Erstellen Sie zuerst Ihre Dienstkonten im Cluster und Google Cloud :

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

    2. Binden Sie dann die beiden Dienstkonten. Ersetzen Sie dabei PROJECT_ID durch Ihre Projekt-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. Weisen Sie dem Google Cloud Dienstkonto als Nächstes die monitoring.viewer Rolle zu:

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

    4. Legen Sie schließlich das Dienstkonto für Ihre Frontend-Bereitstellungen als Ihr neues Dienstkonto im Cluster fest:

      kubectl set serviceaccount deployment frontend prom-frontend-sa

  3. Richten Sie die messwertbasierte HPA-Ressource ein. Stellen Sie eine HPA-Ressource bereit, die auf Ihrem bevorzugten Servermesswert basiert. Weitere Informationen finden Sie unter Horizontales Pod-Autoscaling in der Dokumentation zu Google Cloud Managed Service for Prometheus. Die spezifische HPA-Konfiguration hängt vom Typ des Messwerts (Server oder TPU) und vom installierten Messwertadapter ab.

    Für alle HPA-Konfigurationen sind einige Werte erforderlich, die festgelegt werden müssen, um eine HPA-Ressource zu erstellen:

    • MIN_REPLICAS: Die Mindestanzahl der zulässigen JetStream-Pod-Replikate. Wenn Sie das JetStream-Bereitstellungsmanifest nicht im Schritt JetStream bereitstellen ändern, empfehlen wir, diesen Wert auf 1 festzulegen.
    • MAX_REPLICAS: Die maximale Anzahl der zulässigen JetStream-Pod-Replikate. Für die JetStream-Beispielbereitstellung sind 8 Chips pro Replikat erforderlich und der Knotenpool enthält 16 Chips. Wenn Sie die Latenz beim Hochskalieren gering halten möchten, legen Sie diesen Wert auf 2 fest. Bei größeren Werten erstellt Cluster Autoscaler neue Knoten im Knotenpool, wodurch sich die Latenz beim Hochskalieren erhöht.

    • TARGET: Der Ziel-Durchschnittswert für diesen Messwert für alle JetStream-Instanzen. Weitere Informationen dazu, wie die Anzahl der Replikate aus diesem Wert bestimmt wird, finden Sie in der Kubernetes-Dokumentation unter Autoscaling.

    Stackdriver-Adapter für benutzerdefinierte Messwerte

    Der Stackdriver-Adapter für benutzerdefinierte Messwerte unterstützt die Skalierung Ihrer Arbeitslast mit dem Durchschnittswert einzelner Messwertabfragen aus Google Cloud Managed Service for Prometheus für alle Pods. Wenn Sie den Stackdriver-Adapter für benutzerdefinierte Messwerte verwenden, empfehlen wir, mit den Servermesswerten jetstream_prefill_backlog_size und jetstream_slots_used_percentage sowie dem TPU-Messwert memory_used zu skalieren.

    Erstellen Sie die folgende Datei hpa.yaml, um ein HPA-Manifest für die Skalierung mit Servermesswerten zu erstellen:

    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

    Wenn Sie den Stackdriver-Adapter für benutzerdefinierte Messwerte mit TPU-Messwerten verwenden, empfehlen wir, nur den Messwert kubernetes.io|node|accelerator|memory_used für die Skalierung zu verwenden. Erstellen Sie die folgende Datei hpa.yaml, um ein HPA-Manifest für die Skalierung mit diesem Messwert zu erstellen:

    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-Adapter

    Der Prometheus-Adapter unterstützt die Skalierung Ihrer Arbeitslast mit dem Wert von PromQL-Abfragen aus Google Cloud Managed Service for Prometheus. Zuvor haben Sie die Servermesswerte jetstream_prefill_backlog_size und jetstream_slots_used_percentage definiert, die den Durchschnittswert für alle Pods darstellen.

    Erstellen Sie die folgende Datei hpa.yaml, um ein HPA-Manifest für die Skalierung mit Servermesswerten zu erstellen:

    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

    Wenn Sie ein HPA-Manifest für die Skalierung mit TPU-Messwerten erstellen möchten, empfehlen wir, nur den in der Helm-Werte-Datei des Prometheus-Adapters definierten Messwert memory_used_percentage zu verwenden. memory_used_percentage ist der Name der folgenden PromQL-Abfrage, die die aktuelle durchschnittliche Speicherauslastung für alle Beschleuniger widerspiegelt:

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

    Erstellen Sie die folgende Datei hpa.yaml, um ein HPA-Manifest für die Skalierung mit memory_used_percentage zu erstellen:

    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

Skalieren mit mehreren Messwerten

Sie können die Skalierung auch anhand mehrerer Messwerte konfigurieren. Informationen dazu, wie die Anzahl der Replikate anhand mehrerer Messwerte bestimmt wird, finden Sie in der Kubernetes-Dokumentation zum Autoscaling. Um diese Art von HPA-Manifest zu erstellen, erfassen Sie alle Einträge aus dem Feld spec.metrics jeder HPA-Ressource in einer einzelnen HPA-Ressource. Das folgende Snippet zeigt ein Beispiel dafür, wie Sie die HPA-Ressourcen bündeln können:

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

Autoscaling überwachen und testen

Sie können beobachten, wie Ihre JetStream-Arbeitslasten basierend auf Ihrer HPA-Konfiguration skaliert werden.

Führen Sie den folgenden Befehl aus, um die Anzahl der Replikate in Echtzeit zu beobachten:

kubectl get hpa --watch

Die Ausgabe dieses Befehls sollte in etwa so aussehen:

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

Verwenden Sie den folgenden Befehl, um die Skalierungsfähigkeit Ihres HPA zu testen. Damit werden 100 Anfragen an den Modellendpunkt gesendet. Dadurch werden die verfügbaren Decodierungs-Slots erschöpft und es entsteht ein Rückstand von Anfragen in der Prefill-Warteschlange. Das HPA wird ausgelöst, um die Größe der Modellbereitstellung zu erhöhen.

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 }'

Nächste Schritte