Configurar o escalonamento automático para cargas de trabalho LLM em TPUs

Autopilot Standard

Nesta página, mostramos como configurar sua infraestrutura de escalonamento automático usando o escalonador automático horizontal de pods (HPA, na sigla em inglês) do GKE para implantar o modelo de linguagem grande (LLM) do Gemma usando o JetStream de host único.

Para saber mais sobre a seleção de métricas para escalonamento automático, consulte Práticas recomendadas para cargas de trabalho LLM com escalonamento automático com TPUs no GKE.

Antes de começar

Antes de começar, verifique se você realizou as tarefas a seguir:

Ativar a API Google Kubernetes Engine.

Ativar a API Google Kubernetes Engine

Se você quiser usar a Google Cloud CLI para essa tarefa, instale e, em seguida, inicialize a CLI gcloud. Se você instalou a CLI gcloud anteriormente, instale a versão mais recente executando gcloud components update.
Observação: para instalações atuais da CLI gcloud, defina as propriedades compute/region e compute/zone. Ao definir locais padrão, é possível evitar erros na CLI gcloud, como: One of [--zone, --region] must be supplied: Please specify location.

Conheça e conclua o fluxo de trabalho em Disponibilizar o Gemma usando TPUs no GKE com o JetStream. Verifique se o argumento PROMETHEUS_PORT está definido no manifesto de implantação do JetStream.

Escalonamento automático usando métricas

É possível usar as métricas de desempenho específicas da carga de trabalho emitidas pelo servidor de inferência do JetStream ou as métricas de desempenho do TPU para direcionar o escalonamento automático para seus pods.

Para configurar o escalonamento automático com métricas, siga estas etapas:

Exportar as métricas do servidor JetStream para o Cloud Monitoring. Você usa o Google Cloud Managed Service para Prometheus, que simplifica a implantação e a configuração do seu coletor Prometheus. O Google Cloud Managed Service para Prometheus é ativado por padrão nos cluster do GKE. Também é possível ativá-la manualmente.

O manifesto de exemplo a seguir mostra como configurar as definições do recurso PodMonitoring para direcionar o Google Cloud Managed Service para Prometheus para coletar métricas dos pods em intervalos recorrentes de 15 segundos:

Se você precisar coletar as métricas do servidor, use o manifesto a seguir. Com as métricas de servidor, são permitidos intervalos de raspagem de dados de até cinco 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
```
Se você precisar extrair métricas de TPU, use o manifesto a seguir. Com as métricas do sistema, há suporte para intervalos de raspagem de dados de até 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
```

Instale um adaptador de métricas. Com esse adaptador, as métricas do servidor exportadas para o Monitoring ficam visíveis para o controlador HPA. Para mais detalhes, consulte Escalonamento automático horizontal de pods na documentação do Google Cloud Managed Service para Prometheus.

Se você quiser que o JetStream seja escalonado com métricas individuais, use o adaptador de métricas personalizadas do Stackdriver.
Se você quiser que o JetStream seja dimensionado com o valor de uma expressão composta por várias métricas distintas, use o Adaptador Prometheus de terceiros.

Adaptador de métricas personalizadas do Stackdriver

O adaptador de métricas personalizadas do Stackdriver dá suporte a métricas de consulta do Google Cloud Managed Service para Prometheus, começando com a versão v0.13.1 do adaptador.

Para instalar o adaptador de métricas personalizadas do Stackdriver, faça o seguinte:

Configure a coleção gerenciada no cluster.

Instale o adaptador de métricas personalizadas do Stackdriver no cluster.

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

Se você tiver a Federated Identity for Workload do GKE ativada no cluster do Kubernetes e usar a federação de identidade da carga de trabalho para o GKE, também precisará conceder o papel de leitor do Monitoring para a conta de serviço em que o adaptador é executado.

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 Prometheus

Considere estas considerações ao usar prometheus-adapter para escalonar com o Google Cloud Managed Service para Prometheus:

Encaminhe consultas pelo proxy da interface do front-end do Prometheus, assim como ao consultar o Google Cloud Managed Service para Prometheus usando a API ou a interface do Prometheus. Esse front-end será instalado em uma etapa posterior.
Por padrão, o argumento prometheus-url da implantação prometheus-adapter está definido como --prometheus-url=http://frontend.default.svc:9090/, em que default é o namespace em que você implantou o front-end. Se você implantou o front-end em outro namespace, configure esse argumento de acordo.
No campo .seriesQuery da configuração de regras, não é possível usar uma correspondência de expressão regular (regex) no nome de uma métrica. Em vez disso, especifique totalmente os nomes das métricas.

Como os dados podem levar um pouco mais de tempo para serem disponibilizados em O Google Cloud Managed Service para Prometheus em comparação com o Prometheus upstream, configurar uma lógica de escalonamento automático excessivamente ávida pode causar comportamentos indesejados. Não há garantia de atualização de dados, mas eles geralmente ficam disponíveis para consulta de três a sete segundos após o envio para o Google Cloud Managed Service para Prometheus, excluindo qualquer latência de rede.

Todas as consultas emitidas por prometheus-adapter têm escopo global. Isso significa que, se você tiver aplicativos em dois namespaces que emitem métricas com nomes idênticos, uma configuração do HPA que usa essa métrica é escalonada usando dados dos dois aplicativos. Para evitar o escalonamento com dados incorretos, sempre use os filtros namespace ou cluster no PromQL.

Para definir um exemplo de configuração do HPA usando prometheus-adapter e a coleta gerenciada, siga estas etapas:

