Configure a escala automática para cargas de trabalho de MDIs em TPUs

Esta página mostra como configurar a sua infraestrutura de escalamento automático através do Horizontal Pod Autoscaler (HPA) do GKE para implementar o grande modelo de linguagem (GML) Gemma através do JetStream de anfitrião único.

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

Antes de começar

Antes de começar, certifique-se de que realizou as seguintes tarefas:

  • Ative a API Google Kubernetes Engine.
    • Ative a API Google Kubernetes Engine
  • Se quiser usar a CLI gcloud para esta tarefa, instale-a e, em seguida, inicialize-a. Se instalou anteriormente a CLI gcloud, execute o comando gcloud components update para obter a versão mais recente. As versões anteriores da CLI gcloud podem não suportar a execução dos comandos neste documento.

Use a escala automática com métricas

Pode usar as métricas de desempenho específicas da carga de trabalho emitidas pelo servidor de inferência JetStream ou as métricas de desempenho da TPU para direcionar o dimensionamento automático dos seus pods.

Para configurar o dimensionamento automático com métricas, siga estes passos:

  1. Exporte as métricas do servidor JetStream para o Cloud Monitoring. Usa o Google Cloud Managed Service for Prometheus, que simplifica a implementação e a configuração do seu coletor do Prometheus. O Google Cloud Managed Service for Prometheus está ativado por predefinição no seu cluster do GKE. Também pode ativá-lo manualmente.

    O exemplo de manifesto seguinte mostra como configurar as definições de recursos PodMonitoring para direcionar o serviço gerido do Google Cloud para o Prometheus de modo a extrair métricas dos seus pods a intervalos recorrentes de 15 segundos:

    Se precisar de extrair métricas do servidor, use o seguinte manifesto. Com as métricas do servidor, são suportados intervalos de recolha tão frequentes quanto 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

    Se precisar de extrair métricas de TPU, use o seguinte manifesto. Com as métricas do sistema, são suportados intervalos de recolha tão frequentes quanto 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. Instale um adaptador de métricas. Este adaptador torna as métricas do servidor que exportou para o Monitoring visíveis para o controlador HPA. Para mais detalhes, consulte o artigo Ajuste de escala automático de pods horizontal na documentação do serviço gerido do Google Cloud para Prometheus.

    Adaptador do Stackdriver de métricas personalizadas

    O adaptador do Stackdriver de métricas personalizadas suporta a consulta de métricas do Google Cloud Managed Service for Prometheus, a partir da versão v0.13.1 do adaptador.

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

    1. Configure a coleção gerida no seu cluster.

    2. Instale o adaptador do Stackdriver de métricas personalizadas no 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 tiver a Workload Identity Federation para o GKE ativada no seu cluster do Kubernetes e usar a Workload Identity Federation para o GKE, também tem de conceder a função de leitor do Monitoring à conta de serviço na qual o adaptador é executado. Substitua PROJECT_ID pelo ID do seu projeto.

    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 do Prometheus

    Tenha em atenção estas considerações quando usar o prometheus-adapter para dimensionar através do Google Cloud Managed Service for Prometheus:

    • Encaminhar consultas através do proxy da IU de front-end do Prometheus, tal como quando consulta o Google Cloud Managed Service for Prometheus através da API ou da IU do Prometheus. Este front-end é instalado num passo posterior.
    • Por predefinição, o argumento prometheus-url da implementação está definido como --prometheus-url=http://frontend.default.svc:9090/, em que default é o espaço de nomes onde implementou o front-end.prometheus-adapter Se implementou o frontend noutro espaço de nomes, configure este argumento em conformidade.
    • No campo .seriesQuery da configuração das regras, não pode usar um correspondente de expressão regular (regex) num nome de métrica. Em alternativa, especifique totalmente os nomes das métricas.

    Uma vez que os dados podem demorar um pouco mais a ficar disponíveis no Google Cloud Managed Service for Prometheus em comparação com o Prometheus a montante, a configuração de uma lógica de escala automática demasiado ansiosa pode causar um comportamento indesejável. Embora não haja garantia de atualização dos dados, os dados estão normalmente disponíveis para consulta 3 a 7 segundos após serem enviados para o serviço gerido da Google Cloud para o Prometheus, excluindo qualquer latência da rede.

    Todas as consultas emitidas pela prometheus-adapter têm âmbito global. Isto significa que, se tiver aplicações em dois espaços de nomes que emitem métricas com nomes idênticos, uma configuração do HPA que use essa métrica é dimensionada com base nos dados de ambas as aplicações. Para evitar a escalabilidade com dados incorretos, use sempre filtros namespace ou cluster no seu PromQL.

    Para configurar uma configuração de HPA de exemplo com prometheus-adapter e recolha gerida, siga estes passos:

    1. Configure a coleção gerida no seu cluster.

    2. Implemente o proxy da IU do frontend do Prometheus no seu 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

    3. Certifique-se de que o prometheus-adapter está instalado no seu cluster instalando o gráfico Helm prometheus-community/prometheus-adapter. Crie o seguinte ficheiro 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"})

      Em seguida, use este ficheiro como o ficheiro de valores para implementar o seu gráfico 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 usar a Federação de identidades de cargas de trabalho para o GKE, também tem de configurar e autorizar uma conta de serviço executando os seguintes comandos:

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

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

    2. Em seguida, associe as duas contas de serviço. Certifique-se de que substitui PROJECT_ID pelo ID do seu projeto:

      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. De seguida, atribua a função monitoring.viewer à Google Cloud conta de serviço:

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

    4. Por último, defina a conta de serviço das implementações de frontend como a nova conta de serviço no cluster:

      kubectl set serviceaccount deployment frontend prom-frontend-sa

  3. Configure o recurso HPA baseado em métricas. Implemente um recurso HPA com base na sua métrica de servidor preferida. Para mais detalhes, consulte o artigo Ajuste de escala automático de pods horizontal na documentação do serviço gerido do Google Cloud para o Prometheus. A configuração específica do HPA depende do tipo de métrica (servidor ou TPU) e do adaptador de métricas instalado.

    São necessários alguns valores em todas as configurações do HPA e têm de ser definidos para criar um recurso do HPA:

    • MIN_REPLICAS: o número mínimo de réplicas de pods do JetStream permitidas. Se não modificar o manifesto de implementação do JetStream a partir do passo Implementar JetStream, recomendamos que defina este valor como 1.
    • MAX_REPLICAS: o número máximo de réplicas do agrupamento JetStream permitidas. A implementação do JetStream de exemplo requer 8 chips por réplica e o conjunto de nós contém 16 chips. Se quiser manter a latência de aumento da escala baixa, defina este valor como 2. Valores maiores acionam o Cluster Autoscaler para criar novos nós no node pool, o que aumenta a latência de expansão.

    • TARGET: a média segmentada para esta métrica em todas as instâncias do JetStream. Consulte a documentação do Kubernetes sobre o redimensionamento automático para mais informações sobre como a contagem de réplicas é determinada a partir deste valor.

    Adaptador do Stackdriver de métricas personalizadas

    O adaptador do Stackdriver de métricas personalizadas suporta o dimensionamento da sua carga de trabalho com o valor médio das consultas de métricas individuais do serviço gerido do Google Cloud para o Prometheus em todos os pods. Quando usar o adaptador do Stackdriver de métricas personalizadas, recomendamos o dimensionamento com as métricas do servidor jetstream_prefill_backlog_size e jetstream_slots_used_percentage e a métrica da TPU memory_used.

    Para criar um manifesto HPA para o dimensionamento com métricas do servidor, crie o seguinte ficheiro 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 usar o adaptador do Stackdriver de métricas personalizadas com métricas de TPU, recomendamos que use apenas a métrica kubernetes.io|node|accelerator|memory_used para o escalamento. Para criar um manifesto HPA para o escalamento com esta métrica, crie o seguinte ficheiro 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 do Prometheus

    O adaptador do Prometheus suporta o dimensionamento da sua carga de trabalho com o valor das consultas PromQL do Google Cloud Managed Service for Prometheus. Anteriormente, 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 HPA para o dimensionamento com métricas do servidor, crie o seguinte ficheiro 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 HPA para o escalamento com métricas de TPU, recomendamos que use apenas o memory_used_percentage definido no ficheiro de valores do Helm do prometheus-adapter. memory_used_percentage é o nome dado à seguinte consulta PromQL que reflete a média atual da memória 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 HPA para o escalamento com memory_used_percentage, crie o seguinte ficheiro 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

Expanda a escala com várias métricas

Também pode configurar o ajuste de escala com base em várias métricas. Para saber como a quantidade de réplicas é determinada através de várias métricas, consulte a documentação do Kubernetes sobre o dimensionamento automático. Para criar este tipo de manifesto de HPA, recolha todas as entradas do campo spec.metrics de cada recurso de HPA num único recurso de HPA. O fragmento seguinte mostra um exemplo de como pode agrupar os recursos 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

Monitorize e teste a escala automática

Pode observar como as suas 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

O resultado deste comando deve ser semelhante ao seguinte:

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

Para testar a capacidade de escalabilidade do HPA, use o seguinte comando, que envia uma rajada de 100 pedidos para o ponto final do modelo. Isto esgota os espaços de descodificação disponíveis e provoca um atraso de pedidos na fila de preenchimento, o que aciona o HPA para aumentar o tamanho da implementaçã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 }'

O que se segue?