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

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

Per saperne di più sulla selezione delle metriche per la scalabilità automatica, consulta 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à, installala e poi inizializza gcloud CLI. Se hai già installato gcloud CLI, scarica l'ultima versione eseguendo gcloud components update.

Scalabilità automatica tramite le metriche

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

Per configurare la scalabilità automatica con le metriche:

  1. Esporta le metriche dal server JetStream a 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 cluster GKE. Puoi anche abilitarlo manualmente.

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

    Se devi estrarre le metriche del server, utilizza il seguente manifest. Con le metriche del server, sono supportati intervalli di scraping frequenti 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 scraping 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 che hai esportato in Monitoring. Per maggiori dettagli, consulta la sezione Scalabilità automatica orizzontale dei pod nella documentazione di Google Cloud Managed Service per Prometheus.

    Adattatore Stackdriver per le metriche personalizzate

    L'adattatore Stackdriver per le metriche personalizzate supporta l'esecuzione di query sulle metriche di Google Cloud Managed Service per Prometheus, a partire dalla versione v0.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 abilitato Workload Identity Federation for GKE sul tuo cluster Kubernetes e utilizzi Workload Identity Federation for GKE, devi anche concedere il ruolo Visualizzatore Monitoring al account di servizio con cui viene eseguito l'adattatore. Sostituisci PROJECT_ID con l'ID 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 scalare utilizzando Google Cloud Managed Service per Prometheus:

    • Instrada le query tramite il proxy della UI frontend di Prometheus, proprio come quando esegui query su Google Cloud Managed Service per Prometheus utilizzando l'API o la UI 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 hai 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 un matcher di espressioni regolari (regex) su un nome di metrica. Specifica invece i nomi delle metriche per intero.

    Poiché i dati potrebbero impiegare un po' più di tempo per essere disponibili in Google Cloud Managed Service per Prometheus rispetto a Prometheus upstream, la configurazione di una logica di scalabilità automatica eccessivamente rapida può causare un comportamento indesiderato. Sebbene non vi sia alcuna garanzia sulla freschezza dei dati, in genere i dati sono disponibili per le query 3-7 secondi dopo l'invio a Google Cloud Managed Service per Prometheus, escluso qualsiasi ritardo di rete.

    Tutte le query emesse da prometheus-adapter hanno portata 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 viene scalata utilizzando i dati di entrambe le applicazioni. Per evitare lo 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, segui questi passaggi:

    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 dei valori 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 service account eseguendo i seguenti comandi:

    1. Innanzitutto, crea gli account di servizio Trusted Cloud by S3NS in-cluster:

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

    2. Poi, associa i due service account, 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.s3ns-system.iam.gserviceaccount.com \
&&
kubectl annotate serviceaccount \
  --namespace default \
  prom-frontend-sa \
  iam.gke.io/gcp-service-account=jetstream-iam-sa@PROJECT_ID.s3ns-system.iam.gserviceaccount.com

    3. Successivamente, assegna al service account Trusted Cloud by S3NS il ruolo monitoring.viewer:

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

    4. Infine, imposta l'account di servizio delle implementazioni frontend in modo che sia il nuovo service account in-cluster:

      kubectl set serviceaccount deployment frontend prom-frontend-sa

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

    Per creare una risorsa HPA, è necessario impostare alcuni valori in tutte le configurazioni HPA:

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

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

    Adattatore Stackdriver per le metriche personalizzate

    L'adattatore Stackdriver delle metriche personalizzate supporta lo scaling del workload con il valore medio delle singole query delle metriche di Google Cloud Managed Service per Prometheus in tutti i pod. Quando utilizzi l'adattatore Stackdriver per le metriche personalizzate, ti consigliamo di scalare 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 lo scaling 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 lo scaling. Per creare un manifest HPA per lo scaling con 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 il ridimensionamento del 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 lo scaling 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 lo scaling con le metriche TPU, ti consigliamo di utilizzare solo memory_used_percentage definito nel file dei valori Helm di prometheus-adapter. memory_used_percentage è il nome assegnato alla seguente query PromQL che riflette la memoria media attuale 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

Scalare utilizzando più metriche

Puoi anche configurare lo scaling 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. Il seguente snippet mostra un esempio di come 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 come vengono scalati i tuoi workload JetStream in base alla configurazione di HPA.

Per osservare il conteggio delle repliche in tempo reale, esegui questo 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à di HPA, utilizza il seguente comando che invia un burst di 100 richieste all'endpoint del modello. In questo modo, gli slot di decodifica disponibili verranno esauriti e si creerà un backlog di richieste nella coda di precompilazione, il che attiverà 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