Questo tutorial mostra come orchestrare carichi di lavoro su più sezioni su Google Kubernetes Engine (GKE). Esegui un Carico di lavoro Jax che utilizza 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.
Questo tutorial mostra come orchestrare più sezioni carichi di lavoro l'esecuzione simultanea di risorse TPU.
Prima di utilizzare le TPU in GKE, ti consigliamo di completare nel seguente percorso di apprendimento:
- Scopri di più sulla disponibilità attuale della versione di TPU con il Architettura di sistema Cloud TPU.
- 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:
- 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. - Crea le risorse Kueue per garantire che le quote vengano condivise in modo equo tra i carichi di lavoro.
- Esegui il carico di lavoro Multisezione.
Prima di iniziare
Prima di iniziare, assicurati di aver eseguito le seguenti attività:
- Attiva l'API Google Kubernetes Engine. Abilita l'API Google Kubernetes Engine
- Se vuoi utilizzare Google Cloud CLI per questa attività,
install e poi
inizializzare
con gcloud CLI. Se hai già installato gcloud CLI, scarica la versione più recente
eseguendo
gcloud components update
.
Installa JobSet v0.2.3 o versioni successive.
Installa Kueue v0.4.1 o versioni successive.
prepara l'ambiente
Nella console Google Cloud, avvia un'istanza di Cloud Shell:
Apri Cloud ShellImposta 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.
Cluster Autopilot che eseguono la versione 1.29.2-gke.1521000 abilitare le TPU per impostazione predefinita. TPU sui cluster Autopilot sono configurate nella specifica del carico di lavoro. Per ulteriori informazioni, consulta la sezione Definisci i carichi di lavoro multisezione 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 posizione in cui vuoi
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. Autopilot i cluster 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
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.
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
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
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
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 TPU
sezioni con topologia
2x4
. Ogni sezione TPU ha una topologia2x4
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 Kuueue è un 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 set di job eseguono un comando Jax carico di lavoro che restituisce il numero globale di chip TPU nella sezione, poi dorme per 60 secondi per simulare il tempo di addestramento del modello, poi esce.
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
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 totali. - 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 corrono
insieme in seguito.
Verifica che Kueue abbia ammesso i carichi di lavoro
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 che ti servono.
Monitora i carichi di lavoro
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 Set di jobmultislice-1slice
emultislice-2slice
.
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.
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
Crea un
PriorityClass
per ogni livello di priorità 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."
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
- Nella console Google Cloud, vai alla pagina Gestisci risorse.
- Nell'elenco dei progetti, seleziona il progetto che vuoi eliminare, quindi fai clic su Elimina.
- Nella finestra di dialogo, digita l'ID del progetto e fai clic su Chiudi per eliminare il progetto.
Elimina la singola risorsa
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
Elimina il manifest Kueue:
VERSION=kueue.x-k8s.io/v1beta1 kubectl delete -f \ https://github.com/kubernetes-sigs/kueue/releases/download/$VERSION/manifests.yaml
Elimina il cluster:
gcloud container clusters delete kueue-cohort --region=COMPUTE_REGION
Passaggi successivi
- Scopri di più su Kueue.
- Scopri come Implementa un sistema di accodamento dei job con condivisione della quota tra gli spazi dei nomi su GKE.