Configurar el autoescalado de cargas de trabajo de LLM en TPUs

En esta página se muestra cómo configurar tu infraestructura de autoescalado mediante el autoescalador horizontal de pods (HPA) de GKE para desplegar el modelo de lenguaje extenso (LLM) Gemma con JetStream de un solo host.

Para obtener más información sobre cómo seleccionar métricas para el autoescalado, consulta Prácticas recomendadas para autoescalar cargas de trabajo de LLMs con TPUs en GKE.

Antes de empezar

Antes de empezar, asegúrate de que has realizado las siguientes tareas:

  • Habilita la API de Google Kubernetes Engine.
    • Habilitar la API de Google Kubernetes Engine
  • Si quieres usar Google Cloud CLI para esta tarea, instálala y, a continuación, inicialízala. Si ya has instalado la gcloud CLI, obtén la versión más reciente ejecutando gcloud components update.

Autoescalar con métricas

Puedes usar las métricas de rendimiento específicas de la carga de trabajo que emite el servidor de inferencia de JetStream o las métricas de rendimiento de TPU para dirigir el autoescalado de tus pods.

Para configurar el escalado automático con métricas, sigue estos pasos:

  1. Exporta las métricas del servidor JetStream a Cloud Monitoring. Usas Managed Service para Prometheus de Google Cloud, que simplifica el despliegue y la configuración de tu recopilador de Prometheus. Google Cloud Managed Service para Prometheus está habilitado de forma predeterminada en tu clúster de GKE. También puedes habilitarlo manualmente.

    En el siguiente manifiesto de ejemplo se muestra cómo configurar las definiciones de recursos PodMonitoring para dirigir Google Cloud Managed Service para Prometheus a que obtenga métricas de tus pods a intervalos recurrentes de 15 segundos:

    Si necesitas raspar métricas del servidor, usa el siguiente manifiesto. Con las métricas de servidor, se admiten intervalos de raspado de hasta 5 segundos.

    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

    Si necesitas raspar métricas de TPU, usa el siguiente manifiesto. Con las métricas del sistema, se admiten intervalos de raspado de hasta 15 segundos.

    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. Instala un adaptador de métricas. Este adaptador hace que las métricas del servidor que has exportado a Monitoring sean visibles para el controlador de HPA. Para obtener más información, consulta Escalado automático horizontal de pods en la documentación de Google Cloud Managed Service para Prometheus.

    Adaptador de métricas personalizadas de Stackdriver

    El adaptador de Stackdriver para métricas personalizadas admite consultas de métricas de Google Cloud Managed Service para Prometheus a partir de la versión 0.13.1 del adaptador.

    Para instalar el adaptador de métricas personalizadas de Stackdriver, siga estos pasos:

    1. Configure la recogida gestionada en su clúster.

    2. Instala el adaptador de métricas personalizadas de Stackdriver en tu clúster.

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

    3. Si tienes habilitada la federación de Workload Identity para GKE en tu clúster de Kubernetes y la usas, también debes conceder el rol Lector de Monitoring a la cuenta de servicio con la que se ejecuta el adaptador. Sustituye PROJECT_ID por el ID de tu proyecto.

    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

    Adaptador de Prometheus

    Ten en cuenta estas consideraciones al usar prometheus-adapter para escalar con Google Cloud Managed Service para Prometheus:

    • Encamina las consultas a través del proxy de la interfaz de usuario de frontend de Prometheus, al igual que cuando consultas Google Cloud Managed Service para Prometheus con la API o la interfaz de usuario de Prometheus. Este frontend se instala en un paso posterior.
    • De forma predeterminada, el argumento prometheus-url de prometheus-adapter Deployment se define como --prometheus-url=http://frontend.default.svc:9090/, donde default es el espacio de nombres en el que has implementado el frontend. Si has implementado el frontend en otro espacio de nombres, configura este argumento en consecuencia.
    • En el campo .seriesQuery de la configuración de las reglas, no puede usar un comparador de expresiones regulares (regex) en un nombre de métrica. En su lugar, especifique los nombres de las métricas por completo.

    Como los datos pueden tardar un poco más en estar disponibles en Google Cloud Managed Service para Prometheus que en el Prometheus upstream, configurar una lógica de escalado automático demasiado ambiciosa puede provocar un comportamiento no deseado. Aunque no se garantiza la actualización de los datos, normalmente se pueden consultar entre 3 y 7 segundos después de enviarlos a Google Cloud Managed Service para Prometheus, sin incluir la latencia de la red.

    Todas las consultas emitidas por prometheus-adapter tienen un ámbito global. Esto significa que, si tienes aplicaciones en dos espacios de nombres que emiten métricas con el mismo nombre, una configuración de HPA que use esa métrica se escalará usando datos de ambas aplicaciones. Para evitar que se escale con datos incorrectos, utilice siempre los filtros namespace o cluster en PromQL.

    Para configurar un ejemplo de HPA con prometheus-adapter y la recogida gestionada, sigue estos pasos:

    1. Configure la recogida gestionada en su clúster.

    2. Despliega el proxy de la interfaz de usuario de frontend de Prometheus en tu clúster. Crea el siguiente archivo de manifiesto llamado 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

      A continuación, aplica el manifiesto:

      kubectl apply -f prometheus-frontend.yaml

    3. Asegúrate de que prometheus-adapter esté instalado en tu clúster instalando el gráfico de Helm prometheus-community/prometheus-adapter. Crea el siguiente archivo 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"})

      A continuación, usa este archivo como archivo de valores para implementar tu gráfico de 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

    Si usas Workload Identity Federation para GKE, también debes configurar y autorizar una cuenta de servicio ejecutando los siguientes comandos:

    1. Primero, crea tus cuentas de servicio y en el clúster: Trusted Cloud by S3NS 

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

    2. A continuación, vincula las dos cuentas de servicio. No olvides sustituir PROJECT_ID por el ID de tu proyecto:

      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. A continuación, asigna el rol monitoring.viewer a la cuenta de servicio: Trusted Cloud by S3NS 

      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. Por último, define la cuenta de servicio de tus implementaciones de frontend como tu nueva cuenta de servicio en el clúster:

      kubectl set serviceaccount deployment frontend prom-frontend-sa

  3. Configura el recurso HPA basado en métricas. Despliega un recurso HPA basado en la métrica de servidor que prefieras. Para obtener más información, consulta Escalado automático horizontal de pods en la documentación de Google Cloud Managed Service para Prometheus. La configuración específica de HPA depende del tipo de métrica (servidor o TPU) y del adaptador de métricas que esté instalado.

    Hay algunos valores que son obligatorios en todas las configuraciones de HPA y que se deben definir para crear un recurso de HPA:

    • MIN_REPLICAS: número mínimo de réplicas de pods de JetStream permitidas. Si no vas a modificar el manifiesto de implementación de JetStream en el paso Implementar JetStream, te recomendamos que le asignes el valor 1.
    • MAX_REPLICAS: número máximo de réplicas de pods de JetStream permitidas. El ejemplo de implementación de JetStream requiere 8 chips por réplica y el pool de nodos contiene 16 chips. Si quieres que la latencia de ampliación siga siendo baja, asigna el valor 2. Cuanto más altos sean los valores, más tiempo tardará la herramienta de adaptación dinámica del clúster en crear nodos en el grupo de nodos, lo que aumentará la latencia de escalado vertical.

    • TARGET: el promedio objetivo de esta métrica en todas las instancias de JetStream. Consulta la documentación de Kubernetes sobre el autoescalado para obtener más información sobre cómo se determina el número de réplicas a partir de este valor.

    Adaptador de métricas personalizadas de Stackdriver

    El adaptador de métricas personalizadas de Stackdriver permite escalar tu carga de trabajo con el valor medio de las consultas de métricas individuales de Google Cloud Managed Service para Prometheus en todos los pods. Cuando se usa el adaptador de métricas personalizadas de Stackdriver, recomendamos escalar con las métricas de servidor jetstream_prefill_backlog_size y jetstream_slots_used_percentage, así como con la métrica de TPU memory_used.

    Para crear un manifiesto de HPA para escalar con métricas del servidor, crea el siguiente archivo 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

    Cuando se usa el adaptador de Stackdriver de métricas personalizadas con métricas de TPU, recomendamos usar solo la métrica kubernetes.io|node|accelerator|memory_used para el escalado. Para crear un manifiesto de HPA para escalar con esta métrica, cree el siguiente archivo 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

    Adaptador de Prometheus

    Prometheus Adapter admite el escalado de tu carga de trabajo con el valor de las consultas de PromQL de Google Cloud Managed Service para Prometheus. Antes, ha definido las métricas de servidor jetstream_prefill_backlog_size y jetstream_slots_used_percentage que representan el valor medio de todos los pods.

    Para crear un manifiesto de HPA para escalar con métricas del servidor, crea el siguiente archivo 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

    Para crear un manifiesto de HPA para escalar con métricas de TPU, te recomendamos que solo uses el memory_used_percentage definido en el archivo de valores de Helm de prometheus-adapter. memory_used_percentage es el nombre que se le ha dado a la siguiente consulta de PromQL, que refleja la media actual de memoria utilizada en todos los aceleradores:

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

    Para crear un manifiesto de HPA para escalar con memory_used_percentage, crea el siguiente archivo 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

Escalar con varias métricas

También puedes configurar el escalado en función de varias métricas. Para saber cómo se determina el número de réplicas mediante varias métricas, consulta la documentación de Kubernetes sobre el autoescalado. Para crear este tipo de manifiesto de HPA, recopila todas las entradas del campo spec.metrics de cada recurso de HPA en un único recurso de HPA. El siguiente fragmento muestra un ejemplo de cómo puedes agrupar los recursos de 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

Monitorizar y probar el autoescalado

Puedes observar cómo se escalan tus cargas de trabajo de JetStream en función de tu configuración de HPA.

Para observar el número de réplicas en tiempo real, ejecuta el siguiente comando:

kubectl get hpa --watch

El resultado de este comando debería ser similar al siguiente:

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

Para probar la capacidad de escalado de tu HPA, usa el siguiente comando, que envía una ráfaga de 100 solicitudes al endpoint del modelo. De esta forma, se agotarán los espacios de decodificación disponibles y se producirá un retraso en la cola de relleno previo, lo que hará que HPA aumente el tamaño de la implementación del modelo.

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

Siguientes pasos