Otimize a priorização da carga de trabalho de IA/ML do GKE

Este documento descreve as ferramentas e as práticas recomendadas para maximizar a utilização de recursos e minimizar o tempo de inatividade das cargas de trabalho de IA/ML heterogéneas no Google Kubernetes Engine (GKE), especialmente quando não existe capacidade nas reservas ou através de recursos a pedido. As cargas de trabalho heterogéneas referem-se a diferentes tipos de cargas de trabalho de IA/ML que são executadas em simultâneo no mesmo cluster do GKE. Por exemplo, pode executar um serviço de inferência online sensível à latência juntamente com uma série de tarefas de preparação em lote interrompíveis.

Este guia oferece recomendações para administradores e operadores da plataforma, bem como especialistas em dados e IA.

Vantagens da priorização da carga de trabalho de IA/ML

As cargas de trabalho heterogéneas têm prioridades diferentes e partilham capacidade e recursos limitados. As práticas recomendadas nesta página descrevem como configurar o GKE e as ferramentas de código aberto para ajudar a obter as seguintes vantagens:

  • Minimize o tempo de inatividade para cargas de trabalho de alta prioridade.
  • Executar rapidamente cargas de trabalho de alta prioridade.
  • Otimize o consumo de recursos.

Contexto

O GKE suporta as seguintes ferramentas de código aberto para otimizar a utilização de recursos.

  • Kueue: um sistema de filas de cargas de trabalho nativo do Kubernetes concebido para cargas de trabalho de lotes, IA e computação de elevado desempenho. O Kueue pode ser expandido para gerir outros tipos de cargas de trabalho, como os definidos por definições de recursos personalizados, como leaderworkerset. O Kueue gere as quotas e a forma como as cargas de trabalho as consomem num cluster do Kubernetes. O Kueue toma decisões sobre quando uma carga de trabalho espera, quando uma carga de trabalho começa (por exemplo, criando o pod) e quando um pod pertencente a uma carga de trabalho é anulado.

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

  • Hotswap: uma técnica que reduz o tempo médio de recuperação (MTTR). A troca a quente permite a preempção com base na prioridade da carga de trabalho quando os recursos do cluster são totalmente utilizados e não está disponível capacidade adicional, quer seja de instâncias a pedido ou de reservas existentes.

    • Quando um nó que aloja uma carga de trabalho fica em mau estado, a carga de trabalho é reagendada em nós sobresselentes elegíveis. Se não estiverem disponíveis nós adicionais, a troca a quente pode substituir uma carga de trabalho de prioridade inferior para dar lugar à carga de trabalho a ser recuperada.
    • Se configurar os seus pods com PriorityClass, a carga de trabalho configurada com uma prioridade mais elevada desalojará uma carga de trabalho de baixa prioridade em execução para adquirir os respetivos recursos. Este processo de despejo é conhecido como preemptiva.

Exemplos de utilização

Use a tabela seguinte para compreender as práticas recomendadas para cada exemplo de utilização:

Exemplo de utilização Prática recomendada Descrição
Várias cargas de trabalho com prioridades diferentes Use o Kueue para definir filas e atribuir prioridades a cargas de trabalho com base na respetiva importância. O Kueue pode gerir a quota para que determinadas equipas ou projetos tenham acesso a uma determinada quantidade de recursos.

O Kueue permite-lhe aplicar as seguintes configurações:

  • Dê prioridade a trabalhos de alta prioridade atribuindo-lhes um valor de Kueue WorkloadPriority mais elevado.
  • Ative a fila de partilha equitativa do Kueue para que todas as cargas de trabalho acabem por receber recursos, mesmo as de baixa prioridade.

Para testar a configuração de práticas recomendadas, consulte o exemplo do Kueue neste documento.
Tem de reduzir o MTTR atual. Use a troca a quente para reagendar cargas de trabalho em recursos saudáveis quando ocorre uma interrupção e para antecipar cargas de trabalho de baixa prioridade em favor de cargas de trabalho de alta prioridade.

A troca a quente permite-lhe aplicar as seguintes configurações:

  • Configure PriorityClasses para definir níveis de prioridade para as suas cargas de trabalho.
  • Atribua um valor de PriorityClasses mais elevado a cargas de trabalho críticas.
  • Agendar novamente automaticamente cargas de trabalho em nós em bom estado quando ocorrem interrupções.

Para testar a configuração de práticas recomendadas, consulte o exemplo de troca a quente neste documento.
Várias cargas de trabalho de IA a competir por recursos limitados Combine o Kueue e o Hotswap. Esta combinação oferece um sistema robusto que dá prioridade às cargas de trabalho críticas durante o agendamento inicial e durante o tempo de execução.

O Kueue e o Hotswap permitem-lhe aplicar as seguintes configurações:

  • Use o Kueue para gerir o agendamento inicial e a admissão de cargas de trabalho com base na prioridade.
  • Use a troca a quente para processar interrupções de carga de trabalho e permitir uma recuperação rápida. A troca a quente 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 das práticas recomendadas, consulte o exemplo de Kueue e Hotswap neste documento.

Exemplos de implementações de práticas recomendadas

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

Kueue

O exemplo de manifesto seguinte 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

Este manifesto faz o seguinte:

  • Define um ResourceFlavor denominado tpu-v6e-slice que especifica as etiquetas dos nós para fatias de TPU v6e.
  • Define um ClusterQueue denominado tpu-training-cq que gere a quota para recursos de TPU.
  • Define um LocalQueue denominado default-queue que permite que as cargas de trabalho no espaço de nomes default usem a fila do cluster tpu-training-cq.

Hotswap

O exemplo seguinte mostra uma configuração de troca a quente que define duas classes de prioridade, low-priority-job e high-priority-job. Esta configuração de troca a quente cria uma carga de trabalho JobSet de alta prioridade e usa 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 nesta configuração, o Hotswap realiza as seguintes ações:

  • Se uma falha de infraestrutura interromper a carga de trabalho de alta prioridade, o JobSet reinicia-a. A troca a quente antecipa a carga de trabalho de baixa prioridade para reagendar a carga de trabalho de alta prioridade antes de a infraestrutura recuperar. A carga de trabalho de baixa prioridade permanece com o estado de falha. Este processo reduz significativamente o tempo de inatividade da carga de trabalho.
  • Quando a infraestrutura é recuperada, o Hotswap reagenda a carga de trabalho de baixa prioridade no conjunto de nós recuperado.

Kueue e Hotswap

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

O exemplo seguinte mostra uma configuração combinada do Kueue e do Hotswap. Este exemplo usa 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 nesta configuração, o Kueue é combinado com o Hotswap e realiza as seguintes ações:

  • O Kueue gere a admissão de low-jax-trillium e high-jax-trillium JobSets na fila do cluster com base nas respetivas prioridades definidas e recursos disponíveis.
  • Se o high-jax-trillium JobSet for interrompido por uma falha de infraestrutura, o Hotswap antecipa o low-jax-trillium JobSet para reagendar o JobSet de alta prioridade.
  • A troca a quente garante que o JobSet de alta prioridade é reiniciado rapidamente, minimizando o respetivo tempo de inatividade.
  • Quando a infraestrutura é recuperada, a troca a quente reagenda o JobSet de baixa prioridade no conjunto de nós recuperado.

O que se segue?