Pianifica i workload GKE con Topology Aware Scheduling (TAS)

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

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

Che cos'è la pianificazione sensibile alla topologia (TAS)?

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

La TAS funziona al meglio con la 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 la TAS potrebbe non funzionare bene in questo scenario.

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti attività:

  • Abilita l'API Google Kubernetes Engine.
    • Abilita 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 gcloud components update comando. 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 cluster.

Prepara il cluster GKE

Per preparare il cluster GKE all'esecuzione dei workload con la TAS, completa i seguenti passaggi:

  1. Installa Kueue con la TAS abilitata

  2. Visualizza la topologia del cluster GKE

  3. Configura Kueue

Installa Kueue con la TAS abilitata

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

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

Manifest Kueue

  1. Installa Kueue:

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

  2. Abilita la 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 la TAS abilitata 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 Kueue, 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 relativa topologia fisica per la TAS. In caso contrario, si verificano 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, fai riferimento alle 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. In genere, questo nome host è anche il nome del nodo GKE.

Più valori di etichetta condividono due VM, più le VM sono fisicamente vicine l'una all'altra. Per ulteriori informazioni su questi termini, consulta Terminologia.

Configura Kueue

Dopo aver installato Kueue, devi configurarlo per specificare l'infrastruttura che gestisce. In genere, Kueue richiede una ClusterQueue definizione della quota di risorse di un'infrastruttura statica o di un'infrastruttura 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 workload utilizzando la TAS nel seguente modo:

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

  • Workload 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 della quota di risorse ClusterQueue a Kueue, esamina i seguenti esempi:

  • Quota molto elevata: Kueue non interrompe quasi mai l'ammissione di un workload in base alle risorse richieste. In base alle definizioni della TAS, Kueue potrebbe ammettere o meno i workload in base alla topologia dell'infrastruttura. Per ulteriori informazioni, consulta Quota di risorse molto elevata.

  • Quota realistica: Kueue ammette il workload solo se le risorse richieste dal workload rientrano in questi limiti di quota di risorse. In base alle definizioni della TAS, Kueue controlla la topologia dell'infrastruttura prima di ammettere il workload. Per ulteriori informazioni, consulta Quota di 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 workload in base alla quota di risorse disponibile. Al contrario, Kueue utilizza le informazioni sulla topologia dei nodi disponibili per provare a far corrispondere la topologia ai requisiti del workload.

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

  1. Apri un editor di file a tua scelta. Quindi, includi la seguente definizione della 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 di risorse realistica

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

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

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

  1. Apri un editor di file a tua scelta. Quindi, includi la seguente definizione della 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 workload con la TAS utilizzando Kueue

I seguenti scenari mostrano come puoi indicare a Kueue e alla TAS di gestire le combinazioni comuni di workload e infrastrutture utilizzando i tipi di richieste di topologia e i livelli di richieste di topologia:

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

    • kueue.x-k8s.io/podset-preferred-topology: Kueue assegna 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 rientrare 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 provare ad ammettere questo workload finché l'intero workload non può rientrare nel livello di topologia scelto.

  • Di seguito sono riportati i livelli di richieste di topologia disponibili, che ti consentono di essere più o meno specifico sull'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 workload 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 per l'esecuzione all'interno di un nome host (consigliato) Questa configurazione ammetterà il workload purché siano disponibili risorse sufficienti a soddisfare i requisiti di risorse del workload, anche se la capacità è frammentata. Kueue pianificherà i 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 workload se e solo se è disponibile un host con risorse sufficienti a soddisfare i requisiti di risorse del workload.

    Questa opzione è utile quando sono presenti più VM per host (ad esempio, tipi di macchine più piccoli) o 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 per l'esecuzione all'interno di un host Questa configurazione ammetterà il workload purché 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 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 sottoblocco Questa configurazione ammetterà il 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 per l'esecuzione all'interno di un sottoblocco Questa configurazione ammetterà il workload purché 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 sottoblocco e utilizzerà sottoblocchi aggiuntivi, se necessario. In questo caso, Kueue assegnerà un rango più alto a un sottoblocco con una capacità disponibile maggiore, anche se frammentata, rispetto a un sottoblocco con una 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 workload se e solo se le risorse disponibili all'interno di un blocco soddisfano i requisiti di risorse del workload. Se ammesso, Kueue ridurrà al minimo il numero di sottoblocchi 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 per l'esecuzione all'interno di un blocco Questa configurazione ammetterà il workload purché 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 workload utilizzando PodGroup con la 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 seguenti mostrano come creare un PodGroup di tre pod con un leader e due worker.

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

Se il leader è uno dei worker e richiede anche una GPU, il leader può avere qualsiasi numero all'interno del 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"

Caso 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 applicherà le assegnazioni di rango.

Vedi l'esempio seguente:

---
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