Orchestra i carichi di lavoro multi-sezione utilizzando JobSet e Kueue

Questo tutorial mostra come orchestrare più carichi di lavoro multislice su Google Kubernetes Engine (GKE) per migliorare l'utilizzo delle risorse. Esegui il deployment di un carico di lavoro Jax come esempio, eseguilo su TPU Multislice e implementa la coda dei job con JobSet e Kueue. Kueue determina quando devono essere eseguiti i job in base alle risorse disponibili, alle quote e a una gerarchia per la condivisione equa tra i team.

Questo tutorial è rivolto a sviluppatori di machine learning (ML), amministratori e operatori della piattaforma interessati alle funzionalità di orchestrazione dei container di Kubernetes per l'addestramento di LLM. Per scoprire di più sui ruoli comuni e sulle attività di esempio a cui facciamo riferimento nei Google Cloud contenuti, consulta Ruoli e attività utente comuni di GKE Enterprise.

Prima di leggere questa pagina, assicurati di conoscere quanto segue:

Obiettivi

  1. Prepara l'ambiente con un cluster GKE con tre slice TPU v5e. Ogni slice TPU ha una topologia 2x4 con 8 chip. Pertanto, in totale 24 chip TPU v5e.
  2. Crea le risorse Kueue per assicurarti che le quote vengano condivise equamente tra i carichi di lavoro.
  3. Esegui il carico di lavoro Multislice.

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à, installa e poi inizializza gcloud CLI. Se hai già installato l'interfaccia a riga di comando gcloud, ottieni la versione più recente eseguendo gcloud components update.

Prepara l'ambiente

  1. Nella console Google Cloud, avvia un'istanza Cloud Shell:
    Apri Cloud Shell

  2. Imposta le variabili di ambiente predefinite utilizzando il comando gcloud config set:

    gcloud config set project PROJECT_ID

    Sostituisci PROJECT_ID con il tuo Google Cloud ID progetto.

I cluster Autopilot che eseguono la versione 1.29.2-gke.1521000 o successive attivano le TPU per impostazione predefinita. Le TPU sui cluster Autopilot vengono configurate nella specifica del workload. Per ulteriori informazioni, consulta la sezione Definire i carichi di lavoro multislice con JobSet.

Crea un cluster GKE

In Cloud Shell, crea un cluster GKE:

Autopilot 

gcloud container clusters create-auto multislice-cluster \
    --location=LOCATION \
    --cluster-version 1.29.2-gke.1521000 \
    --release-channel rapid

In questo comando:

  • Il flag --location specifica la località Compute Engine del cluster.
  • Il flag --cluster-version specifica la versione di Kubernetes per il cluster.
  • Il flag --release-channel specifica il canale di rilascio per il cluster. In questo caso, il canale rapido supporta le versioni più recenti disponibili in GKE.

Standard 

gcloud container clusters create multislice-cluster \
    --location=LOCATION

Sostituisci LOCATION con la posizione in cui vuoi creare il cluster. Assicurati che abbia la capacità per il tipo di macchina ct5lp-hightpu-4t. La creazione del cluster potrebbe richiedere diversi minuti.

Se utilizzi la modalità GKE Autopilot, vai alla sezione Creare le risorse Kueue. I cluster Autopilot che eseguono la versione 1.29.2-gke.1521000 o successive attivano le TPU per impostazione predefinita.

Crea tre pool di nodi di slice TPU in modalità standard

In questa sezione, crei pool di nodi TPU utilizzando il comando gcloud beta container node-pools create.

  1. Crea il primo node pool denominato nodepool1:

    gcloud beta container node-pools create nodepool1 \
    --location=LOCATION \
    --cluster=multislice-cluster \
    --node-locations=NODE_LOCATION \
    --machine-type=ct5lp-hightpu-4t \
    --tpu-topology=2x4 \
    --project=PROJECT_ID

    Sostituisci NODE_LOCATION con una o più zone nella regione del cluster in cui vuoi creare i nodi.

  2. Crea il secondo node pool denominato nodepool2:

    gcloud beta container node-pools create nodepool2 \
    --location=LOCATION \
    --cluster=multislice-cluster \
    --node-locations=NODE_LOCATION \
    --machine-type=ct5lp-hightpu-4t \
    --tpu-topology=2x4 \
    --project=PROJECT_ID

  3. Crea il terzo node pool denominato nodepool3:

    gcloud beta container node-pools create nodepool3 \
    --location=LOCATION \
    --cluster=multislice-cluster \
    --node-locations=NODE_LOCATION \
    --machine-type=ct5lp-hightpu-4t \
    --tpu-topology=2x4 \
    --project=PROJECT_ID

