Pianifica i workload GKE con Topology Aware Scheduling (TAS)

I workload di AI e ML richiedono una comunicazione significativa tra pod. A causa di questo requisito, la larghezza di banda di rete tra i pod influisce direttamente sul tempo di esecuzione e sul costo del carico di lavoro. Questa larghezza di banda dipende dal posizionamento delle istanze di macchine virtuali (VM) nel cluster.

Questo documento spiega come ottimizzare la pianificazione dei carichi di lavoro di AI o ML su larga scala su un cluster Google Kubernetes Engine (GKE) per prestazioni e affidabilità. In particolare, configuri il cluster in modo che utilizzi la pianificazione consapevole della topologia (TAS) per la comunicazione a bassa latenza. Questo approccio riduce al minimo il sovraccarico di comunicazione e contribuisce a massimizzare le prestazioni dei tuoi workload.

Che cos'è la pianificazione topologia consapevole (TAS)?

TAS può migliorare significativamente l'efficienza dell'addestramento di modelli linguistici di grandi dimensioni (LLM). TAS posiziona strategicamente i worker nella topologia di rete per ridurre al minimo l'overhead di comunicazione durante l'aggregazione del gradiente, che richiede ai worker di comunicare in un ordine di classificazione specifico. Riducendo al minimo gli hop di rete tra i worker che comunicano in sequenza, TAS riduce la contesa di rete e ottimizza l'utilizzo della larghezza di banda, il che porta a una convergenza più rapida e a tempi di addestramento più brevi. Con modelli LLM sempre più grandi, TAS è essenziale per massimizzare le prestazioni e la scalabilità dell'addestramento distribuito.

TAS funziona al meglio con una capacità posizionata in modo denso, che può essere ottenuta tramite le prenotazioni. Con le VM con avvio flessibile o le VM spot, è meno probabile che la capacità venga allocata in modo ravvicinato, quindi TAS potrebbe non funzionare bene in questo scenario.

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti operazioni:

  • Attiva l'API Google Kubernetes Engine.
  • Attiva l'API Google Kubernetes Engine
  • Se vuoi utilizzare Google Cloud CLI per questa attività, installala e poi inizializza gcloud CLI. Se hai già installato gcloud CLI, scarica l'ultima versione eseguendo il comando gcloud components update. Le versioni precedenti di gcloud CLI potrebbero non supportare l'esecuzione dei comandi in questo documento.
  • Per connetterti al cluster, esegui questo comando:

    gcloud container clusters get-credentials CLUSTER_NAME
    

    Sostituisci CLUSTER_NAME con il nome del tuo cluster.

Prepara il cluster GKE

Per preparare il cluster GKE a eseguire carichi di lavoro con TAS, completa i seguenti passaggi:

  1. Installa Kueue con TAS abilitato

  2. Visualizza la topologia del cluster GKE

  3. Configura Kueue

Installa Kueue con TAS abilitato

Ti consigliamo di utilizzare TAS con Kueue, un sistema nativo di Kubernetes che gestisce le quote e il modo in cui i job dovrebbero utilizzarle. TAS richiede Kueue versione 0.10.0 o successive e devi abilitarlo esplicitamente.

Per installare Kueue e attivare TAS, seleziona una delle seguenti opzioni:

Manifest di Kueue

  1. Installa Kueue:

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/v0.10.0/manifests.yaml
    
  2. Abilita TAS in Kueue:

    kubectl -n kueue-system patch deployment kueue-controller-manager \
        --type json -p='[{"op": "add", "path": "/spec/template/spec/containers/0/args/-", "value": "--feature-gates=TopologyAwareScheduling=true"}]'
    

Grafico Helm

Installa Kueue con TAS abilitato utilizzando un grafico Helm:

