Otimizar a priorização de cargas de trabalho de IA/ML do GKE

Este documento descreve ferramentas e práticas recomendadas para maximizar a utilização de recursos e minimizar o tempo de inatividade de cargas de trabalho heterogêneas de IA/ML no Google Kubernetes Engine (GKE), especialmente quando não há capacidade em reservas ou por recursos sob demanda. Cargas de trabalho heterogêneas se referem a diferentes tipos de cargas de trabalho de IA/ML que são executadas simultaneamente no mesmo cluster do GKE. Por exemplo, é possível executar um serviço de inferência on-line sensível à latência junto com uma série de jobs de treinamento em lote interrompíveis.

Este guia fornece recomendações para administradores e operadores de plataforma e especialistas em dados e IA.

Benefícios da priorização de cargas de trabalho de IA/ML

As cargas de trabalho heterogêneas têm prioridades diferentes e compartilham capacidade e recursos limitados. As práticas recomendadas nesta página descrevem como configurar o GKE e ferramentas de código aberto para ajudar você a ter os seguintes benefícios:

  • Minimizar o tempo de inatividade para cargas de trabalho de alta prioridade.
  • Execute rapidamente cargas de trabalho de alta prioridade.
  • Otimizar o consumo de recursos.

Contexto

O GKE é compatível com as seguintes ferramentas de código aberto para otimizar o uso de recursos.

  • Kueue:um sistema de enfileiramento de cargas de trabalho nativo do Kubernetes projetado para cargas de trabalho de computação em lote, de IA e de alto desempenho. O Kueue pode ser estendido para gerenciar outros tipos de carga de trabalho, como os definidos por definições de recursos personalizados como leaderworkerset. O Kueue gerencia cotas e como as cargas de trabalho as consomem em um cluster do Kubernetes. O Kueue decide quando uma carga de trabalho espera, quando uma carga de trabalho é iniciada (por exemplo, criando o pod) e quando um pod pertencente a uma carga de trabalho é interrompido.

    Para mais informações sobre o Kueue, consulte a documentação de conceitos do Kueue.

  • Troca a quente:uma técnica que reduz o tempo médio de recuperação (MTTR). O Hotswap permite a remoção com base na prioridade da carga de trabalho quando os recursos do cluster estão totalmente utilizados e não há capacidade adicional disponível, seja de instâncias sob demanda ou reservas atuais.

    • Quando um nó que hospeda uma carga de trabalho fica sem integridade, ela é reprogramada em nós sobressalentes qualificados. Se não houver nós extras disponíveis, o Hotswap poderá interromper uma carga de trabalho de baixa prioridade para abrir espaço para a carga de trabalho que está sendo recuperada.
    • Se você configurar seus pods com PriorityClass, a carga de trabalho configurada com maior prioridade vai remover uma carga de trabalho de baixa prioridade em execução para adquirir os recursos dela. Esse processo de remoção é conhecido como remoção por substituição.

Casos de uso

Use a tabela a seguir para entender as práticas recomendadas de cada caso de uso:

Caso de uso Prática recomendada Descrição
Várias cargas de trabalho com prioridades diferentes Use o Kueue para definir filas e atribuir prioridades às cargas de trabalho com base na importância delas. O Kueue pode gerenciar a cota para que determinadas equipes ou projetos tenham acesso a uma quantidade definida de recursos.

Com o Kueue, é possível aplicar as seguintes configurações:

  • Priorize jobs de alta prioridade atribuindo a eles um WorkloadPriority mais alto do Kueue.
  • Ative o enfileiramento de compartilhamento justo do Kueue para que todas as cargas de trabalho recebam recursos, mesmo as de baixa prioridade.

Para testar a configuração de prática recomendada, consulte o exemplo do Kueue neste documento.
É preciso reduzir o MTTR atual. Use o Hotswap para reagendar cargas de trabalho em recursos íntegros quando ocorrer uma interrupção e para remover cargas de trabalho de baixa prioridade em favor de cargas de trabalho de alta prioridade.

O hot swap permite aplicar as seguintes configurações:

  • Configure PriorityClasses para definir níveis de prioridade para suas cargas de trabalho.
  • Atribua um PriorityClasses maior a cargas de trabalho críticas.
  • Remarcar automaticamente cargas de trabalho em nós íntegros quando ocorrerem interrupções.

