Configurare la scalabilità automatica per i carichi di lavoro LLM sulle TPU

Questa pagina mostra come configurare l'infrastruttura di scalabilità automatica utilizzando l'Horizontal Pod Autoscaler (HPA) di GKE per eseguire il deployment del modello linguistico di grandi dimensioni (LLM) di Gemma utilizzando JetStream a host singolo.

Per scoprire di più sulla selezione delle metriche per la scalabilità automatica, consulta le best practice per la scalabilità automatica dei carichi di lavoro LLM con le TPU su GKE.

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti operazioni:

  • Attiva l'API Google Kubernetes Engine.
    • Attiva l'API Google Kubernetes Engine
  • Se vuoi utilizzare Google Cloud CLI per questa attività, installa e poi inizializza gcloud CLI. Se hai già installato gcloud CLI, ottieni la versione più recente eseguendo gcloud components update.

Scalabilità automatica mediante metriche

Puoi utilizzare le metriche sul rendimento specifiche per il carico di lavoro emesse dal server di inferenza JetStream o le metriche sul rendimento delle TPU per indirizzare la scalabilità automatica dei pod.

Per configurare l'autoscaling con le metriche:

  1. Esporta le metriche dal server JetStream in Cloud Monitoring. Utilizzi Google Cloud Managed Service per Prometheus, che semplifica il deployment e la configurazione del raccoglitore Prometheus. Google Cloud Managed Service per Prometheus è abilitato per impostazione predefinita nel tuo cluster GKE. Puoi anche attivarlo manualmente.

    Il seguente manifest di esempio mostra come configurare le definizioni delle risorse PodMonitoring per indicare a Google Cloud Managed Service per Prometheus di estrarre le metriche dai pod a intervalli ricorrenti di 15 secondi:

    Se devi eseguire lo scraping delle metriche del server, utilizza il seguente manifest. Con le metriche del server, sono supportati intervalli di scansione fino a 5 secondi.

    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

    Se devi eseguire lo scraping delle metriche TPU, utilizza il seguente manifest. Con le metriche di sistema, sono supportati intervalli di scansione fino a 15 secondi.

    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. Installa un adattatore delle metriche. Questo adattatore rende visibili al controller HPA le metriche del server esportate in Monitoring. Per maggiori dettagli, consulta Scalabilità automatica dei pod orizzontali nella documentazione di Google Cloud Managed Service per Prometheus.

    Adattatore Stackdriver per le metriche personalizzate

    L'adattatore Stackdriver per le metriche personalizzate supporta le query sulle metriche di Google Cloud Managed Service per Prometheus, a partire dalla versione 0.13.1 dell'adattatore.

    Per installare l'adattatore Stackdriver delle metriche personalizzate:

    1. Configura la raccolta gestita nel cluster.

    2. Installa l'adattatore Stackdriver delle metriche personalizzate nel tuo cluster.

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

    3. Se hai attivato Workload Identity Federation for GKE sul tuo cluster Kubernetes e utilizzi Workload Identity Federation for GKE, devi concedere il ruolo Visualizzatore monitoraggio anche all'account di servizio in cui viene eseguito l'adattatore. Sostituisci PROJECT_ID con l'ID del tuo progetto.

    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

    Adattatore Prometheus

    Tieni presente queste considerazioni quando utilizzi prometheus-adapter per eseguire la scalabilità utilizzando Google Cloud Managed Service per Prometheus:

    • Inoltra le query tramite il proxy dell'interfaccia utente frontend di Prometheus, proprio come quando esegui query su Google Cloud Managed Service per Prometheus utilizzando l'API o l'interfaccia utente di Prometheus. Questo frontend viene installato in un passaggio successivo.
    • Per impostazione predefinita, l'argomento prometheus-url del deployment prometheus-adapter è impostato su --prometheus-url=http://frontend.default.svc:9090/, dove default è lo spazio dei nomi in cui è stato eseguito il deployment del frontend. Se hai eseguito il deployment del frontend in un altro spazio dei nomi, configura questo argomento di conseguenza.
    • Nel campo .seriesQuery della configurazione delle regole, non puoi utilizzare una corrispondenza di espressioni regolari (regex) su un nome metrica. Specifica invece i nomi completi delle metriche.

    Poiché i dati possono richiedere un po' più di tempo per essere disponibili in Google Cloud Managed Service per Prometheus rispetto a Prometheus a monte, la configurazione di una logica di scalabilità automatica troppo esuberante può causare comportamenti indesiderati. Sebbene non sia garantita l'aggiornamento dei dati, in genere è possibile eseguire query sui dati 3-7 secondi dopo che sono stati inviati a Google Cloud Managed Service per Prometheus, esclusa la latenza della rete.

    Tutte le query inviate da prometheus-adapter hanno ambito globale. Ciò significa che se hai applicazioni in due spazi dei nomi che emettono metriche con lo stesso nome, una configurazione HPA che utilizza questa metrica esegue il ridimensionamento utilizzando i dati di entrambe le applicazioni. Per evitare di eseguire il scaling utilizzando dati errati, utilizza sempre i filtri namespace o cluster in PromQL.

    Per configurare una configurazione HPA di esempio utilizzando prometheus-adapter e la raccolta gestita:

    1. Configura la raccolta gestita nel cluster.

    2. Esegui il deployment del proxy dell'interfaccia utente frontend di Prometheus nel cluster. Crea il seguente manifest denominato 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

      Quindi applica il manifest:

      kubectl apply -f prometheus-frontend.yaml

    3. Assicurati che prometheus-adapter sia installato nel cluster installando il grafico Helm prometheus-community/prometheus-adapter. Crea il seguente file 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"})

      Quindi, utilizza questo file come file values per il deployment del grafico 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

    Se utilizzi Workload Identity Federation for GKE, devi anche configurare e autorizzare un account di servizio eseguendo i seguenti comandi:

    1. Per prima cosa, crea gli account di servizio in-cluster e Google Cloud:

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

    2. Quindi, associa i due account di servizio, assicurandoti di sostituire PROJECT_ID con il tuo ID progetto:

      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. Quindi, assegna all'account di servizio Google Cloud il ruolo monitoring.viewer:

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

    4. Infine, imposta l'account di servizio per i deployment frontend come nuovo account di servizio all'interno del cluster:

      kubectl set serviceaccount deployment frontend prom-frontend-sa

  3. Configura la risorsa HPA basata su metriche. Esegui il deployment di una risorsa HPA basata sulla metrica del server che preferisci. Per maggiori dettagli, consulta Scalabilità automatica dei pod orizzontali nella documentazione di Google Cloud Managed Service per Prometheus. La configurazione specifica dell'HPA dipende dal tipo di metrica (server o TPU) e dall'adattatore metrico installato.

    Alcuni valori sono obbligatori in tutte le configurazioni HPA e devono essere impostati per creare una risorsa HPA:

    • MIN_REPLICAS: il numero minimo di repliche del pod JetStream consentito. Se non modifichi il manifest di deployment di JetStream dal passaggio Esegui il deployment di JetStream, ti consigliamo di impostare questo valore su 1.
    • MAX_REPLICAS: il numero massimo di repliche del pod JetStream consentito. Il deployment di JetStream di esempio richiede 8 chip per replica e il pool di nodi contiene 16 chip. Se vuoi mantenere bassa la latenza del ridimensionamento, imposta questo valore su 2. Valori più elevati attivano il gestore della scalabilità automatica dei cluster per creare nuovi nodi nel pool di nodi, aumentando così la latenza dello scale up.

    • TARGET: la media target per questa metrica in tutte le istanze JetStream. Per ulteriori informazioni su come viene determinato il numero di repliche da questo valore, consulta la documentazione di Kubernetes per la scalabilità automatica.

    Adattatore Stackdriver per le metriche personalizzate

    L'adattatore Stackdriver delle metriche personalizzate supporta il ridimensionamento del tuo carico di lavoro con il valore medio delle singole query sulle metriche di Google Cloud Managed Service per Prometheus in tutti i pod. Quando utilizzi l'adattatore Stackdriver per le metriche personalizzate, ti consigliamo di eseguire la scalabilità con le metriche del server jetstream_prefill_backlog_size e jetstream_slots_used_percentage e la metrica TPU memory_used.

    Per creare un manifest HPA per il ridimensionamento con le metriche del server, crea il seguente file 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

    Quando utilizzi l'adattatore Stackdriver per le metriche personalizzate con le metriche TPU, ti consigliamo di utilizzare solo la metrica kubernetes.io|node|accelerator|memory_used per il ridimensionamento. Per creare un manifest HPA per il ridimensionamento in base a questa metrica, crea il seguente file 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

    Adattatore Prometheus

    Prometheus Adapter supporta la scalabilità del tuo carico di lavoro con il valore delle query PromQL di Google Cloud Managed Service per Prometheus. In precedenza, hai definito le metriche del server jetstream_prefill_backlog_size e jetstream_slots_used_percentage che rappresentano il valore medio in tutti i pod.

    Per creare un manifest HPA per il ridimensionamento con le metriche del server, crea il seguente file 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

    Per creare un manifest HPA per il ridimensionamento con le metriche TPU, ti consigliamo di utilizzare solo memory_used_percentage definito nel file dei valori di Helm prometheus-adapter. memory_used_percentage è il nome assegnato alla seguente query PromQL che riflette la memoria media corrente utilizzata in tutti gli acceleratori:

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

    Per creare un manifest HPA per la scalabilità con memory_used_percentage, crea il seguente file 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

Eseguire la scalatura utilizzando più metriche

Puoi anche configurare la scalabilità in base a più metriche. Per scoprire come viene determinato il numero di repliche utilizzando più metriche, consulta la documentazione di Kubernetes sulla scalabilità automatica. Per creare questo tipo di manifest HPA, raccogli tutte le voci del campo spec.metrics di ogni risorsa HPA in un'unica risorsa HPA. Lo snippet seguente mostra un esempio di come puoi raggruppare le risorse 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

Monitorare e testare la scalabilità automatica

Puoi osservare la scalabilità dei carichi di lavoro JetStream in base alla configurazione dell'HPA.

Per osservare il conteggio delle repliche in tempo reale, esegui il seguente comando:

kubectl get hpa --watch

L'output di questo comando dovrebbe essere simile al seguente:

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

Per testare la capacità di scalabilità dell'HPA, utilizza il seguente comando che invia una raffica di 100 richieste all'endpoint del modello. In questo modo verranno esauriti gli slot di decodifica disponibili e verrà generato un backlog di richieste nella coda di precompilazione, attivando l'HPA per aumentare le dimensioni del deployment del modello.

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

Passaggi successivi