Ottimizza la definizione delle priorità dei 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 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.

Vantaggi della definizione delle priorità dei workload AI/ML

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

  • Ridurre al minimo i tempi di inattività per i carichi di lavoro ad alta priorità.
  • Esegui rapidamente i carichi di lavoro ad alta priorità.
  • Ottimizza il consumo delle 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 carichi di lavoro batch, AI e di computing ad alte prestazioni. Kueue può essere esteso per gestire altri tipi di carichi di lavoro, ad esempio quelli definiti da definizioni di risorse personalizzate come leaderworkerset. Kueue gestisce le quote e il modo in cui i carichi di lavoro le utilizzano in un cluster Kubernetes. Kueue decide quando un workload deve attendere, quando deve iniziare (ad esempio creando il pod) e quando un pod appartenente a un workload viene preempted.

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

  • Sostituzione a caldo:una tecnica che riduce il tempo medio di ripristino (MTTR). Lo scambio a caldo consente la preemptive 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 carico di lavoro non è integro, il carico di lavoro 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 tuoi pod con PriorityClass, il workload configurato con priorità più elevata espelle un workload a bassa priorità in esecuzione per acquisirne le risorse. Questa procedura di espulsione è nota 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 priorità ai carichi di lavoro 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 di Kueue più elevato.
  • Attiva la coda di condivisione equa di Kueue in modo che tutti i workload ricevano risorse, anche quelli a bassa priorità.

Per testare la configurazione delle best practice, consulta l'esempio di Kueue in questo documento.
Devi ridurre l'MTTR attuale. Utilizza Hotswap per riprogrammare i workload nelle risorse integre quando si verifica un'interruzione e per prerilasciare i workload a bassa priorità a favore di quelli ad alta priorità.

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 di PriorityClasses più alto ai carichi di lavoro 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 dei carichi di lavoro e consentire un rapido recupero. Lo scambio a caldo consente di ridurre il tempo di ripristino di un carico di lavoro ad alta priorità in caso di interruzione.

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

Esempi di implementazioni delle best practice

Gli esempi seguenti 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 svolge 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 del cluster tpu-training-cq.

Sostituzione a caldo

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 carico di lavoro 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à, JobSet lo riavvia. Lo scambio a caldo ha la precedenza sul carico di lavoro a bassa priorità per riprogrammare il carico di lavoro ad alta priorità prima che l'infrastruttura venga ripristinata. Il carico di lavoro 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 Hotswap

Combina Kueue e Hotswap quando operi in un ambiente complesso con risorse limitate. Questa combinazione fornisce un sistema solido che assegna la priorità ai carichi di lavoro 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 low-jax-trillium e high-jax-trillium JobSet nella coda del cluster in base alle priorità definite e alle risorse disponibili.
  • Se il high-jax-trillium JobSet viene interrotto da un errore dell'infrastruttura, Hotswap esegue l'interruzione preventiva del JobSet low-jax-trillium per riprogrammare il JobSet a 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 ripristinato.

Passaggi successivi