helm install kueue oci://us-central1-docker.pkg.dev/k8s-staging-images/charts/kueue \
    --version="v0.10.0" \
    --create-namespace \
    --namespace=kueue-system \
    --set="controllerManager.featureGates[0].name=TopologyAwareScheduling,controllerManager.featureGates[0].enabled=true"

Dopo aver installato Kueke, devi configurarlo in modo che comprenda l'infrastruttura che gestisce, come spiegato nella sezione successiva.

Visualizza la topologia del cluster GKE

Prima di visualizzare la topologia dei nodi A4X, A4, A3 Ultra, A3 Mega e A3 High (8 GPU) di cui è stato eseguito il provisioning come VM spot, devi definire il posizionamento compatto sui nodi GKE per esporre la loro topologia fisica per TAS. In caso contrario, si verificano degli errori.

Per visualizzare la topologia dei nodi del cluster GKE in un pool di nodi specifico, esegui questo comando:

kubectl get nodes -l cloud.google.com/gke-nodepool=NODE_POOL_NAME \
    -ocustom-columns='NAME:.metadata.name,BLOCK:.metadata.labels.cloud\.google\.com/gce-topology-block,SUBBLOCK:.metadata.labels.cloud\.google\.com/gce-topology-subblock,HOST:.metadata.labels.cloud\.google\.com/gce-topology-host' | sort -k2,4

Sostituisci NODE_POOL_NAME con il nome del pool di nodi.

Per comprendere la topologia fisica dei nodi GKE nelle VM nell'output, consulta le seguenti etichette dei nodi:

  • cloud.google.com/gce-topology-block: l'ID specifico dell'organizzazione del blocco riservato in cui si trova la VM.

  • cloud.google.com/gce-topology-subblock: l'ID specifico dell'organizzazione del sottoblocco in cui si trova la VM.

  • cloud.google.com/gce-topology-host: l'ID dell'host su cui si trova la VM.

  • kubernetes.io/hostname: il nome host del nodo Kubernetes. Questo nome host in genere corrisponde anche al nome del nodo GKE.

Più valori di etichetta condividono due VM, più le VM sono fisicamente vicine tra loro. Per saperne di più su questi termini, consulta Terminologia.

Configura Kueue

Dopo aver installato Kueue, devi configurarlo per specificare l'infrastruttura che gestisce. In genere, Kueue richiede una definizione di quota di risorse ClusterQueue di infrastruttura statica o dinamica con la scalabilità automatica del cluster abilitata. ClusterQueue ammette un Workload solo se le risorse richieste dal workload sono inferiori o uguali al pool di risorse definito in ClusterQueue. Dopo aver configurato Kueue come descritto in questa sezione, Kueue ammette i carichi di lavoro utilizzando TAS nel seguente modo:

  • Carichi di lavoro TAS: Kueue controlla sia la topologia dell'infrastruttura fisica sia il suo utilizzo attuale.

  • Carichi di lavoro non TAS: Kueue non controlla la topologia dell'infrastruttura fisica. Kueue gestisce l'intera quota definita nella configurazione e lascia l'assegnazione dei nodi a kube-scheduler.

Per capire come fornire una definizione di quota di risorse ClusterQueue a Kueue, esamina i seguenti esempi:

  • Quota molto elevata: Kueue non interrompe praticamente mai l'ammissione di un carico di lavoro in base alle risorse richieste. In base alle definizioni TAS, Kueue può o non ammettere carichi di lavoro in base alla topologia dell'infrastruttura. Per ulteriori informazioni, consulta Quota delle risorse molto elevata.

  • Quota realistica: Kueue ammette il workload solo se le risorse richieste rientrano nei limiti di queste quote di risorse. In base alle definizioni TAS, Kueue controlla quindi la topologia dell'infrastruttura prima di ammettere il workload. Per saperne di più, consulta Quota delle risorse realistica.

Tutti i riferimenti alla quota di risorse nelle sezioni seguenti si riferiscono alla quota di risorse ClusterQueue.

Quota di risorse molto elevata

