Autoscaling für LLM-Arbeitslasten auf TPUs konfigurieren

Hier erfahren Sie, wie Sie das Autoscaling für Ihre LLM-Arbeitslasten (Large Language Model) in Google Kubernetes Engine (GKE) konfigurieren, die auf TPUs ausgeführt werden. Dazu verwenden Sie den Horizontal 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 sicherstellen.

HPA passt die Anzahl der Pod-Replikate automatisch an die beobachtete CPU-Auslastung oder andere benutzerdefinierte Messwerte 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 diese Aufgabe verwenden möchten, müssen Sie die gcloud CLI installieren und dann initialisieren. Wenn Sie die gcloud CLI bereits installiert haben, rufen Sie die neueste Version mit dem Befehl gcloud components update ab. Ältere gcloud CLI-Versionen unterstützen möglicherweise nicht die Ausführung der Befehle in diesem Dokument.

Automatische Skalierung 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. Messwerte vom JetStream-Server nach Cloud Monitoring exportieren Sie verwenden Google Cloud Managed Service for Prometheus, was die Bereitstellung und Konfiguration Ihres Prometheus-Collectors vereinfacht. 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 so einrichten, dass Google Cloud Managed Service for Prometheus Messwerte aus Ihren Pods in wiederkehrenden Intervallen von 15 Sekunden erfasst:

    Wenn Sie Servermesswerte erfassen 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 erfassen 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. Metrics Adapter installieren Mit diesem Adapter werden die Servermesswerte, die Sie nach Monitoring exportiert haben, für den HPA-Controller sichtbar. Weitere Informationen finden Sie in der Dokumentation zu Google Cloud Managed Service for Prometheus unter Horizontales Pod-Autoscaling.

    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 Folgendes, wenn Sie prometheus-adapter für die Skalierung mit Google Cloud Managed Service for Prometheus verwenden:

    • 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 Argument prometheus-url des Deployments prometheus-adapter 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 für einen Messwertnamen verwenden. Stattdessen müssen Sie Messwertnamen vollständig angeben.

    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. Um die Skalierung mit falschen Daten zu vermeiden, sollten Sie in Ihrer PromQL immer die Filter namespace oder cluster verwenden.

    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 Datei mit Werten 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 ein Kubernetes-ServiceAccount:

      kubectl create sa prom-frontend-sa

    2. Weisen Sie dem Kubernetes-ServiceAccount die Rolle „Monitoring-Betrachter“ (roles/monitoring.viewer) für das Projekt zu:

      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/default/sa/prom-frontend-sa \
  --condition=None

      Ersetzen Sie PROJECT_ID durch Ihre Projekt-ID.

    3. Konfigurieren Sie das frontend-Deployment für die Verwendung des prom-frontend-sa-ServiceAccount:

      kubectl set serviceaccount deployment frontend prom-frontend-sa

  3. Messwertbasierte HPA-Ressource einrichten: Stellen Sie eine HPA-Ressource bereit, die auf dem bevorzugten Servermesswert basiert. Weitere Informationen finden Sie in der Dokumentation zu Google Cloud Managed Service for Prometheus unter Horizontales Pod-Autoscaling. 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 zulässige Mindestanzahl von JetStream-Pod-Replikaten. Wenn Sie das JetStream-Bereitstellungsmanifest nicht im Schritt JetStream bereitstellen ändern, empfehlen wir, diesen Wert auf 1 festzulegen.
    • MAX_REPLICAS: Die maximal zulässige Anzahl von JetStream-Pod-Replikaten. Für die Beispielbereitstellung von JetStream sind 8 Chips pro Replikat erforderlich und der Knotenpool enthält 16 Chips. Wenn Sie die Latenz beim Hochskalieren niedrig halten möchten, legen Sie diesen Wert auf 2 fest. Bei größeren Werten wird der Cluster Autoscaler ausgelöst, um neue Knoten im Knotenpool zu erstellen, wodurch die Latenz beim Hochskalieren erhöht wird.

    • TARGET: Der angestrebte Durchschnitt für diesen Messwert über alle JetStream-Instanzen hinweg. Weitere Informationen dazu, wie die Anzahl der Replikate aus diesem Wert bestimmt wird, finden Sie in der Kubernetes-Dokumentation zu 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, die Skalierung mit den Servermesswerten jetstream_prefill_backlog_size und jetstream_slots_used_percentage sowie dem TPU-Messwert memory_used vorzunehmen.

    So erstellen Sie ein HPA-Manifest für das Skalieren mit Servermesswertenhpa.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

    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. Um ein HPA-Manifest für die Skalierung mit diesem Messwert zu erstellen, erstellen Sie die folgende hpa.yaml-Datei:

    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 aller Pods darstellen.

    So erstellen Sie ein HPA-Manifest für das Skalieren mit Servermesswertenhpa.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

    Wenn Sie ein HPA-Manifest für die Skalierung mit TPU-Messwerten erstellen möchten, empfehlen wir, nur die in der Helm-Wertedatei des Prometheus-Adapters definierten memory_used_percentage zu verwenden. memory_used_percentage ist der Name der folgenden PromQL-Abfrage, die den aktuellen durchschnittlichen Arbeitsspeicherverbrauch aller 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 hpa.yaml-Datei, 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

Mit mehreren Messwerten skalieren

Sie können die Skalierung auch auf Grundlage mehrerer Messwerte konfigurieren. Informationen dazu, wie die Replikatanzahl anhand mehrerer Messwerte bestimmt wird, finden Sie in der Kubernetes-Dokumentation zum Autoscaling. Um dieses HPA-Manifest zu erstellen, müssen Sie alle Einträge aus dem Feld spec.metrics jeder HPA-Ressource in einer einzelnen HPA-Ressource zusammenfassen. Das folgende Snippet zeigt ein Beispiel für das Bündeln der HPA-Ressourcen:

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 wird ein Burst von 100 Anfragen an den Modellendpunkt gesendet. Dadurch werden die verfügbaren Decodeslots erschöpft und es kommt zu einem Rückstand von Anfragen in der Vorauffüllungswarteschlange. Die HPA wird dann die Größe der Modellbereitstellung 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