GKE crea tre node pool. Ogni pool di nodi è una sezione TPU separata.

Nei passaggi precedenti, hai utilizzato il comando gcloud beta container node-pools create per creare i pool di nodi. Questi comandi utilizzano i seguenti flag:

  • --node-locations: l'elenco separato da virgole di una o più zone in cui GKE crea i node pool.
  • --machine-type: il tipo di macchina da utilizzare per i nodi. In questo caso, hai utilizzato ct5lp-hightpu-4t. Per ulteriori informazioni sui tipi di macchine compatibili con le TPU, consulta la tabella in Scegliere la versione TPU.
  • --tpu-topology: la topologia TPU da utilizzare per il pool di nodi. In questo caso, hai utilizzato 2x4. Per saperne di più sulle topologie TPU, consulta Scegliere la topologia TPU.

Crea le risorse Kueue

  1. Crea il seguente manifest kueue.yaml:

    apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "vlp-24"
spec:
  nodeLabels:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: "cluster-queue"
spec:
  namespaceSelector: {}
  queueingStrategy: BestEffortFIFO
  resourceGroups:
  - coveredResources: ["google.com/tpu"]
    flavors:
    - name: "vlp-24"
      resources:
      - name: "google.com/tpu"
        nominalQuota: 24

---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  namespace: default
  name: multislice-queue
spec:
  clusterQueue: cluster-queue

  2. Applica il manifest kueue.yaml:

    kubectl apply -f kueue.yaml

GKE crea le seguenti risorse Kueue:

  • ResourceFlavor: un'astrazione delle risorse in un cluster. In questo esempio, GKE crea tre sezioni TPU con topologia 2x4. Ogni sezione TPU ha una topologia 2x4 con 8 chip (24 chip TPU in totale).
  • ClusterQueue: una coda globale che gestisce i carichi di lavoro e le risorse del cluster.
  • LocalQueue: agrupa i workload strettamente correlati che in genere vengono eseguiti da un singolo tenant (utente). Ogni coda locale rimanda a una coda cluster da cui vengono allocate le risorse per l'esecuzione dei relativi carichi di lavoro. Un carico di lavoro Kueue è un'astrazione che rappresenta un carico di lavoro batch. In questo caso, ogni carico di lavoro è un JobSet.

Definisci i tuoi carichi di lavoro multislice con JobSet

In questa sezione, crei tre JobSet. Un jobset è un'API di carico di lavoro che ti consente di gestire un gruppo di job Kubernetes come un'unità. Il caso d'uso più comune per un JobSet è l'addestramento distribuito, ma puoi utilizzarlo anche per eseguire carichi di lavoro batch.

I seguenti JobSet eseguono un carico di lavoro Jax che restituisce il numero globale di chip TPU nel segmento, quindi rimane inattivo per 60 secondi per simulare il tempo di addestramento del modello, quindi esce.

  1. Installa l'API JobSet nel cluster:

    VERSION=v0.8.1
kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/$VERSION/manifests.yaml

  2. Crea il seguente manifest jobsets-multislice.yaml:

    Autopilot 

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-1slice
  labels:
    kueue.x-k8s.io/queue-name: multislice-queue
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: 1
      template:
        spec:
          parallelism: 2
          completions: 2
          backoffLimit: 0
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 2x4
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                command:
                - bash
                - -c
                - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  python -c 'import jax; print("Global device count:", jax.device_count())'
                resources:
                  limits:
                    google.com/tpu: 4

