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

Autopilota Standard

---

Questo tutorial mostra come orchestrare più carichi di lavoro multislice su Google Kubernetes Engine (GKE). Esegui un carico di lavoro Jax utilizzando TPU Multislice, JobSet e Kueue. Kueue implementa l'accodamento dei job, decidendo quando I job devono attendere e quando devono avviarsi, in base alle quote e alla gerarchia la condivisione equa delle risorse tra i team.

Ti consigliamo di completare questo tutorial se utilizzi carichi di lavoro che richiedono Risorse TPU da eseguire contemporaneamente.

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

1. Scopri la disponibilità attuale delle versioni TPU con l'architettura di sistema Cloud TPU.
2. Scopri di più su TPU Multislice in GKE.

Obiettivi

Questo tutorial è destinato agli amministratori GKE che hanno cluster GKE esistenti e vuoi eseguire più sezioni carichi di lavoro con scale out impegnativi.

Questo tutorial illustra i seguenti passaggi:

1. Prepara l'ambiente con un cluster GKE con tre slice TPU v5e. Ogni sezione TPU ha una topologia 2x4 con 8 chip. Pertanto, 24 chip TPU v5e TPU in totale.
2. Crea le risorse Kueue per assicurarti che le quote vengano 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 operazioni:

• Attiva l'API Google Kubernetes Engine.
Abilita l'API Google Kubernetes Engine
• Se vuoi utilizzare Google Cloud CLI per questa attività, installa e poi inizializza gcloud CLI. Se hai già installato gcloud CLI, scarica la versione più recente eseguendo gcloud components update.
  Nota: per gcloud CLI esistente di installazione, assicurati di impostare compute/region e compute/zone proprietà. Se imposti località predefinite, puoi evitare errori in gcloud CLI come i seguenti: One of [--zone, --region] must be supplied: Please specify location.

• Installa JobSet v0.2.3 o versioni successive.

• Installa Kueue versione 0.4.1 o 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 attivano le TPU per impostazione predefinita. TPU sui cluster Autopilot sono configurate nella specifica del carico di lavoro. 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

Standard

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

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

Se utilizzi dalla modalità GKE Autopilot, passa Crea la sezione 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

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 nel cluster regione 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 una ClusterQueue da cui le risorse allocato per eseguire i suoi 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 carichi di lavoro multislice con JobSet

In questa sezione, crei tre JobSet. Questi 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 e poi esce.

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 JobSet multislice-1slice crea un job che richiede un totale di una sezione TPU.
• Il JobSet multislice-2slice crea due job che richiedono due sezioni TPU in totale.
• Il JobSet multislice-3slice crea tre job che richiedono tre sezioni TPU in totale.

Poiché il cluster ha solo tre sezioni 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 attendono e corrono insieme in seguito.

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 accoda uno o più carichi di lavoro, a seconda delle risorse TPU che ti servono.

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
   

   Scopri che GKE ha pianificato, creato ed eseguito i pod multislice-3slice prima. Poi, GKE ha eseguito i pod multislice-1slice e multislice-2slice set di job.

Risultato: hai eseguito correttamente tre carichi di lavoro Jax Multislice.

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

Facoltativamente, puoi assegnare priorità ai 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 prelazione di ClusterQueue, potrebbero verificarsi alcuni scenari di prelazione. Ad esempio, GKE può prerilasciare un carico di lavoro in esecuzione per eseguire un carico di lavoro con priorità più elevata se non sono presenti e risorse sufficienti per eseguire entrambi i carichi di lavoro contemporaneamente. GKE interrompe l'elaborazione e lo inserisce nuovamente in un carico di lavoro prerilasciato per continuare a elaborare possibile.

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 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à descritte in questo documento, quando lo elimini elimini anche tutto il 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 appspot.com URL, elimina le risorse selezionate all'interno del progetto anziché eliminare l'intero 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 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

• Scopri di più su Kueue.
• Scopri come implementare un sistema di coda dei job con condivisione delle quote tra gli spazi dei nomi su GKE.