L'esempio seguente utilizza una quota di risorse molto elevata, in modo che Kueue non interrompa mai un carico di lavoro in base alla quota di risorse disponibile. Kueue utilizza invece le informazioni sulla topologia dei nodi disponibili per cercare di abbinare la topologia ai requisiti del workload.

Per utilizzare la seguente definizione di quota delle risorse, completa i seguenti passaggi:

  1. Apri un editor di file a tua scelta. Quindi, includi la seguente definizione di quota in un file YAML denominato kueue-tas-config-very-high-quota.yaml:

      apiVersion: kueue.x-k8s.io/v1alpha1
      kind: Topology
      metadata:
        name: "gke-default"
      spec:
        levels:
        - nodeLabel: "cloud.google.com/gce-topology-block"
        - nodeLabel: "cloud.google.com/gce-topology-subblock"
        - nodeLabel: "cloud.google.com/gce-topology-host"
        - nodeLabel: "kubernetes.io/hostname"
    ---
      kind: ResourceFlavor
      apiVersion: kueue.x-k8s.io/v1beta1
      metadata:
        name: "tas-flavor"
      spec:
        nodeLabels:
          cloud.google.com/gke-nodepool: "NODE_POOL_NAME"
        topologyName: "gke-default"
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: NoSchedule
    ---
      apiVersion: kueue.x-k8s.io/v1beta1
      kind: ClusterQueue
      metadata:
        name: "tas-cluster-queue"
      spec:
        namespaceSelector: {}
        resourceGroups:
        - coveredResources: ["nvidia.com/gpu"]
          flavors:
          - name: "tas-flavor"
            resources:
            - name: "nvidia.com/gpu"
              nominalQuota: 10000000
    ---
      apiVersion: kueue.x-k8s.io/v1beta1
      kind: LocalQueue
      metadata:
        namespace: "default"
        name: "tas-user-queue"
      spec:
        clusterQueue: "tas-cluster-queue"
    

    Sostituisci NODE_POOL_NAME con il nome del pool di nodi.

  2. Crea e applica la configurazione della quota di risorse per il sistema di accodamento dei job Kueue:

    kubectl create -f kueue-tas-config-very-high-quota.yaml
    

Quota per le risorse realistica

L'esempio precedente configurava solo le risorse GPU. Tuttavia, Kueue può gestire tutte le risorse compatibili con Kubernetes.

L'esempio seguente definisce una quota di risorse più realistica, tra cui CPU, memoria e GPU. Questo è per 100 macchine a3-ultragpu-8g. Una singola macchina ha 224 vCPU, 2944 GB di memoria e 8 GPU.

Per utilizzare la seguente definizione di quota delle risorse, completa i seguenti passaggi:

  1. Apri un editor di file a tua scelta. Quindi, includi la seguente definizione di quota in un file YAML denominato kueue-tas-config-real-quota.yaml:

      apiVersion: kueue.x-k8s.io/v1alpha1
      kind: Topology
      metadata:
        name: "gke-default"
      spec:
        levels:
        - nodeLabel: "cloud.google.com/gce-topology-block"
        - nodeLabel: "cloud.google.com/gce-topology-subblock"
        - nodeLabel: "cloud.google.com/gce-topology-host"
        - nodeLabel: "kubernetes.io/hostname"
    ---
      kind: ResourceFlavor
      apiVersion: kueue.x-k8s.io/v1beta1
      metadata:
        name: "tas-flavor"
      spec:
        nodeLabels:
          cloud.google.com/gke-nodepool: "NODE_POOL_NAME"
        topologyName: "gke-default"
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: NoSchedule
    ---
      apiVersion: kueue.x-k8s.io/v1beta1
      kind: ClusterQueue
      metadata:
        name: "tas-cluster-queue"
      spec:
        namespaceSelector: {} # match all
        resourceGroups:
        - coveredResources: ["cpu", "memory", "nvidia.com/gpu"]
          flavors:
          - name: "tas-flavor"
            resources:
            # numbers below represent quota of 100 a3-ultragpu-8g machines
            - name: "cpu"
              nominalQuota: 22400
            - name: "memory"
              nominalQuota: 294400Gi
            - name: "nvidia.com/gpu"
              nominalQuota: 800
    ---
      apiVersion: kueue.x-k8s.io/v1beta1
      kind: LocalQueue
      metadata:
        namespace: "default"
        name: "tas-user-queue"
      spec:
        clusterQueue: "tas-cluster-queue"
    

    Sostituisci NODE_POOL_NAME con il nome del pool di nodi.

  2. Crea e applica una configurazione della quota di risorse per il sistema di accodamento dei job Kueue:

    kubectl create -f kueue-tas-config-real-quota.yaml
    

    L'output è simile al seguente:

    topology.kueue.x-k8s.io/gke-default created
    resourceflavor.kueue.x-k8s.io/tas-flavor created
    clusterqueue.kueue.x-k8s.io/tas-cluster-queue created
    localqueue.kueue.x-k8s.io/tas-user-queue created
    

