Autoscaling für LLM-Arbeitslasten auf TPUs konfigurieren


Auf dieser Seite wird gezeigt, wie Sie die Autoscaling-Infrastruktur mithilfe des Horizontalen Pod-Autoscalings (HPA) von GKE einrichten, um das Gemma-Large Language Model (LLM) mit JetStream mit einzelnem Host bereitzustellen.

Weitere Informationen zum Auswählen von Messwerten für das Autoscaling finden Sie unter Best Practices für das Autoscaling von LLM-Arbeitslasten mit TPUs in GKE.

Hinweise

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 gcloud components update ab.

Autoscaling mithilfe von Messwerten

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

So richten Sie Autoscaling mit Messwerten ein:

  1. Exportieren Sie die Messwerte vom JetStream-Server nach Cloud Monitoring. 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.

    Im folgenden Beispielmanifest wird gezeigt, wie Sie Ihre PodMonitoring-Ressourcendefinitionen einrichten, damit der Google Cloud Managed Service for Prometheus in regelmäßigen Abständen von 15 Sekunden Messwerte aus Ihren Pods abruft:

    Wenn Sie Servermesswerte erfassen möchten, verwenden Sie das folgende Manifest. Bei Servermesswerten werden Extraktionsintervalle von nur 5 Sekunden unterstützt.

    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
    

    Wenn Sie TPU-Messwerte abrufen möchten, 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. Mit diesem Adapter werden die in Monitoring exportierten Servermesswerte 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.

    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 die folgenden Punkte, wenn Sie prometheus-adapter mit Google Cloud Managed Service for Prometheus 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 Argument prometheus-url der prometheus-adapter-Bereitstellung auf --prometheus-url=http://frontend.default.svc:9090/ festgelegt. Dabei ist default der Namespace, in dem Sie das Frontend bereitgestellt haben. Wenn Sie das Frontend in einem anderen Namespace bereitgestellt haben, konfigurieren Sie dieses Argument entsprechend.
    • Sie können keinen Regex-Matcher für einen Messwertnamen im Feld .seriesQuery der Regelkonfiguration 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.

    So richten Sie eine HPA-Beispielkonfiguration mit prometheus-adapter und einer verwalteten Sammlung ein:

    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. Installieren Sie dazu das Helm-Diagramm prometheus-community/prometheus-adapter. 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 Wertedatei 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 In-Cluster- und Google Cloud-Dienstkonten:

      gcloud iam service-accounts create prom-frontend-sa && kubectl create sa prom-frontend-sa
      
    2. Binden Sie dann die beiden Dienstkonten:

      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. Weisen Sie als Nächstes dem Google Cloud-Dienstkonto die Rolle monitoring.viewer zu:

      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. Legen Sie abschließend das Dienstkonto für die Frontend-Bereitstellungen als neues In-Cluster-Dienstkonto 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 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 dem installierten Messwertadapter ab.

    Einige Werte sind für alle HPA-Konfigurationen erforderlich und müssen festgelegt werden, um eine HPA-Ressource zu erstellen:

    • MIN_REPLICAS: Die Mindestanzahl zulässiger JetStream-Pod-Replikate. Wenn Sie das JetStream-Bereitstellungsmanifest aus dem Schritt JetStream bereitstellen nicht ä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 Replikate erforderlich und der Knotenpool enthält 16 Chips. Wenn Sie die Latenz für die vertikale Skalierung niedrig halten möchten, legen Sie den Wert auf 2 fest. Größere Werte lösen das Cluster Autoscaler dazu aus, neue Knoten im Knotenpool zu erstellen, wodurch die Latenz hochskaliert wird.
    • TARGET: Der Zieldurchschnitt für diesen Messwert in allen JetStream-Instanzen. Weitere Informationen dazu, wie die Anzahl der Replikate anhand dieses Werts ermittelt wird, finden Sie in der Kubernetes-Dokumentation zum 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 von Google Cloud Managed Service for Prometheus über alle Pods hinweg. Wenn Sie den Stackdriver-Adapter für benutzerdefinierte Messwerte verwenden, empfehlen wir, die Skalierung anhand der Servermesswerte jetstream_prefill_backlog_size und jetstream_slots_used_percentage sowie des TPU-Messwerts memory_used vorzunehmen.

    Wenn Sie ein HPA-Manifest zum Skalieren mit Servermesswerten erstellen möchten, 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: 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 hpa.yaml-Datei, 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 von Google Cloud Managed Service for Prometheus. Sie haben bereits die Servermesswerte jetstream_prefill_backlog_size und jetstream_slots_used_percentage definiert, die den durchschnittlichen Wert aller Pods darstellen.

    Wenn Sie ein HPA-Manifest zum Skalieren mit Servermesswerten erstellen möchten, 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: jetstream_METRIC
          target:
            type: AverageValue
            averageValue: TARGET
    

    Wenn Sie ein HPA-Manifest zum Skalieren mit TPU-Messwerten erstellen möchten, empfehlen wir, nur die memory_used_percentage zu verwenden, die in der Helm-Wertdatei „prometheus-adapter“ definiert ist. 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 zum Erstellen eines HPA-Manifests für die Skalierung mit memory_used_percentage 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: memory_used_percentage
          target:
            type: AverageValue
            averageValue: TARGET
    

Skalierung mit mehreren Messwerten

Sie können die Skalierung auch basierend auf mehreren Messwerten konfigurieren. Informationen dazu, wie die Anzahl der Replikate anhand mehrerer Messwerte ermittelt wird, finden Sie in der Kubernetes-Dokumentation zum Autoscaling. Erfassen Sie alle Einträge aus dem Feld spec.metrics jeder HPA-Ressource in einer einzigen HPA-Ressource, um diese Art von HPA-Manifest zu erstellen. 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 Replikatanzahl 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

Mit dem folgenden Befehl können Sie die Skalierbarkeit Ihres HPA testen. Dabei werden 100 Anfragen an den Modellendpunkt gesendet. Dadurch werden die verfügbaren Dekodierungsslots aufgebraucht und es kommt zu einem Rückstau von Anfragen in der Prefill-Warteschlange. Das wiederum führt dazu, dass die HPA die Größe der Modellbereitstellung erhöht.

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