Para testar a configuração de prática recomendada, consulte o exemplo de troca a quente neste documento.
Várias cargas de trabalho de IA competindo por recursos limitados Combine Kueue e Hotswap. Essa combinação oferece um sistema robusto que prioriza cargas de trabalho críticas durante o agendamento inicial e o tempo de execução.

Com o Kueue e o Hotswap, é possível aplicar as seguintes configurações:

  • Use o Kueue para gerenciar o agendamento e a admissão iniciais de cargas de trabalho com base na prioridade.
  • Use o Hotswap para lidar com interrupções de carga de trabalho e permitir uma recuperação rápida. O hotswap ajuda a reduzir o tempo de recuperação de uma carga de trabalho de alta prioridade quando ocorre uma interrupção.

Para testar a configuração de prática recomendada, consulte o exemplo de Kueue e Hotswap neste documento.

Exemplos de implementações de práticas recomendadas

Os exemplos a seguir demonstram como implementar o Kueue e o Hotswap e como combiná-los para as práticas recomendadas descritas na seção anterior.

Kueue

O exemplo de manifesto a seguir mostra uma configuração do Kueue:

  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ResourceFlavor
  metadata:
    name: tpu-v6e-slice
  spec:
    nodeLabels:
      cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: tpu-training-cq
  spec:
    resourceGroups:
    - flavors:
      - name: tpu-v6e-slice
        resources:
        - name: google.com/tpu
          nominalQuota: 32
    queueingStrategy: BestEffortFIFO
    preemption:
      reclaimWithinCohort: Never
      reclaimOutOfCohort:
        enable: true
        reclaimMoreThanNominalQuota: false
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    name: default-queue
    namespace: default
  spec:
    clusterQueue: tpu-training-cq

Esse manifesto faz o seguinte:

  • Define um ResourceFlavor chamado tpu-v6e-slice que especifica os rótulos de nó para frações de TPU v6e.
  • Define um ClusterQueue chamado tpu-training-cq que gerencia a cota para recursos de TPU.
  • Define um LocalQueue chamado default-queue que permite que cargas de trabalho no namespace default usem a fila do cluster tpu-training-cq.

Hotswap

O exemplo a seguir mostra uma configuração de troca a quente que define duas classes de prioridade, low-priority-job e high-priority-job. Essa configuração do Hotswap cria uma carga de trabalho do JobSet de alta prioridade e usa o MaxText.

  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: low-priority-job
  value: 1000000
  globalDefault: false
  description: "This priority class should be used for low priority pods only."
  ---
  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: high-priority-job
  value: 2000000
  globalDefault: false
  description: "This priority class should be used for critical pods only."
  ---
  apiVersion: jobset.x-k8s.io/v1alpha2
  kind: JobSet
  metadata:
    name: high-jax-trillium
    annotations:
      alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
  spec:
    failurePolicy:
      maxRestarts: 10
      restartStrategy: BlockingRecreate
    replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          backoffLimit: 0
          completions: 4
          parallelism: 4
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              priorityClassName: high-priority-job
              containers:
              - name: jax-program
                image: <IMAGE LOCATION>
                command:
                -   python3
                -   MaxText/train.py
                -   MaxText/configs/base.yml
                -   model_name=llama2-7b
                -   run_name=<UNIQUE RUN NAME>
                -   steps=300
                -   base_output_directory=gs://<OUTPUT BUCKET>
                -   dataset_path=gs://max-datasets-rogue
                -   max_target_length=4096
                -   dataset_type=synthetic
                -   enable_checkpointing=False
                resources:
                  limits:
                    google.com/tpu: 4

Com base nessa configuração, o Hotswap executa as seguintes ações:

  • Se uma falha de infraestrutura interromper a carga de trabalho de alta prioridade, o JobSet vai reiniciá-la. O Hotswap força a interrupção da carga de trabalho de baixa prioridade para reagendar a de alta prioridade antes que a infraestrutura se recupere. A carga de trabalho de baixa prioridade permanece com status de falha. Esse processo reduz significativamente o tempo ocioso da carga de trabalho.
  • Quando a infraestrutura é recuperada, o Hotswap reagenda a carga de trabalho de baixa prioridade no pool de nós recuperado.