Pianifica i carichi di lavoro con TAS utilizzando Kueue

Gli scenari seguenti mostrano come puoi indicare a Kueue e TAS di gestire combinazioni comuni di carichi di lavoro e infrastrutture utilizzando i tipi di richieste di topologia e i livelli di richieste di topologia:

  • Di seguito sono riportati i tipi di richiesta di topologia disponibili (preferiti o obbligatori):

    • kueue.x-k8s.io/podset-preferred-topology: Kueue dà la priorità alla pianificazione dell'intero workload all'interno di un determinato livello di topologia, ma ammette comunque un workload che non rientra in questo livello di topologia. Per un workload che potrebbe essere stato inserito in un singolo livello di topologia, Kueue potrebbe pianificare il workload su più istanze di quel livello di topologia.

    • kueue.x-k8s.io/podset-required-topology: Kueue continua a tentare di ammettere questo workload finché l'intero workload non può rientrare nel livello di topologia scelto.

  • Di seguito sono riportati i livelli disponibili per la richiesta di topologia, che ti consentono di essere più o meno specifico in merito all'infrastruttura fisica in cui preferisci o richiedi l'esecuzione del job:

    • cloud.google.com/gce-topology-block

    • cloud.google.com/gce-topology-subblock

    • cloud.google.com/gce-topology-host

    • kubernetes.io/hostname

Per pianificare i carichi di lavoro utilizzando questi valori, utilizza il seguente file YAML del job:

apiVersion: batch/v1
kind: Job
metadata:
  generateName: JOB_NAME
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
spec:
  parallelism: NUMBER_OF_REPLICAS
  completions: NUMBER_OF_REPLICAS
  completionMode: Indexed
  template:
    metadata:
      annotations:
        ANNOTATIONS_STRING
    spec:
      containers:
      - name: dummy-job
        image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
        args: ["60s"]
        resources:
          requests:
            nvidia.com/gpu: "1"
          limits:
            nvidia.com/gpu: "1"
      restartPolicy: Never

