In dieser Anleitung erfahren Sie, wie Sie mehrere Multislice-Arbeitslasten in der Google Kubernetes Engine (GKE) orchestrieren. Sie führen eine Jax-Arbeitslast mit TPU-Multislice, JobSet und Kueue aus. Kueue implementiert die Jobwarteschlangen und entscheidet anhand der Kontingente und einer Hierarchie für die faire Freigabe von Ressourcen, wann Jobs warten und wann sie gestartet werden sollen.
Wir empfehlen Ihnen, diese Anleitung durchzugehen, wenn Sie Arbeitslasten verwenden, für die TPU-Ressourcen gleichzeitig ausgeführt werden müssen.
Bevor Sie TPUs in GKE verwenden, sollten Sie den folgenden Lernpfad durcharbeiten:
- Lernen Sie mehr über die aktuelle Verfügbarkeit von TPU-Versionen unter Cloud TPU-Systemarchitektur.
- TPU-Multislice in GKE
Ziele
Diese Anleitung richtet sich an GKE-Administratoren, die bereits GKE-Cluster haben und zum ersten Mal Multislice-Arbeitslasten ausführen möchten.
Diese Anleitung umfasst die folgenden Schritte:
- Bereiten Sie Ihre Umgebung mit einem GKE-Cluster mit drei v5e-TPU-Slices vor. Jedes TPU-Slice hat eine
2x4
-Topologie mit 8 Chips. Daher insgesamt 24 TPU v5e TPU-Chips. - Erstellen Sie die Kueue-Ressourcen, um für eine faire Verteilung der Kontingente zwischen den Arbeitslasten zu sorgen.
- Führen Sie die Multislice-Arbeitslast aus.
Vorbereitung
Führen Sie die folgenden Schritte durch, bevor Sie beginnen:
- Aktivieren Sie die Google Kubernetes Engine API. Google Kubernetes Engine API aktivieren
- Wenn Sie die Google Cloud CLI für diese Aufgabe verwenden möchten, müssen Sie die gcloud CLI installieren und dann initialisieren. Wenn Sie die gcloud CLI bereits installiert haben, rufen Sie die neueste Version mit
gcloud components update
ab.
Installieren Sie JobSet v0.2.3 oder höher.
Installieren Sie Kueue v0.4.1 oder höher.
Umgebung vorbereiten
Starten Sie in der Google Cloud Console eine Cloud Shell-Instanz:
Cloud Shell öffnenLegen Sie die Standardumgebungsvariablen fest:
gcloud config set project PROJECT_ID gcloud config set compute/region COMPUTE_REGION
Ersetzen Sie die folgenden Werte:
- PROJECT_ID: Ihre Google Cloud-Projekt-ID.
- COMPUTE_REGION: Die Compute Engine-Region.
Autopilot-Cluster, auf denen Version 1.29.2-gke.1521000 oder höher ausgeführt wird, aktivieren standardmäßig TPUs. TPUs in Autopilot-Clustern werden in der Arbeitslastspezifikation konfiguriert. Weitere Informationen finden Sie im Abschnitt Multislice-Arbeitslasten mit JobSets definieren.
GKE-Cluster erstellen
Erstellen Sie in Cloud Shell einen GKE-Cluster:
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
Ersetzen Sie LOCATION durch den Standort, an dem Sie den Cluster erstellen möchten. Achten Sie darauf, dass es Kapazität für den Maschinentyp ct5lp-hightpu-4t
hat.
Die Erstellung eines Clusters kann einige Minuten dauern.
Wenn Sie den GKE-Autopilot-Modus verwenden, fahren Sie mit dem Abschnitt Kuee-Ressourcen erstellen fort. Autopilot-Cluster, die Version 1.29.2-gke.1521000 oder höher ausführen, aktivieren standardmäßig TPUs.
Drei TPU-Slice-Knotenpools im Standardmodus erstellen
Erstellen Sie den ersten Knotenpool mit dem Namen
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
Ersetzen Sie NODE_LOCATION durch eine oder mehrere Zonen in der Clusterregion, in der Sie die Knoten erstellen möchten.
Erstellen Sie den zweiten Knotenpool mit dem Namen
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
Erstellen Sie den dritten Knotenpool mit dem Namen
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 erstellt drei Knotenpools. Jeder Knotenpool ist ein separates TPU-Slice.
Kueue-Ressourcen erstellen
Erstellen Sie das folgende
kueue.yaml
-Manifest: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
Wenden Sie das
kueue.yaml
-Manifest an:kubectl apply -f kueue.yaml
GKE erstellt die folgenden Kueue-Ressourcen:
- ResourceFlavor: Eine Abstraktion der Ressourcen in einem Cluster. In diesem Beispiel erstellt GKE drei TPU-Slices mit
2x4
-Topologie. Jedes TPU-Slice hat eine2x4
-Topologie mit 8 Chips (insgesamt 24 TPU-Chips). - ClusterQueue: Eine globale Warteschlange, die Arbeitslasten und Clusterressourcen verwaltet.
- LocalQueue: Gruppiert eng verwandte Arbeitslasten, die in der Regel von einem einzelnen Mandanten (Nutzer) ausgeführt werden. Jede LocalQueue verweist auf eine ClusterQueue, aus der Ressourcen zum Ausführen ihrer Arbeitslasten zugewiesen werden. Eine Kueue-Arbeitslast ist eine Abstraktion, die eine Batcharbeitslast darstellt. In diesem Fall ist jede Arbeitslast ein JobSet.
Multislice-Arbeitslasten mit JobSets definieren
In diesem Abschnitt erstellen Sie drei JobSets. Diese JobSets führen eine Jax-Arbeitslast aus, die die globale Anzahl von TPU-Chips im Slice ausgibt, 60 Sekunden lang pausiert, um eine gewisse Modelltrainingszeit zu simulieren, und dann beendet wird.
Erstellen Sie das folgende
jobsets-multislice.yaml
-Manifest: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
Wenden Sie das
jobsets-multislice.yaml
-Manifest an:kubectl apply -f jobsets-multislice.yaml
GKE erstellt die Jobs mit den folgenden Ressourcenanfragen:
- Der JobSet
multislice-1slice
erstellt einen Job, für den insgesamt ein TPU-Slice erforderlich ist. - Das JobSet
multislice-2slice
erstellt zwei Jobs, für die insgesamt zwei TPU-Slices erforderlich sind. - Das JobSet
multislice-3slice
erstellt drei Jobs, für die insgesamt drei TPU-Slices erforderlich sind.
Da der Cluster nur drei TPU-Slices hat, können nicht alle JobSets gleichzeitig ausgeführt werden.
Wenn Kueue alle drei multislice-3slice
-JobSets in die Warteschlange einreiht, werden seine Jobs bis zum Abschluss allein ausgeführt. Das multislice-1slice
und multislice-2slice
warten und werden gemeinsam ausgeführt.
Prüfen, ob Kueue die Arbeitslasten zugelassen hat
Prüfen Sie die Arbeitslasten in der Warteschlange in Kueue:
kubectl get workloads
Die Ausgabe sieht in etwa so aus:
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 stellt eine oder mehrere Arbeitslasten je nach den erforderlichen TPU-Ressourcen in die Warteschlange.
Arbeitslasten überwachen
Überwachen Sie, welche Pods ausgeführt werden:
kubectl get pods
Die Ausgabe sieht in etwa so aus:
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
Beachten Sie, dass GKE zuerst die Pods für
multislice-3slice
geplant, erstellt und ausgeführt hat. Anschließend hat GKE die Pods aus den JobSetsmultislice-1slice
undmultislice-2slice
ausgeführt.
Kueue-Arbeitslastprioritäten und vorzeitiges Beenden aktivieren
Optional können Sie Kueue-Arbeitslastprioritäten zuweisen, die die Reihenfolge bestimmen, in der Arbeitslasten in der Warteschlange von Kueue zugelassen werden.
Aktualisieren Sie Ihre
ClusterQueue
mit einer Richtlinie zum vorzeitigen Beendigen: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
Erstellen Sie eine
PriorityClass
für jede Prioritätsstufe, die Sie Arbeitslasten zuweisen möchten: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."
Weisen Sie Ihrem JobSet
priorityClassName
zu: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 ```
Bereinigen
Damit Ihrem Google Cloud-Konto die in dieser Anleitung verwendeten Ressourcen nicht in Rechnung gestellt werden, löschen Sie entweder das Projekt, das die Ressourcen enthält, oder Sie behalten das Projekt und löschen die einzelnen Ressourcen.
Projekt löschen
- In the Google Cloud console, go to the Manage resources page.
- In the project list, select the project that you want to delete, and then click Delete.
- In the dialog, type the project ID, and then click Shut down to delete the project.
Einzelne Ressource löschen
Das Kueue-Kontingentsystem löschen:
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
Das Kueue-Manifest löschen:
VERSION=kueue.x-k8s.io/v1beta1 kubectl delete -f \ https://github.com/kubernetes-sigs/kueue/releases/download/$VERSION/manifests.yaml
Löschen Sie den Cluster:
gcloud container clusters delete kueue-cohort --region=COMPUTE_REGION