Kueue e Hotswap

Combine o Kueue e o Hotswap quando você opera em um ambiente complexo com recursos limitados. Essa combinação oferece um sistema robusto que prioriza cargas de trabalho críticas durante o agendamento inicial e o tempo de execução.

O exemplo a seguir mostra uma configuração combinada do Kueue e do Hotswap. Este exemplo usa o MaxText:

  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: low-priority-job
  value: 1000000
  globalDefault: false
  description: "This priority class should be used for low priority pods only."
  ---
  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: high-priority-job
  value: 2000000
  globalDefault: false
  description: "This priority class should be used for critical pods only."
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ResourceFlavor
  metadata:
    name: tpu-v6e-slice
  spec:
    nodeLabels:
      cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: tpu-training-cq
  spec:
    resourceGroups:
    - flavors:
      - name: tpu-v6e-slice
        resources:
        - name: google.com/tpu
          nominalQuota: 32
    queueingStrategy: BestEffortFIFO
    preemption:
      reclaimWithinCohort: Never
      reclaimOutOfCohort:
        enable: true
        reclaimMoreThanNominalQuota: false
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    name: default-queue
    namespace: default
  spec:
    clusterQueue: tpu-training-cq
  ---
  apiVersion: jobset.x-k8s.io/v1alpha2
  kind: JobSet
  metadata:
    name: low-jax-trillium
    annotations:
      kueue.x-k8s.io/queue-name: default-queue
      alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
  spec:
    failurePolicy:
      maxRestarts: 10
      restartStrategy: BlockingRecreate
    replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          backoffLimit: 0
          completions: 4
          parallelism: 4
          template:
            metadata:
              labels:
                kueue.x-k8s.io/managed-by: kueue
                kueue.x-k8s.io/priority-class: low-priority-job
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              priorityClassName: low-priority-job
              containers:
              - name: jax-program
                image: <IMAGE LOCATION>
                command:
                - python3
                - MaxText/train.py
                - MaxText/configs/base.yml
                - model_name=llama2-7b
                - run_name=low-priority-run
                - steps=30000
                - base_output_directory=gs://<OUTPUT BUCKET>
                - dataset_path=gs://max-datasets-rogue
                - max_target_length=4096
                - dataset_type=synthetic
                - enable_checkpointing=False
                resources:
                  limits:
                    google.com/tpu: 4
  ---
  apiVersion: jobset.x-k8s.io/v1alpha2
  kind: JobSet
  metadata:
    name: high-jax-trillium
    annotations:
      kueue.x-k8s.io/queue-name: default-queue
      alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
  spec:
    failurePolicy:
      maxRestarts: 10
      restartStrategy: BlockingRecreate
    replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          backoffLimit: 0
          completions: 4
          parallelism: 4
          template:
            metadata:
              labels:
                kueue.x-k8s.io/managed-by: kueue
                kueue.x-k8s.io/priority-class: high-priority-job
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              priorityClassName: high-priority-job
              containers:
              - name: jax-program
                image: <IMAGE LOCATION>
                command:
                - python3
                - MaxText/train.py
                - MaxText/configs/base.yml
                - model_name=llama2-7b
                - run_name=high-priority-run
                - steps=300
                - base_output_directory=gs://<OUTPUT BUCKET>
                - dataset_path=gs://max-datasets-rogue
                - max_target_length=4096
                - dataset_type=synthetic
                - enable_checkpointing=False
                resources:
                  limits:
                    google.com/tpu: 4

Com base nessa configuração, o Kueue é combinado com o Hotswap e realiza as seguintes ações:

  • O Kueue gerencia a admissão de low-jax-trillium e high-jax-trillium JobSets na fila do cluster com base nas prioridades definidas e nos recursos disponíveis.
  • Se o JobSet high-jax-trillium for interrompido por uma falha de infraestrutura, o Hotswap vai antecipar o JobSet low-jax-trillium para reprogramar o JobSet de alta prioridade.
  • O hotswap garante que o JobSet de alta prioridade seja reiniciado rapidamente, minimizando o tempo ocioso.
  • Quando a infraestrutura é recuperada, o Hotswap reagenda o JobSet de baixa prioridade no pool de nós recuperado.

A seguir