Sostituisci le seguenti variabili:

  • JOB_NAME: un nome per il job.

  • NUMBER_OF_REPLICAS: il numero di pod in esecuzione in parallelo.

  • ANNOTATIONS_STRING: consulta la tabella seguente:

    Tipo e livello di topologia richiesti Descrizione ANNOTATIONS_STRING
    Preferito da eseguire all'interno di un nome host (opzione consigliata) Questa configurazione ammetterà il tuo workload a condizione che siano disponibili risorse sufficienti a soddisfare i requisiti di risorse del workload, anche se la capacità è frammentata. Kueue pianificherà i tuoi pod nel modo più compatto possibile. kueue.x-k8s.io/podset-preferred-topology: "kubernetes.io/hostname"
    Obbligatorio per l'esecuzione all'interno di un host

    Questa configurazione ammetterà il tuo workload se e solo se è disponibile un host con risorse sufficienti a soddisfare i requisiti di risorse del workload.

    Ciò è utile quando ci sono più VM per host (ad esempio, tipi di macchine più piccoli) o quando più pod possono essere eseguiti su un singolo nodo. In questi casi, se il workload viene ammesso, verrà eseguito su un singolo host.

    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-host"
    Preferito da eseguire all'interno di un host Questa configurazione ammetterà il tuo workload a condizione che siano disponibili risorse sufficienti a soddisfare i requisiti di risorse del workload, anche se la capacità è frammentata. Kueue tenterà di pianificare i tuoi pod all'interno di un host e utilizzerà host aggiuntivi se necessario. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-host"
    Obbligatorio per l'esecuzione all'interno di un blocco secondario Questa configurazione ammetterà il tuo workload se e solo se è disponibile un sottoblocco con risorse sufficienti a soddisfare i requisiti di risorse del workload. kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-subblock"
    Preferito da eseguire all'interno di un blocco secondario Questa configurazione ammetterà il tuo workload a condizione che siano disponibili risorse sufficienti a soddisfare i requisiti di risorse del workload, anche se la capacità è frammentata. Kueue tenterà di pianificare i tuoi pod all'interno di un sottoblocco e utilizzerà sottoblocchi aggiuntivi, se necessario. In questo caso, Kueue assegnerà un ranking più alto a un sottoblocco con più capacità disponibile, anche se è frammentato, rispetto a un sottoblocco con capacità appena sufficiente a soddisfare i requisiti. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-subblock"
    Obbligatorio per l'esecuzione all'interno di un blocco Questa configurazione ammetterà il tuo workload se e solo se le risorse disponibili all'interno di un blocco soddisfano i requisiti di risorse del tuo workload. Se viene accettato, Kueue ridurrà al minimo il numero di blocchi secondari e host per pianificare il workload. Ciò potrebbe comportare la frammentazione della capacità disponibile. kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
    Preferito da eseguire all'interno di un blocco Questa configurazione ammetterà il tuo workload a condizione che siano disponibili risorse sufficienti a soddisfare i requisiti di risorse del workload, anche se la capacità è frammentata. Kueue tenterà di pianificare i pod all'interno di un blocco e utilizzerà blocchi aggiuntivi se necessario. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-block"

Pianifica i carichi di lavoro utilizzando PodGroup con TAS utilizzando Kueue

Quando utilizzi PodGroup, devi specificare tre campi aggiuntivi per ogni pod in un PodGroup:

A seconda del framework ML che utilizzi, un leader di un PodGroup può richiedere o meno una GPU. A causa di una limitazione di Kueue, questi casi devono essere gestiti in modo diverso. Gli esempi riportati di seguito mostrano come creare un PodGroup di tre pod con un leader e due worker.

Scenario 1: il leader è anche un worker e richiede una GPU

Se il leader è uno dei worker e richiede anche una GPU, può avere qualsiasi numero all'interno di PodGroup. Per semplicità, nell'esempio seguente l'indice del leader è 0:

apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-leader-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "0"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  containers:
  - name: leader
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
  restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-1-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "1"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-2-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "2"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

Scenario 2: il leader non è un worker e non richiede una GPU

Se il leader non è uno dei worker a causa della limitazione di Kueue, il leader deve avere l'ultimo indice nel PodGroup, a causa del modo in cui Kueue crea i PodSet. Se il leader non ha l'ultimo indice e il primo worker non utilizza il primo indice, Kueue non applica le assegnazioni di ranking.

Vedi il seguente esempio:

---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-leader-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "2"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  containers:
  - name: leader
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        cpu: "1"
      limits:
        cpu: "1"
  restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-0-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "0"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-1-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "1"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

Passaggi successivi