---
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-2slice
  labels:
    kueue.x-k8s.io/queue-name: multislice-queue
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          parallelism: 2
          completions: 2
          backoffLimit: 0
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 2x4
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                command:
                - bash
                - -c
                - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  python -c 'import jax; print("Global device count:", jax.device_count())'
                  sleep 60
                resources:
                  limits:
                    google.com/tpu: 4
---
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-3slice
  labels:
    kueue.x-k8s.io/queue-name: multislice-queue
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: 3
      template:
        spec:
          parallelism: 2
          completions: 2
          backoffLimit: 0
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 2x4
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                command:
                - bash
                - -c
                - |
                  sleep 60
                resources:
                  limits:
                    google.com/tpu: 4

    Standard 

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-1slice
  labels:
    kueue.x-k8s.io/queue-name: multislice-queue
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: 1
      template:
        spec:
          parallelism: 2
          completions: 2
          backoffLimit: 0
          template:
            spec:
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 2x4
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                securityContext:
                  privileged: true
                command:
                - bash
                - -c
                - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  python -c 'import jax; print("Global device count:", jax.device_count())'
                resources:
                  limits:
                    google.com/tpu: 4

---
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-2slice
  labels:
    kueue.x-k8s.io/queue-name: multislice-queue
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          parallelism: 2
          completions: 2
          backoffLimit: 0
          template:
            spec:
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 2x4
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                securityContext:
                  privileged: true
                command:
                - bash
                - -c
                - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  python -c 'import jax; print("Global device count:", jax.device_count())'
                  sleep 60
                resources:
                  limits:
                    google.com/tpu: 4
---
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-3slice
  labels:
    kueue.x-k8s.io/queue-name: multislice-queue
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: 3
      template:
        spec:
          parallelism: 2
          completions: 2
          backoffLimit: 0
          template:
            spec:
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 2x4
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                securityContext:
                  privileged: true
                command:
                - bash
                - -c
                - |
                  sleep 60
                resources:
                  limits:
                    google.com/tpu: 4

  3. Applica il manifest jobsets-multislice.yaml:

    kubectl apply -f jobsets-multislice.yaml

GKE crea i job con le seguenti richieste di risorse:

  • Il JobSet multislice-1slice crea un job che richiede un totale di una sezione TPU.
  • Il JobSet multislice-2slice crea due job che richiedono in totale due sezioni TPU.
  • Il JobSet multislice-3slice crea tre job che richiedono in totale tre sezioni TPU.

Poiché il cluster ha solo tre slice TPU, non tutti i JobSet possono essere eseguiti contemporaneamente. Quando Kueue mette in coda tutti e tre i JobSet di multislice-3slice, i relativi job vengono eseguiti da soli fino al completamento. multislice-1slice e multislice-2slice rimangono in attesa e vengono eseguiti insieme in un secondo momento.

Verifica che Kueue abbia accettato i carichi di lavoro

  1. Controlla i carichi di lavoro in coda in Kueue:

    kubectl get workloads

    L'output è simile al seguente:

    NAME                             QUEUE              ADMITTED BY     AGE
jobset-multislice-1slice-2530a   multislice-queue                   3s
jobset-multislice-2slice-ffb02   multislice-queue                   4s
jobset-multislice-3slice-8c695   multislice-queue   cluster-queue   10s

Kueue mette in coda uno o più workload, a seconda delle risorse TPU richieste.

Monitora i workload

  1. Monitora i pod in esecuzione:

    kubectl get pods

    L'output è simile al seguente:

    NAME                                READY   STATUS      RESTARTS   AGE
