Orchestrare carichi di lavoro Multislice utilizzando JobSet e Kueue

Autopilot Standard

---

Questo tutorial mostra come eseguire un carico di lavoro Jax utilizzando più sezioni TPU in Google Kubernetes Engine (GKE) e Kueue. Kueue implementa l'accodamento dei job, decidendo quando i job devono attendere e quando devono iniziare, in base alle quote e a una gerarchia per la condivisione equa delle risorse tra i team.

Questo tutorial mostra come orchestrare più carichi di lavoro multisettore che richiedono l'esecuzione simultanea di risorse TPU.

Prima di utilizzare le TPU in GKE, ti consigliamo di completare il seguente percorso di apprendimento:

1. Scopri di più sulla disponibilità attuale della versione di TPU con l'architettura di sistema Cloud TPU.
2. Scopri di più su TPU Multislice in GKE.

Obiettivi

Questo tutorial è destinato agli amministratori di GKE che hanno cluster GKE esistenti e vogliono eseguire carichi di lavoro multislice per la prima volta.

Questo tutorial illustra i seguenti passaggi:

1. Prepara l'ambiente con un cluster GKE con tre sezioni TPU v5e. Ogni sezione TPU ha una topologia 2x4 con 8 chip. Pertanto, 24 chip TPU v5e TPU in totale.
2. Creare le risorse Kueue per garantire che le quote siano condivise equamente tra i carichi di lavoro.
3. Esegui il carico di lavoro Multisezione.

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à, installa e quindi initialize gcloud CLI. Se hai già installato gcloud CLI, ottieni la versione più recente eseguendo gcloud components update.
  Nota: per le installazioni esistenti di gcloud CLI, assicurati di impostare le proprietà compute/region e compute/zone. Se imposti località predefinite, puoi evitare errori in gcloud CLI, come quelli che seguono: One of [--zone, --region] must be supplied: Please specify location.

• Installa JobSet v0.2.3 o versioni successive.

• Installa Kueue v0.4.1 o versioni successive.

prepara l'ambiente

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

2. Imposta le variabili di ambiente predefinite:

   gcloud config set project PROJECT_ID
   gcloud config set compute/region COMPUTE_REGION
   

   Sostituisci i seguenti valori:

   • PROJECT_ID: il tuo ID progetto Google Cloud.
   • COMPUTE_REGION: la regione di Compute Engine.

I cluster Autopilot che eseguono la versione 1.29.2-gke.1521000 o successive abilitano le TPU per impostazione predefinita. Le TPU sui cluster Autopilot sono configurate nella specifica dei carichi di lavoro. Per ulteriori informazioni, consulta la sezione Definire i carichi di lavoro multisettori con i set di job.

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

Standard

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

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

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

Crea tre pool di nodi della sezione TPU in modalità Standard

1. Crea il primo pool di nodi 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 \
       --num-nodes=2 \
       --project=PROJECT_ID
   

   Sostituisci NODE_LOCATION con una o più zone dell'area geografica del cluster in cui vuoi creare i nodi.

2. Crea il secondo pool di nodi 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 \
       --num-nodes=2 \
       --project=PROJECT_ID
   

3. Crea il terzo pool di nodi 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 \
       --num-nodes=2 \
       --project=PROJECT_ID
   

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

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 dei cluster.
• LocalQueue: raggruppa carichi di lavoro strettamente correlati, generalmente eseguiti da un singolo tenant (utente). Ogni LocalQueue punta a un ClusterQueue da cui vengono allocate le risorse per eseguire i propri carichi di lavoro. Un carico di lavoro Kuueue è un'astrazione che rappresenta un carico di lavoro batch. In questo caso, ogni carico di lavoro è un JobSet.

Definisci i carichi di lavoro multisettore con i jobset

In questa sezione, creerai tre set di job. Questi oggetti JobSet eseguono un carico di lavoro Jax che restituisce il numero globale di chip TPU nella sezione, quindi viene sospeso per 60 secondi per simulare un tempo di addestramento del modello, quindi si chiude.

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

2. Applica il manifest jobsets-multislice.yaml:

   kubectl apply -f jobsets-multislice.yaml
   

GKE crea i job con le seguenti richieste di risorse:

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

Poiché il cluster ha solo tre sezioni TPU, non tutti i JobSet possono essere eseguiti contemporaneamente. Quando Kueue accoda tutti e tre i jobSet multislice-3slice, i relativi job vengono eseguiti da soli fino al completamento. multislice-1slice e multislice-2slice attendono e poi vengono eseguite insieme.

Verifica che Kueue abbia ammesso i carichi di lavoro

1. Controlla i carichi di lavoro accodati 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 accoda uno o più carichi di lavoro, a seconda delle risorse TPU richieste.

Monitora i carichi di lavoro

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
   

   Per prima cosa, GKE ha pianificato, creato ed eseguito i pod per multislice-3slice. Poi, GKE ha eseguito i pod da multislice-1slice e multislice-2slice JobSet.

Operazione riuscita: hai eseguito tre carichi di lavoro Jax Multislice.

Abilita le priorità e il prerilascio dei carichi di lavoro Kueue

Facoltativamente, puoi assegnare le priorità dei carichi di lavoro Kueue che determinano l'ordine in cui Kueue ammette i carichi di lavoro in coda.

Nota: a seconda della configurazione del criterio di prerilascio di ClusterQueue, potrebbero verificarsi alcuni scenari di prerilascio. Ad esempio, GKE può prerilasciare un carico di lavoro in esecuzione per eseguire un carico di lavoro con priorità più elevata se non ci sono risorse sufficienti per eseguire entrambi i carichi di lavoro contemporaneamente. GKE arresta e ripete la coda di un carico di lavoro prerilasciato per proseguire con l'elaborazione quando possibile.

1. Aggiorna ClusterQueue per avere un criterio di prerilascio:

   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 carichi di lavoro:

   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 set di job:

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

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

Attenzione: l'eliminazione di un progetto ha i seguenti effetti:
• L'intero contenuto del progetto viene eliminato. Se hai utilizzato un progetto esistente per le attività in questo documento, quando lo elimini elimini anche qualsiasi altro lavoro svolto nel progetto.
• Gli ID progetto personalizzati non sono più disponibili. Quando hai creato questo progetto, potresti aver creato un ID progetto personalizzato che vuoi utilizzare in futuro. Per conservare gli URL che utilizzano l'ID progetto, ad esempio un URL appspot.com, elimina le risorse selezionate all'interno del progetto anziché eliminare l'intero progetto.
1. Nella console Google Cloud, vai alla pagina Gestisci risorse.

   Vai a Gestisci risorse

2. Nell'elenco dei progetti, seleziona il progetto che vuoi eliminare, quindi fai clic su Elimina.
3. Nella finestra di dialogo, digita l'ID del progetto e fai clic su Chiudi per eliminare il progetto.

Elimina la singola risorsa

1. Elimina il sistema di quote 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 manifest 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

• Scopri di più su Kueue
• Scopri come implementare un sistema di accodamento dei job con la condivisione della quota tra gli spazi dei nomi su GKE.