Best practice per l'assegnazione di priorità ai workload AI/ML di GKE

Questo documento descrive gli strumenti e le best practice per massimizzare l'utilizzo delle risorse e ridurre al minimo i tempi di inattività dei carichi di lavoro AI/ML eterogenei in Google Kubernetes Engine (GKE), in particolare quando non è disponibile capacità nelle prenotazioni o tramite risorse on demand. I workload eterogenei si riferiscono a diversi tipi di workload AI/ML che vengono eseguiti contemporaneamente nello stesso cluster GKE. Ad esempio, potresti eseguire un servizio di inferenza online sensibile alla latenza insieme a una serie di job di addestramento batch interrompibili.

Questa guida fornisce consigli per gli amministratori e gli operatori della piattaforma e per gli specialisti di dati e AI. Per una panoramica consolidata di tutte le best practice di GKE, consulta Best practice per GKE.

Vantaggi della definizione delle priorità dei workload AI/ML

I workload eterogenei hanno priorità diverse e condividono capacità e risorse limitate. Le best practice riportate in questa pagina descrivono come configurare GKE e gli strumenti open source per ottenere i seguenti vantaggi:

  • Ridurre al minimo i tempi di inattività per i carichi di lavoro ad alta priorità.
  • Esegui rapidamente i workload ad alta priorità.
  • Ottimizzare il consumo di risorse.

Sfondo

GKE supporta i seguenti strumenti open source per ottimizzare l'utilizzo delle risorse.

  • Kueue:un sistema di accodamento dei workload nativo di Kubernetes progettato per i workload batch, di AI e di computing ad alte prestazioni. Kueue può essere esteso per gestire altri tipi di carichi di lavoro, come quelli definiti da definizioni di risorse personalizzate come leaderworkerset. Kueue gestisce le quote e il modo in cui i workload le consumano in un cluster Kubernetes. Kueue prende decisioni su quando un workload attende, quando un workload viene avviato (ad esempio, creando il pod) e quando un pod appartenente a un workload viene prerilasciato.

    Per saperne di più su Kueue, consulta la documentazione sui concetti di Kueue.

  • Sostituzione a caldo:una tecnica che riduce il tempo medio di ripristino (MTTR). Lo scambio a caldo consente il preempt in base alla priorità del workload quando le risorse del cluster sono utilizzate completamente e non è disponibile capacità aggiuntiva, né da istanze on demand né da prenotazioni esistenti.

    • Quando un nodo che ospita un workload diventa non integro, il workload viene ripianificato sui nodi di riserva idonei. Se non sono disponibili nodi di riserva, Hotswap può interrompere un workload con priorità inferiore per fare spazio al workload in fase di recupero.
    • Se configuri i pod con PriorityClass, il workload configurato con priorità più alta espelle un workload in esecuzione con priorità bassa per acquisire le sue risorse. Questo processo di espulsione è noto come prerilascio.

Casi d'uso

Utilizza la seguente tabella per comprendere le best practice per ogni caso d'uso:

Caso d'uso Best practice Descrizione
Più workload con priorità diverse Utilizza Kueue per definire le code e assegnare le priorità ai workload in base alla loro importanza. Kueue può gestire le quote in modo che determinati team o progetti abbiano accesso a una quantità fissa di risorse.

Kueue ti consente di applicare le seguenti configurazioni:

  • Dai la priorità ai job ad alta priorità assegnando loro un valore WorkloadPriority più alto in Kueue.
  • Attiva l'accodamento con condivisione equa di Kueue in modo che tutti i workload ricevano risorse, anche quelli con priorità bassa.

Per testare la configurazione delle best practice, consulta l'esempio di Kueue in questo documento.

Devi ridurre l'MTTR attuale. Utilizza Hotswap per ripianificare i workload in risorse integre quando si verifica un'interruzione e prerilascia i workload con priorità bassa a favore di quelli con priorità alta.

Lo scambio a caldo ti consente di applicare le seguenti configurazioni:

  • Configura PriorityClasses per definire i livelli di priorità per i tuoi workload.
  • Assegna un valore PriorityClasses più alto ai workload critici.
  • Ripianifica automaticamente i workload sui nodi integri quando si verificano interruzioni.

Per testare la configurazione delle best practice, consulta l'esempio di sostituzione a caldo in questo documento.

Più workload di AI in competizione per risorse limitate Combina Kueue e Hotswap. Questa combinazione fornisce un sistema solido che assegna la priorità ai workload critici sia durante la pianificazione iniziale sia durante l'esecuzione.

Kueue e Hotswap ti consentono di applicare le seguenti configurazioni:

  • Utilizza Kueue per gestire la pianificazione e l'ammissione iniziali dei workload in base alla priorità.
  • Utilizza Hotswap per gestire le interruzioni del workload e consentire un rapido ripristino. La sostituzione a caldo aiuta a ridurre il tempo di ripristino di un workload ad alta priorità quando si verifica un'interruzione.

Per testare la configurazione delle best practice, consulta l'esempio di Kueue e Hotswap in questo documento.

Esempi di implementazioni di best practice

I seguenti esempi mostrano come implementare Kueue e Hotswap e come combinarli per le best practice descritte nella sezione precedente.

Kueue

Il seguente manifest di esempio mostra una configurazione di 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

Questo manifest esegue le seguenti operazioni:

  • Definisce un ResourceFlavor denominato tpu-v6e-slice che specifica le etichette dei nodi per gli slice TPU v6e.
  • Definisce un ClusterQueue denominato tpu-training-cq che gestisce la quota per le risorse TPU.
  • Definisce un LocalQueue denominato default-queue che consente ai workload nello spazio dei nomi default di utilizzare la coda cluster tpu-training-cq.

Hotswap

L'esempio seguente mostra una configurazione Hotswap che definisce due classi di priorità, low-priority-job e high-priority-job. Questa configurazione Hotswap crea un workload JobSet ad alta priorità e utilizza 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

In base a questa configurazione, Hotswap esegue le seguenti azioni:

  • Se un errore dell'infrastruttura interrompe il workload ad alta priorità, il JobSet lo riavvia. Hotswap prerilascia il workload a bassa priorità per ripianificare quello ad alta priorità prima che l'infrastruttura si ripristini. Il workload a bassa priorità rimane nello stato Non riuscito. Questo processo riduce notevolmente il tempo di inattività del workload.
  • Quando l'infrastruttura viene ripristinata, Hotswap riprogramma il workload a bassa priorità nel pool di nodi ripristinato.

Kueue e hot-swap

Combina Kueue e Hotswap quando operi in un ambiente complesso con risorse limitate. Questa combinazione fornisce un sistema solido che assegna la priorità ai workload critici durante la pianificazione iniziale e durante l'esecuzione.

L'esempio seguente mostra una configurazione combinata di Kueue e Hotswap. Questo esempio utilizza 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

In base a questa configurazione, Kueue viene combinato con Hotswap ed esegue le seguenti azioni:

  • Kueue gestisce l'ammissione di JobSet low-jax-trillium e high-jax-trillium nella coda del cluster in base alle priorità definite e alle risorse disponibili.
  • Se il JobSet high-jax-trillium viene interrotto da un errore dell'infrastruttura, l'hot swap prerilascia il JobSet low-jax-trillium per ripianificare il JobSet con priorità elevata.
  • Lo scambio a caldo garantisce il riavvio rapido di JobSet ad alta priorità, riducendo al minimo il suo tempo di inattività.
  • Quando l'infrastruttura viene ripristinata, Hotswap riprogramma il JobSet a bassa priorità nel pool di nodi recuperato.

Passaggi successivi