Configure a coleção gerenciada no cluster.

Implante o proxy de IU do front-end do Prometheus no cluster. Crie o seguinte manifesto com o nome 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

Em seguida, aplique o manifesto:

kubectl apply -f prometheus-frontend.yaml

Verifique se o prometheus-adapter está instalado no cluster. Para isso, instale o gráfico do helm prometheus-community/prometheus-adapter. Crie o arquivo values.yaml a seguir:

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"})

Em seguida, use este arquivo como o arquivo de valores para implantar o gráfico do 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 você usar a Federação de identidade da carga de trabalho para o GKE, também precisará configurar e autorizar uma conta de serviço executando os seguintes comandos:

Primeiro, crie as contas de serviço no cluster e do Google Cloud:

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

Em seguida, vincule as duas contas de serviço:

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

Em seguida, atribua o papel monitoring.viewer à conta de serviço do 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 fim, defina a conta de serviço de implantações do front-end como a nova conta de serviço no cluster:
```
kubectl set serviceaccount deployment frontend prom-frontend-sa
```

Configurar o recurso HPA com base em métricas. Implante um recurso HPA baseado na métrica de servidor que preferir. Para mais detalhes, consulte Escalonamento automático horizontal de pods na documentação do Google Cloud Managed Service para Prometheus. A configuração específica do HPA depende do tipo de métrica (servidor ou TPU) e de qual adaptador de métrica está instalado.

Alguns valores são necessários em todas as configurações do HPA e devem ser definidos para criar um recurso do HPA:
- MIN_REPLICAS: o número mínimo de réplicas de pod do JetStream permitidas. Se você não modificar o manifesto de implantação do JetStream na etapa Implantar o JetStream, recomendamos definir esse valor como 1.
- MAX_REPLICAS: o número máximo de réplicas de pod do JetStream permitido. O exemplo de implantação do JetStream requer 8 chips por réplica, e o pool de nós contém 16 chips. Se você quiser manter a latência de escalonamento vertical baixa, defina como 2. Valores maiores acionam o escalonador automático de cluster para criar novos nós no pool de nós, aumentando assim a latência de escalonamento vertical.
- TARGET: a média desejada para essa métrica em todas as instâncias do JetStream. Consulte a Documentação do Kubernetes para escalonamento automático para mais informações sobre como o número de réplicas é determinado com base nesse valor.
Adaptador de métricas personalizadas do Stackdriver
O adaptador de métricas personalizadas do Stackdriver permite o escalonamento da carga de trabalho com o valor médio das consultas de métricas individuais do Google Cloud Managed Service para Prometheus em todos os pods. Ao usar o adaptador de métricas personalizadas do Stackdriver, recomendamos o escalonamento com as métricas de servidor jetstream_prefill_backlog_size e jetstream_slots_used_percentage e a métrica do TPU memory_used.

Para criar um manifesto de HPA para escalonamento com métricas do servidor, crie o seguinte arquivo 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
```
Ao usar o adaptador de métricas personalizadas do Stackdriver com métricas de TPU, recomendamos usar apenas a métrica kubernetes.io|node|accelerator|memory_used para escalonamento. Para criar um manifesto HPA para escalonamento com essa métrica, crie o seguinte arquivo 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 Prometheus
O adaptador do Prometheus permite o escalonamento da carga de trabalho com o valor de consultas PromQL do Google Cloud Managed Service para Prometheus. Anteriormente, você definiu as métricas de servidor jetstream_prefill_backlog_size e jetstream_slots_used_percentage que representam o valor médio em todos os pods.

Para criar um manifesto de HPA para escalonamento com métricas do servidor, crie o seguinte arquivo 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 criar um manifesto do HPA para escalonamento com métricas de TPU, recomendamos usar apenas o memory_used_percentage definido no arquivo de valores do helm do adaptador do Prometheus. memory_used_percentage é o nome dado à consulta do PromQL a seguir, que reflete a memória média atual usada em todos os aceleradores:
```
avg(kubernetes_io:node_accelerator_memory_used{cluster_name="CLUSTER_NAME"}) / avg(kubernetes_io:node_accelerator_memory_total{cluster_name="CLUSTER_NAME"})
```
Para criar um manifesto de HPA para escalonamento com memory_used_percentage, crie o seguinte arquivo 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
```

Dimensionar usando várias métricas

Também é possível configurar o escalonamento com base em várias métricas. Para saber como o número de réplicas é determinado usando várias métricas, consulte a documentação do Kubernetes sobre o escalonamento automático. Para criar esse tipo de manifesto, colete todas as entradas do campo spec.metrics de cada recurso HPA em um único recurso HPA. O snippet a seguir mostra um exemplo de como agrupar os recursos do 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

Monitorar e testar o escalonamento automático

É possível observar como as cargas de trabalho do JetStream são dimensionadas com base na configuração do HPA.

Para observar a contagem de réplicas em tempo real, execute o seguinte comando:

kubectl get hpa --watch

A saída deste comando será semelhante a esta:

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

Para testar a capacidade de escalonamento do HPA, use o comando a seguir, que envia um burst de 100 solicitações para o endpoint do modelo. Isso esgotará os slots de decodificação disponíveis e causará um backlog de solicitações na fila de preenchimento, acionando o HPA a aumentar o tamanho da implantação do 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 }'

A seguir

Saiba como otimizar o escalonamento automático de pods com base nas métricas do Cloud Monitoring.
Saiba mais sobre o escalonamento automático horizontal de pods na documentação do Kubernetes de código aberto.