Configura el ajuste de escala automático para cargas de trabajo de LLM en TPU

Autopilot Estándar

En esta página, se muestra cómo configurar tu infraestructura de ajuste de escala automático mediante Horizontal Pod Autoscaler (HPA) de GKE para implementar el modelo de lenguaje grande (LLM) de Gemma con JetStream de un solo host.

Si deseas obtener más información sobre cómo seleccionar métricas para el ajuste de escala automático, consulta Prácticas recomendadas para el ajuste de escala automático de cargas de trabajo de LLM con TPU en GKE.

Antes de comenzar

Antes de comenzar, asegúrate de haber realizado las siguientes tareas:

Habilita la API de Google Kubernetes Engine.

Habilitar la API de Google Kubernetes Engine

Si deseas usar Google Cloud CLI para esta tarea, instala y, luego, inicializa gcloud CLI. Si ya instalaste gcloud CLI, ejecuta gcloud components update para obtener la versión más reciente.
Nota: Para las instalaciones de gcloud CLI existentes, asegúrate de configurar las propiedades compute/region y compute/zone. Cuando configuras las ubicaciones predeterminadas, puedes evitar errores en la CLI de gcloud como el siguiente: One of [--zone, --region] must be supplied: Please specify location.

Familiarízate con el flujo de trabajo y complétalo en Entrega Gemma con TPU en GKE con JetStream. Asegúrate de que el argumento PROMETHEUS_PORT esté configurado en el manifiesto de implementación de JetStream.

Ajusta la escala automáticamente 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 ajuste de escala automático de tus Pods.

Para configurar el ajuste de escala automático con métricas, sigue estos pasos:

Exporta las métricas del servidor de JetStream a Cloud Monitoring. Usas Google Cloud Managed Service para Prometheus, que simplifica la implementación y la configuración de tu colector 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 tus definiciones de recursos PodMonitoring a fin de dirigir a Google Cloud Managed Service para Prometheus para que recopile métricas de tus Pods en intervalos recurrentes de 15 segundos:

Si necesitas raspar métricas del servidor, usa el siguiente manifiesto. Con las métricas del servidor, se admiten intervalos de recopilación de hasta 5 segundos.
```
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
```
Si necesitas hacer scraping de métricas de TPU, usa el siguiente manifiesto. Con las métricas del sistema, se admiten intervalos de scraping de tan solo 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
```

Instala un adaptador de métricas. Este adaptador hace que las métricas del servidor que exportaste a Monitoring sean visibles para el controlador de HPA. Si deseas obtener más detalles, consulta Ajuste automático de escala horizontal de Pods en la documentación de Google Cloud Managed Service para Prometheus.

Si deseas que JetStream escale con métricas individuales, usa el adaptador de métricas personalizadas de Stackdriver.
Si quieres que JetStream escale con el valor de una expresión compuesta por varias métricas distintas, usa el adaptador de Prometheus de terceros.

Adaptador de métricas personalizadas de Stackdriver

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

Para instalar el adaptador de métricas personalizadas de Stackdriver, haz lo siguiente:

Configurar la recopilación administrada para tu clúster.

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

Si tienes habilitada la federación de identidades para cargas de trabajo para GKE en tu clúster de Kubernetes y la usas, también debes otorgar el rol de visualizador de Monitoring a la cuenta de servicio bajo la que se ejecuta el adaptador.

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 cuando uses prometheus-adapter para escalar con Google Cloud Managed Service para Prometheus:

Enruta las consultas a través del proxy de IU de frontend de Prometheus, al igual que cuando consultas Google Cloud Managed Service para Prometheus con la API o la IU de Prometheus. Este frontend se instalará en un paso posterior.
De forma predeterminada, el argumento prometheus-url del Deployment prometheus-adapter se establece en --prometheus-url=http://frontend.default.svc:9090/, en el que default es el espacio de nombres en el que implementaste el frontend. Si implementaste el frontend en otro espacio de nombres, configura este argumento según corresponda.
En el campo .seriesQuery de la configuración de reglas, no puedes usar un comparador de expresiones regulares (regex) en un nombre de métrica. En su lugar, especifica por completo los nombres de las métricas.

Como los datos pueden tardar un poco más en estar disponibles dentro de Google Cloud Managed Service para Prometheus en comparación con el procesamiento ascendente de Prometheus, la configuración de una lógica de ajuste de escala automático demasiado apresurada puede provocar un comportamiento no deseado. Aunque no hay garantía de que los datos se actualicen, por lo general, los datos están disponibles para consultar entre 3 y 7 segundos después de enviarlos a Google Cloud Managed Service para Prometheus, sin incluir ninguna latencia de red.

Todas las consultas emitidas por prometheus-adapter son de alcance global. Esto significa que, si tienes aplicaciones en dos espacios de nombres que emiten métricas con nombre idéntico, una configuración de HPA que usa esa métrica se escala con datos de ambas aplicaciones. Para evitar el escalamiento con datos incorrectos, usa siempre los filtros namespace o cluster en tu PromQL.

Para establecer una configuración de HPA de ejemplo con prometheus-adapter y la colección administrada, sigue estos pasos:

Configurar la recopilación administrada para tu clúster.

Implementa el proxy de la IU del frontend de Prometheus en tu clúster. Crea el siguiente 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

Luego, aplica el manifiesto:

kubectl apply -f prometheus-frontend.yaml

Instala el gráfico de Helm prometheus-community/prometheus-adapter para asegurarte de que prometheus-adapter esté instalado en tu clúster. 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"})

Luego, usa este archivo como el archivo de valores para implementar el 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 la federación de identidades para cargas de trabajo para GKE, también debes configurar y autorizar una cuenta de servicio. Para ello, ejecuta los siguientes comandos:

Primero, crea tus cuentas de servicio de Google Cloud en el clúster:

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

Luego, vincula las dos cuentas de servicio:

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

A continuación, otorga el rol monitoring.viewer a la cuenta de servicio de Google Cloud:

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

Por último, configura tu cuenta de servicio de implementaciones de frontend para que sea tu nueva cuenta de servicio en el clúster:
```
kubectl set serviceaccount deployment frontend prom-frontend-sa
```

Configura el recurso de HPA basado en métricas. Implementa un recurso de HPA que se base en tu métrica de servidor preferida. Si deseas obtener más detalles, consulta Ajuste de escala 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 instalado.

Algunos valores se requieren en todas las configuraciones de HPA y se deben establecer para crear un recurso de HPA:
- MIN_REPLICAS: Es la cantidad mínima de réplicas de Pods de JetStream permitidas. Si no modificas el manifiesto de implementación de JetStream del paso Implementa JetStream, te recomendamos que lo establezcas en 1.
- MAX_REPLICAS: Es la cantidad máxima de réplicas de Pods de JetStream permitidas. La implementación de ejemplo de JetStream requiere 8 chips por réplica y el grupo de nodos contiene 16 chips. Si quieres mantener baja la latencia de escalamiento vertical, establece este valor en 2. Los valores más grandes activan el escalador automático del clúster para crear nodos nuevos en el grupo de nodos, lo que aumenta la latencia de escalamiento vertical.
- TARGET: Es el promedio objetivo de esta métrica en todas las instancias de JetStream. Consulta la documentación de Kubernetes sobre el ajuste de escala automático para obtener más información sobre cómo se determina el recuento de réplicas a partir de este valor.
Adaptador de métricas personalizadas de Stackdriver
El adaptador de métricas personalizadas de Stackdriver admite el escalamiento de tu carga de trabajo con el valor promedio de las consultas de métricas individuales de Google Cloud Managed Service para Prometheus en todos los Pods. Cuando usas el adaptador de métricas personalizadas de Stackdriver, recomendamos escalar con las métricas del servidor jetstream_prefill_backlog_size y jetstream_slots_used_percentage y la métrica de TPU memory_used.

Si deseas 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 usas el adaptador de métricas personalizadas de Stackdriver con métricas de TPU, recomendamos usar solo la métrica kubernetes.io|node|accelerator|memory_used para el escalamiento. Si deseas crear un manifiesto de HPA para escalar con esta métrica, 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: 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
El adaptador de Prometheus admite el escalamiento de tu carga de trabajo con el valor de las consultas de PromQL desde Google Cloud Managed Service para Prometheus. Anteriormente, definiste las métricas del servidor jetstream_prefill_backlog_size y jetstream_slots_used_percentage que representan el valor promedio de todos los Pods.

Si deseas 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 prometheus-adapter. memory_used_percentage es el nombre que se le asignó a la siguiente consulta de PromQL, que refleja la memoria promedio actual que se usa 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"})
```
Si deseas 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
```

Escala con varias métricas

También puedes configurar el escalamiento en función de varias métricas. Para obtener información sobre cómo se determina el recuento de réplicas con varias métricas, consulta la documentación de Kubernetes sobre el ajuste de escala automático. Para compilar este tipo de manifiesto de HPA, recopila todas las entradas del campo spec.metrics de cada recurso de HPA en un solo recurso de HPA. En el siguiente fragmento, se 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

Supervisa y prueba el ajuste de escala automático

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

Para observar el recuento 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 escalamiento de tu HPA, usa el siguiente comando, que envía un aumento de actividad de 100 solicitudes al extremo del modelo. Esto agotará las ranuras de decodificación disponibles y causará una acumulación de solicitudes en la cola de llenado previo, lo que activará HPA para aumentar 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 }'

¿Qué sigue?

Obtén información para optimizar el ajuste de escala automático de Pods según las métricas de Cloud Monitoring.
Obtén más información sobre el ajuste de escala automático horizontal de Pods en la documentación de código abierto de Kubernetes.