multislice-1slice-slice-0-0-pf2ll   1/1     Running     0          1s
multislice-1slice-slice-0-1-55g62   1/1     Running     0          1s
multislice-2slice-slice-0-0-f4hf7   1/1     Running     0          3s
multislice-2slice-slice-0-1-c8kv7   1/1     Running     0          3s
multislice-2slice-slice-1-0-7h46t   1/1     Running     0          3s
multislice-2slice-slice-1-1-lj9hb   1/1     Running     0          3s
multislice-3slice-slice-0-0-wzq9t   0/1     Completed   0          2m31s
multislice-3slice-slice-0-1-zf4dp   0/1     Completed   0          2m30s
multislice-3slice-slice-1-0-hbfn5   0/1     Completed   0          2m31s
multislice-3slice-slice-1-1-45fgl   0/1     Completed   0          2m30s
multislice-3slice-slice-2-0-wjbp4   0/1     Completed   0          2m30s
multislice-3slice-slice-2-1-lwnvs   0/1     Completed   0          2m30s

    Verifica che GKE abbia pianificato, creato ed eseguito i pod per multislice-3slice. GKE ha poi eseguito i pod da multislice-1slice e multislice-2slice JobSet.

Attivare le priorità e la preemption dei carichi di lavoro Kueue

Se vuoi, puoi assegnare priorità ai carichi di lavoro di Kueue che determinano l'ordine in cui i carichi di lavoro in coda vengono ammessi da Kueue.

  1. Aggiorna ClusterQueue in modo da avere un criterio di prelazione:

    apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "vlp-24"
spec:
  nodeLabels:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: "cluster-queue"
spec:
  namespaceSelector: {}
  resourceGroups:
  - coveredResources: ["google.com/tpu"]
    flavors:
    - name: "vlp-24"
      resources:
      - name: "google.com/tpu"
        nominalQuota: 24
 preemption:
    reclaimWithinCohort: Any
    withinClusterQueue: LowerPriority
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  namespace: default
  name: multislice-queue
spec:
  clusterQueue: cluster-queue

  2. Crea un PriorityClass per ogni livello di priorità distinto che vuoi assegnare ai workload:

    apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
  name: low-priority
value: 100
globalDefault: false
description: "This low priority class should be used for some Pods only."

  3. Assegna priorityClassName al tuo JobSet:

    Autopilot 

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: low-priority
  labels:
    kueue.x-k8s.io/queue-name: multislice-queue
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: 1
      template:
        spec:
          parallelism: 2
          completions: 2
          backoffLimit: 0
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 2x4
              priorityClassName: low-priority
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                command:
                - bash
                - -c
                - |
                  sleep 60
                resources:
                  limits:
                    google.com/tpu: 4 # Number of TPU chips per worker

    Standard 

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: low-priority
  labels:
    kueue.x-k8s.io/queue-name: multislice-queue
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: 1
      template:
        spec:
          parallelism: 2
          completions: 2
          backoffLimit: 0
          template:
            spec:
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 2x4
              priorityClassName: low-priority
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                securityContext:
                  privileged: true
                command:
                - bash
                - -c
                - |
                  sleep 60
                resources:
                  limits:
                    google.com/tpu: 4 # Number of TPU chips per worker

GKE include un criterio di prelazione che definisce in che modo Kueue assegna le risorse disponibili. Il criterio specifica che un carico di lavoro può essere prelevato se un carico di lavoro con priorità più elevata ha bisogno delle risorse. I carichi di lavoro con un valore di priorità inferiore hanno maggiori probabilità di essere prelevati da carichi di lavoro con priorità più alta.

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo tutorial, elimina il progetto che contiene le risorse oppure mantieni il progetto ed elimina le singole risorse.

Elimina il progetto

  1. In the Google Cloud console, go to the Manage resources page.

    Go to Manage resources

  2. In the project list, select the project that you want to delete, and then click Delete.
  3. In the dialog, type the project ID, and then click Shut down to delete the project.

Elimina la singola risorsa

  1. Elimina il sistema di quote di Kueue:

    kubectl delete -n team-a localqueue
kubectl delete -n team-b localqueue
kubectl delete clusterqueue
kubectl delete clusterqueue
kubectl delete clusterqueue
kubectl delete resourceflavor
kubectl delete resourceflavor
kubectl delete resourceflavor

  2. Elimina il file manifest di Kueue:

    VERSION=kueue.x-k8s.io/v1beta1
kubectl delete -f \
    https://github.com/kubernetes-sigs/kueue/releases/download/$VERSION/manifests.yaml

  3. Elimina il cluster:

    gcloud container clusters delete kueue-cohort --region=COMPUTE_REGION

Passaggi successivi