Einführung in TPUs in GKE

Kunden der Google Kubernetes Engine (GKE) können jetzt Kubernetes-Knotenpools erstellen, die TPU v4- und v5e-Pods enthalten. Ein TPU-Pod ist eine Gruppe von TPU-Geräten, die über Hochgeschwindigkeits-Interconnect-Verbindungen verbunden sind. Dies unterscheidet sich von einem Kubernetes-Pod, der kleinsten bereitstellbaren Recheneinheit, die Sie in Kubernetes erstellen und verwalten können. Für Arbeitslasten, für die kein vollständiger TPU-Pod erforderlich ist, können Sie eine Teilmenge eines vollständigen TPU-Pods verwenden, die als TPU-Slice bezeichnet wird. Wie vollständige TPU-Pods hat jedes TPU-Gerät in einem Slice eine eigene TPU-VM. Eine TPU-VM und das verbundene Gerät werden als Host oder TPU-Knoten bezeichnet. Weitere Informationen zu TPU-Pods finden Sie unter Systemarchitektur.

Da der Begriff Pod im Kontext von GKE in der Regel einen Kubernetes-Pod bezeichnet, bezeichnen wir eine Sammlung von einem oder mehreren TPU-Geräten immer als Slice, um Verwirrung zu vermeiden.

Wenn Sie mit GKE arbeiten, müssen Sie zuerst einen GKE-Cluster erstellen.

Anschließend fügen Sie dem Cluster Knotenpools hinzu. GKE-Knotenpools sind Sammlungen von VMs, die dieselben Attribute haben. Bei TPU-Arbeitslasten bestehen Knotenpools aus TPU-VMs.

Knotenpooltypen

GKE unterstützt zwei Arten von TPU-Knotenpools:

TPU-Slice-Knotenpool mit mehreren Hosts

Ein TPU-Slice-Knotenpool mit mehreren Hosts ist ein Knotenpool, der zwei oder mehr verbundene TPU-VMs enthält. Mit jeder VM ist ein TPU-Gerät verbunden. Die TPUs in einem Multi-Host-Slice sind über eine High-Speed Interconnect-Verbindung (ICI) verbunden. Ein TPU-Slice-Knotenpool mit mehreren Hosts ist unveränderlich. Nachdem ein Knotenpool mit mehreren Hostsegmenten erstellt wurde, können Sie ihm keine Knoten mehr hinzufügen. Sie können beispielsweise keinen v4-32-Knotenpool erstellen und ihm später einen zusätzlichen Kubernetes-Knoten (TPU-VM) hinzufügen. Wenn Sie einem GKE-Cluster ein zusätzliches TPU-Slice hinzufügen möchten, müssen Sie einen neuen Knotenpool erstellen.

Die Hosts in einem TPU-Slice-Knotenpool mit mehreren Hosts werden als einzelne atomare Einheit behandelt. Wenn GKE keinen Knoten im Slice bereitstellen kann, schlägt die Bereitstellung aller Knoten im Slice fehl.

Wenn ein Knoten in einem TPU-Slice mit mehreren Hosts repariert werden muss, fährt GKE alle TPU-VMs (Knoten) im Slice herunter, wodurch alle Kubernetes-Pods in der Arbeitslast entfernt werden. Sobald alle TPU-VMs im Slice ausgeführt werden, können die Kubernetes-Pods auf den TPU-VMs im neuen Slice geplant werden.

Das folgende Diagramm zeigt ein Beispiel für einen TPU-Slice des v5litepod-16 (v5e) mit mehreren Hosts. Dieses Slice hat vier TPU-VMs. Jede TPU-VM hat vier TPU v5e-Chips, die mit Highspeed Interconnect-Verbindungen (ICI) verbunden sind, und jeder TPU v5e-Chip hat einen TensorCore.

TPU-Slice-Diagramm mit mehreren Hosts

Das folgende Diagramm zeigt einen GKE-Cluster, der ein TPU-Slice v5litepod-16 (v5e) (Topologie: 4x4) und ein TPU-Segment vom Typ v5litepod-8 (v5e) (Topologie: 2x4) enthält:

TPU v5e Pod-Diagramm

Ein Beispiel für das Ausführen einer Arbeitslast in einem TPU-Slice mit mehreren Hosts finden Sie unter Arbeitslast in einem TPU-Slice mit mehreren Hosts ausführen.

TPU-Slice-Knotenpools mit einzelnem Host

Ein Knotenpool mit einem einzelnen Hostsegment ist ein Knotenpool, der eine oder mehrere unabhängige TPU-VMs enthält. Mit jeder dieser VMs ist ein TPU-Gerät verbunden. Während die VMs innerhalb eines Knotenpools mit einem einzelnen Hostsegment über das Rechenzentrumsnetzwerk (DCN) kommunizieren können, sind die an die VMs angehängten TPUs nicht miteinander verbunden. Das folgende Diagramm zeigt ein Beispiel für ein TPU-Slice mit einem einzelnen Host und sieben v4-8-Maschinen:

Diagramm eines Knotenpools mit einem einzelnen Hostsegment

Ein Beispiel für das Ausführen einer Arbeitslast in einem TPU-Slice mit einem einzelnen Host finden Sie unter Arbeitslasten auf TPU-Knoten ausführen.

TPU-Maschinentypen für GKE-Knotenpools

Bevor Sie Knotenpools erstellen, müssen Sie die TPU-Version und die Größe des TPU-Slice auswählen, den Ihre Arbeitslast benötigt. TPU v4 wird in der GKE-Standardversion 1.26.1-gke.1500 und höher, v5e in der GKE-Standardversion 1.27.2-gke.2100 und höher sowie v5p in der GKE-Standardversion 1.28.3-gke.1024000 und höher unterstützt.

TPU v4, v5e und v5p werden ab GKE Autopilot-Version 1.29.2-gke.1521000 unterstützt.

Weitere Informationen zu den Hardwarespezifikationen der verschiedenen TPU-Versionen finden Sie unter Systemarchitektur. Wählen Sie beim Erstellen eines TPU-Knotenpools eine TPU-Slice-Größe (eine TPU-Topologie) aus, die auf der Größe Ihres Modells und dem erforderlichen Arbeitsspeicher basiert. Der Maschinentyp, den Sie beim Erstellen der Knotenpools angeben, hängt von der Version und Größe der Segmente ab.

V5e

Im Folgenden finden Sie die TPU v5e-Maschinentypen und -Topologien, die für Trainings- und Inferenz-Anwendungsfälle unterstützt werden:

Maschinentyp Topologie Anzahl der TPU-Chips Anzahl der VMs Empfohlener Anwendungsfall
ct5lp-hightpu-1t 1x1 1 1 Training, Inferenz auf einen einzelnen Host
ct5lp-hightpu-4t 2x2 4 1 Training, Inferenz auf einen einzelnen Host
ct5lp-hightpu-8t 2x4 8 1 Training, Inferenz auf einen einzelnen Host
ct5lp-hightpu-4t 2x4 8 2 Training, Inferenz mehrerer Hosts
ct5lp-hightpu-4t 4x4 16 4 Umfangreiches Training, Inferenz mit mehreren Hosts
ct5lp-hightpu-4t 4x8 32 8 Umfangreiches Training, Inferenz mit mehreren Hosts
ct5lp-hightpu-4t 8x8 64 16 Umfangreiches Training, Inferenz mit mehreren Hosts
ct5lp-hightpu-4t 8x16 128 32 Umfangreiches Training, Inferenz mit mehreren Hosts
ct5lp-hightpu-4t 16x16 256 64 Umfangreiches Training, Inferenz mit mehreren Hosts

Cloud TPU v5e ist ein kombiniertes Trainings- und Inferenzprodukt. Trainingsjobs sind für Durchsatz und Verfügbarkeit optimiert, während Inferenzjobs für Latenz optimiert sind. Weitere Informationen finden Sie unter Beschleunigertypen für v5e Training und Beschleunigertypen für Inferenzen von v5e.

TPU v5e-Maschinen sind in us-west4-a, us-east5-b und us-east1-c verfügbar. GKE-Standardcluster müssen Version 1.27.2-gke.2100 der Steuerungsebene oder höher ausführen. GKE Autopilot muss Version 1.29.2-gke.1521000 der Steuerungsebene oder höher ausführen. Weitere Informationen zu v5e finden Sie unter Cloud TPU v5e-Training.

Vergleich der Maschinentypen:

Maschinentyp ct5lp-hightpu-1t ct5lp-hightpu-4t ct5lp-hightpu-8t
Anzahl der v5e-Chips 1 4 8
Anzahl der vCPUs 24 112 224
RAM (GB) 48 192 384
Anzahl der NUMA-Knoten 1 1 2
Wahrscheinlichkeit eines vorzeitigen Beendens Hoch Mittel Niedrig

Um Platz für VMs mit mehr Chips zu schaffen, kann der GKE-Planer VMs mit weniger Chips vorzeitig planen und neu planen. 8-Chip-VMs führen also wahrscheinlich dazu, dass 1- und 4-Chip-VMs präemptiv beendet werden.

v4 und v5p

Dies sind die TPU v4- und v5p-Maschinentypen:

Maschinentyp Anzahl der vCPUs Arbeitsspeicher (GB) Anzahl der NUMA-Knoten
ct4p-hightpu-4t 240 407 2
ct5p-hightpu-4t 208 448 2

Verwenden Sie beim Erstellen eines TPU v4-Slice den Maschinentyp ct4p-hightpu-4t, der einen Host und vier Chips enthält. Weitere Informationen finden Sie unter v4-Topologien und TPU-Systemarchitektur. TPU v4 Pod-Maschinentypen sind in us-central2-b verfügbar. Ihre GKE-Standardcluster müssen Version 1.26.1-gke.1500 oder höher der Steuerungsebene ausführen. GKE Autopilot-Cluster müssen Version 1.29.2-gke.1521000 oder höher der Steuerungsebene ausführen.

Verwenden Sie beim Erstellen eines TPU v5p-Slice den Maschinentyp ct5p-hightpu-4t. Dieser hat einen Host und vier Chips. TPU v5p Pod-Maschinentypen sind in us-west4-a und us-east5-a verfügbar. GKE-Standardcluster müssen Version 1.28.3-gke.1024000 der Steuerungsebene oder höher ausführen. GKE Autopilot muss 1.29.2-gke.1521000 oder höher ausführen. Weitere Informationen zu v5p finden Sie in der Einführung zum v5p-Training.

Bekannte Probleme und Beschränkungen

  • Maximale Anzahl von Kubernetes-Pods: Sie können maximal 256 Kubernetes-Pods in einer einzelnen TPU-VM ausführen.
  • Nur Spezifische Reservierungen: Wenn Sie TPUs in GKE verwenden, ist SPECIFIC der einzige unterstützte Wert für das Flag --reservation-affinity des Befehls gcloud container node-pools create.
  • Nur die Spot-VMs-Variante von TPUs auf Abruf wird unterstützt: Spot-VMs ähneln VMs auf Abruf und unterliegen denselben Verfügbarkeitsbeschränkungen, haben jedoch keine maximale Dauer von 24 Stunden.
  • Keine Unterstützung bei der Kostenzuweisung: Die GKE-Kostenzuweisung und die Nutzungsmessung enthalten keine Daten zur Nutzung oder Kosten von TPUs.
  • Autoscaling berechnet möglicherweise die Kapazität: Cluster-Autoscaling berechnet die Kapazität neuer TPU-Knoten möglicherweise falsch, bevor diese Knoten verfügbar sind. Cluster-Autoscaling führt dann möglicherweise eine zusätzliche Hochskalierung durch und erstellt dadurch mehr Knoten als erforderlich. Cluster-Autoscaling skaliert nach dem regulären Herunterskalieren weitere Knoten herunter, wenn sie nicht benötigt werden.
  • Autoscaling bricht das Hochskalieren ab: Cluster-Autoscaling bricht das Hochskalieren von TPU-Knotenpools ab, die sich länger als 10 Stunden im Wartestatus befinden. Cluster Autoscaler wiederholt solche Hochskalierungsvorgänge später. Dieses Verhalten kann die TPU-Verfügbarkeit für Kunden verringern, die keine Reservierungen verwenden.
  • Herunterskalierung kann das Herunterskalieren verhindern: Nicht-TPU-Arbeitslasten, die eine Toleranz für die TPU-Markierung haben, können das Herunterskalieren des Knotenpools verhindern, wenn sie während des Leerens des TPU-Knotenpools neu erstellt werden.

Ausreichende TPU- und GKE-Kontingente sicherstellen

Möglicherweise müssen Sie bestimmte GKE-bezogene Kontingente in den Regionen erhöhen, in denen Ihre Ressourcen erstellt werden.

Die folgenden Kontingente haben Standardwerte, die wahrscheinlich erhöht werden müssen:

  • SSD-Kontingent für nichtflüchtigen Speicher (GB): Das Bootlaufwerk jedes Kubernetes-Knotens benötigt standardmäßig 100 GB. Daher sollte dieses Kontingent auf mindestens (die maximale Anzahl von GKE-Knoten, die Sie voraussichtlich erstellen) × 100 GB betragen.
  • Kontingent für verwendete IP-Adressen: Jeder Kubernetes-Knoten verwendet genau eine IP-Adresse. Daher sollte dieses Kontingent mindestens so hoch sein wie die maximale Anzahl von GKE-Knoten, die Sie erstellen möchten.

Informationen zum Anfordern einer Erhöhung des Kontingents finden Sie unter Höheres Kontingent anfordern. Weitere Informationen zu den Arten von TPU-Kontingenten finden Sie unter TPU-Kontingente.

Es kann einige Tage dauern, bis Ihre Anfragen zur Kontingenterhöhung genehmigt werden. Wenn Ihre Anfragen zur Kontingenterhöhung innerhalb weniger Tage nicht genehmigt werden, wenden Sie sich an Ihr Google Konten-Team.

TPU-Reservierung migrieren

Wenn Sie keine vorhandene TPU-Reservierung mit TPUs in GKE verwenden möchten, überspringen Sie diesen Abschnitt und fahren Sie mit Google Kubernetes Engine-Cluster erstellen fort.

Damit Sie reservierte TPUs mit GKE verwenden können, müssen Sie zuerst Ihre TPU-Reservierung zu einem neuen Compute Engine-basierten Reservierungssystem migrieren.

Beachten Sie bei dieser Migration einige wichtige Informationen:

  • TPU-Kapazität, die zum neuen Compute Engine-basierten Reservierungssystem migriert wurde, kann nicht mit der Queued Resource API von Cloud TPU verwendet werden. Wenn Sie TPU-Ressourcen in der Warteschlange mit Ihrer Reservierung verwenden möchten, müssen Sie nur einen Teil Ihrer TPU-Reservierung zum neuen Compute Engine-basierten Reservierungssystem migrieren.
  • Auf den TPUs können keine Arbeitslasten aktiv ausgeführt werden, wenn diese zum neuen Compute Engine-basierten Reservierungssystem migriert werden.
  • Wählen Sie eine Zeit für die Migration aus und planen Sie die Migration in Zusammenarbeit mit Ihrem Google Cloud-Account-Management-Team. Das Zeitfenster für die Migration muss während der Geschäftszeiten (Montag bis Freitag, 9:00 bis 17:00 Uhr Pacific Time) liegen.

Google Kubernetes Engine-Cluster erstellen

Siehe Cluster erstellen in der Dokumentation zu Google Kubernetes Engine.

TPU-Knotenpool erstellen

Siehe Knotenpool erstellen in der Dokumentation zu Google Kubernetes Engine.

Ohne privilegierten Modus ausführen

Informationen zum Reduzieren des Berechtigungsbereichs für Ihren Container finden Sie unter TPU-Berechtigungsmodus.

Arbeitslasten auf TPU-Knoten ausführen

Weitere Informationen finden Sie in der Google Kubernetes Engine-Dokumentation unter Arbeitslasten auf TPU-Knoten ausführen.

Knotenselektoren

Damit Kubernetes Ihre Arbeitslast auf TPU-Knoten planen kann, müssen Sie in Ihrem Google Kubernetes Engine-Manifest für jeden TPU-Knoten zwei Selektoren angeben:

  • Legen Sie cloud.google.com/gke-accelerator-type auf tpu-v5-lite-podslice oder tpu-v4-podslice fest.
  • Legen Sie für cloud.google.com/gke-tpu-topology die TPU-Topologie des TPU-Knotens fest.

Die Abschnitte Trainingsarbeitslasten und Inferenzarbeitslasten enthalten Beispielmanifeste, die die Verwendung dieser Knotenselektoren veranschaulichen.

Überlegungen zur Arbeitslastplanung

TPUs haben besondere Merkmale, die eine spezielle Arbeitslastplanung und -verwaltung in Kubernetes erfordern. Weitere Informationen finden Sie unter Überlegungen zur Arbeitslastplanung in der GKE-Dokumentation.

TPU-Knotenreparatur

Wenn ein TPU-Knoten in einem TPU-Slice-Knotenpool mit mehreren Hosts fehlerhaft ist, wird der gesamte Knotenpool neu erstellt. Weitere Informationen finden Sie unter Automatische Knotenreparatur in der GKE-Dokumentation.

In mehreren Scheiben – über ein einzelnes Segment hinaus

Sie können kleinere Segmente in einem Mehrfachsegment aggregieren, um größere Trainingsarbeitslasten zu verarbeiten. Weitere Informationen finden Sie unter Cloud TPU-MultiSlice.

Anleitungen zum Trainieren von Arbeitslasten

Der Schwerpunkt dieser Anleitungen liegt auf Trainingsarbeitslasten auf einem TPU-Slice mit mehreren Hosts (z. B. 4 v5e-Maschinen). Sie decken die folgenden Modelle ab:

  • FLAX-Modelle umarmen: Trainiere in Pokémon
  • PyTorch/XLA: GPT2 in WikiText

Anleitungsressourcen herunterladen

Laden Sie die Python-Skripts und YAML-Spezifikationen der Anleitung für jedes vortrainierte Modell mit dem folgenden Befehl herunter:

git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

Cluster erstellen und Verbindung herstellen

Erstellen Sie einen regionalen GKE-Cluster, damit die Kubernetes-Steuerungsebene in drei Zonen repliziert wird und so eine höhere Verfügbarkeit möglich ist. Erstellen Sie den Cluster in us-west4, us-east1 oder us-central2, je nachdem, welche TPU-Version Sie verwenden. Weitere Informationen zu TPUs und Zonen finden Sie unter Cloud TPU-Regionen und -Zonen.

Mit dem folgenden Befehl wird ein neuer regionaler GKE-Cluster erstellt, der den Rapid Release Channel abonniert und einen Knotenpool hat, der anfänglich einen Knoten pro Zone enthält. Der Befehl aktiviert auch CSI-Treiberfeatures für Workload Identity und Cloud Storage FUSE in Ihrem Cluster, da die Beispielinferenzarbeitslasten in dieser Anleitung Cloud Storage-Buckets zum Speichern von vortrainierten Modellen verwenden.

gcloud container clusters create cluster-name \
  --region your-region \
  --release-channel rapid \
  --num-nodes=1 \
  --workload-pool=project-id.svc.id.goog \
  --addons GcsFuseCsiDriver

Führen Sie den folgenden Befehl aus, um CSI-Treiberfeatures für Workload Identity und Cloud Storage FUSE für vorhandene Cluster zu aktivieren:

gcloud container clusters update cluster-name \
  --region your-region \
  --update-addons GcsFuseCsiDriver=ENABLED \
  --workload-pool=project-id.svc.id.goog

Die Beispielarbeitslasten werden mit den folgenden Annahmen konfiguriert:

  • Der Knotenpool verwendet tpu-topology=4x4 mit vier Knoten
  • Der Knotenpool verwendet machine-type ct5lp-hightpu-4t

Führen Sie den folgenden Befehl aus, um eine Verbindung zum neu erstellten Cluster herzustellen:

gcloud container clusters get-credentials cluster-name \
--location=cluster-region

FLAX-Modelle umarmen: Trainiere in Pokémon

In diesem Beispiel wird das Stable Diffusion-Modell von HuggingFace mit dem Pokémon-Dataset trainiert.

Das Stable Diffusion-Modell ist ein latentes Text-zu-Bild-Modell, das aus jeder Texteingabe fotorealistische Bilder generiert. Weitere Informationen zur stabilen Diffusion finden Sie unter:

Docker-Image erstellen

Das Dockerfile befindet sich im Ordner ai-on-gke/tutorials-and-examples/tpu-examples/training/diffusion/. Führen Sie die folgenden Befehle aus, um das Docker-Image zu erstellen und per Push zu übertragen.

cd ai-on-gke/tutorials-and-examples/tpu-examples/training/diffusion/
docker build -t gcr.io/project-id/diffusion:latest .
docker push gcr.io/project-id/diffusion:latest

Arbeitslast bereitstellen

Erstellen Sie eine Datei mit folgendem Inhalt und nennen Sie sie tpu_job_diffusion.yaml. Füllen Sie das Bildfeld mit dem Bild aus, das Sie gerade erstellt haben.

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-diffusion
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-diffusion
spec:
  backoffLimit: 0
  # Completions and parallelism should be the number of chips divided by 4.
  # (e.g. 4 for a v5litepod-16)
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job-diffusion
        image: gcr.io/${project-id}/diffusion:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU usage metrics, if supported
        command:
        - bash
        - -c
        - |
          cd examples/text_to_image
          python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

Stellen Sie es dann folgendermaßen bereit:

kubectl apply -f tpu_job_diffusion.yaml

Bereinigung

Nachdem der Job ausgeführt wurde, können Sie ihn folgendermaßen löschen:

kubectl delete -f tpu_job_diffusion.yaml

PyTorch/XLA: GPT2 in WikiText

In dieser Anleitung wird gezeigt, wie GPT2 auf v5e-TPUs mit HuggingFace für PyTorch/XLA und mithilfe des Wikitext-Datasets ausgeführt wird.

Docker-Image erstellen

Das Dockerfile befindet sich im Ordner ai-on-gke/tutorials-and-examples/tpu-examples/training/gpt/. Führen Sie die folgenden Befehle aus, um das Docker-Image zu erstellen und per Push zu übertragen.

cd ai-on-gke/tutorials-and-examples/tpu-examples/training/gpt/
docker build -t gcr.io/project-id/gpt:latest .
docker push gcr.io/project-id/gpt:latest

Arbeitslast bereitstellen

Kopieren Sie die folgende YAML-Datei und speichern Sie sie in einer Datei mit dem Namen tpu_job_gpt.yaml. Füllen Sie das Bildfeld mit dem Bild aus, das Sie gerade erstellt haben.

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-gpt
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-gpt
spec:
  backoffLimit: 0
  # Completions and parallelism should be the number of chips divided by 4.
  # (for example, 4 for a v5litepod-16)
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      volumes:
      # Increase size of tmpfs /dev/shm to avoid OOM.
      - name: shm
        emptyDir:
          medium: Memory
          # consider adding `sizeLimit: XGi` depending on needs
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job-gpt
        image: gcr.io/$(project-id)/gpt:latest
        ports:
        - containerPort: 8479
        - containerPort: 8478
        - containerPort: 8477
        - containerPort: 8476
        - containerPort: 8431 # Port to export TPU usage metrics, if supported.
        env:
        - name: PJRT_DEVICE
          value: 'TPU'
        - name: XLA_USE_BF16
          value: '1'
        command:
        - bash
        - -c
        - |
          numactl --cpunodebind=0 python3 -u examples/pytorch/xla_spawn.py   --num_cores 4 examples/pytorch/language-modeling/run_clm.py    --num_train_epochs 3 --dataset_name wikitext     --dataset_config_name wikitext-2-raw-v1 --per_device_train_batch_size 16    --per_device_eval_batch_size 16 --do_train --do_eval  --output_dir /tmp/test-clm     --overwrite_output_dir --config_name my_config_2.json --cache_dir /tmp --tokenizer_name gpt2  --block_size 1024 --optim adafactor --adafactor true --save_strategy no --logging_strategy no --fsdp "full_shard" --fsdp_config fsdp_config.json
        volumeMounts:
        - mountPath: /dev/shm
          name: shm
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

Stellen Sie den Workflow bereit mit:

kubectl apply -f tpu_job_gpt.yaml

Bereinigung

Nachdem der Job ausgeführt wurde, können Sie ihn folgendermaßen löschen:

kubectl delete -f tpu_job_gpt.yaml

Anleitung: Arbeitslasten von Einzelhost-Inferenzen

In dieser Anleitung wird gezeigt, wie Sie eine Inferenzarbeitslast mit einem einzelnen Host auf GKE v5e-TPUs für vortrainierte Modelle mit JAX, TensorFlow und PyTorch ausführen. Im Allgemeinen sind im GKE-Cluster vier separate Schritte auszuführen:

  1. Erstellen Sie einen Cloud Storage-Bucket und richten Sie den Zugriff auf den Bucket ein. Sie verwenden einen Cloud Storage-Bucket, in dem das vortrainierte Modell gespeichert wird.

  2. Laden Sie ein vortrainiertes Modell herunter und konvertieren Sie es in ein TPU-kompatibles Modell. Wenden Sie einen Kubernetes-Pod an, der das vortrainierte Modell herunterlädt, den Cloud TPU Converter verwendet und die konvertierten Modelle mit dem CSI-Treiber für Cloud Storage FUSE in einem Cloud Storage-Bucket speichert. Der Cloud TPU Converter benötigt keine spezielle Hardware. In dieser Anleitung erfahren Sie, wie Sie das Modell herunterladen und den Cloud TPU Converter im CPU-Knotenpool ausführen.

  3. Starten Sie den Server für das konvertierte Modell. Wenden Sie ein Deployment an, das das Modell mit einem Server-Framework bereitstellt, das von dem im nichtflüchtigen ReadOnlyMany-Volume (ROX) gespeicherten Volume unterstützt wird. Die Bereitstellungsreplikate müssen auf einem v5e-Pod-TPU-Knoten mit einem Kubernetes-Pod pro Knoten ausgeführt werden.

  4. Stellen Sie einen Load-Balancer bereit, um den Modellserver zu testen. Der Server wird über den LoadBalancer-Dienst externen Anfragen ausgesetzt. Ein Python-Skript enthält eine Beispielanfrage zum Testen des Modellservers.

Das folgende Diagramm zeigt, wie Anfragen vom Load-Balancer weitergeleitet werden.

Ein Diagramm, das Load-Balancer-Routing zeigt

Beispiele für die Serverbereitstellung

Diese Beispielarbeitslasten werden mit den folgenden Annahmen konfiguriert:

  • Der Cluster wird mit einem TPU v5-Knotenpool mit 3 Knoten ausgeführt
  • Der Knotenpool verwendet den Maschinentyp ct5lp-hightpu-1t, wobei Folgendes gilt:
    • Topologie ist 1 × 1
    • Anzahl der TPU-Chips ist 1

Das folgende GKE-Manifest definiert ein Deployment mit einem einzelnen Hostserver.

apiVersion: apps/v1
kind: Deployment
metadata:
  name: bert-deployment
spec:
  selector:
    matchLabels:
      app: tf-bert-server
  replicas: 3 # number of nodes in node pool
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
      labels:
        app: tf-bert-server
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 1x1  # target topology
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice  # target version
      containers:
      - name: serve-bert
        image: us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
        env:
        - name: MODEL_NAME
          value: "bert"
        volumeMounts:
        - mountPath: "/models/"
          name: bert-external-storage
        ports:
        - containerPort: 8500
        - containerPort: 8501
        - containerPort: 8431 # Port to export TPU usage metrics, if supported.
        resources:
          requests:
            google.com/tpu: 1 # TPU chip request
          limits:
            google.com/tpu: 1 # TPU chip request
      volumes:
      - name: bert-external-storage
        persistentVolumeClaim:
          claimName: external-storage-pvc

Wenn Sie eine andere Anzahl von Knoten in Ihrem TPU-Knotenpool verwenden, ändern Sie das Feld replicas in die Anzahl der Knoten.

Wenn auf Ihrem Standardcluster die GKE-Version 1.27 oder niedriger ausgeführt wird, fügen Sie Ihrem Manifest das folgende Feld hinzu: 

spec:
  securityContext:
    privileged: true

Sie müssen Kubernetes-Pods in GKE-Version 1.28 oder höher nicht im privilegierten Modus ausführen. Weitere Informationen finden Sie unter Container ohne privilegierten Modus ausführen.

Wenn Sie einen anderen Maschinentyp verwenden:

  • Legen Sie für cloud.google.com/gke-tpu-topology die Topologie für den verwendeten Maschinentyp fest.
  • Legen Sie beide google.com/tpu-Felder unter resources fest, damit sie der Anzahl der Chips für den entsprechenden Maschinentyp entsprechen.

Einrichtung

Laden Sie die Python-Skripts und YAML-Manifeste der Anleitung mit dem folgenden Befehl herunter:

git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

Wechseln Sie in das Verzeichnis single-host-inference:

cd ai-on-gke/gke-tpu-examples/single-host-inference/

Python-Umgebung einrichten

Die Python-Skripts, die Sie in dieser Anleitung verwenden, erfordern Python Version 3.9 oder höher. Denken Sie daran, das requirements.txt für jede Anleitung zu installieren, bevor Sie die Python-Testskripts ausführen.

Wenn Python in Ihrer lokalen Umgebung nicht ordnungsgemäß eingerichtet ist, können Sie die Python-Skripts in dieser Anleitung mit Cloud Shell herunterladen und ausführen.

Cluster einrichten

  1. Erstellen Sie einen Cluster mit dem Maschinentyp e2-standard-4.

    gcloud container clusters create cluster-name \
--region your-region \
--release-channel rapid \
--num-nodes=1 \
--machine-type=e2-standard-4 \
--workload-pool=project-id.svc.id.goog \
--addons GcsFuseCsiDriver

  2. Erstellen Sie den TPU-Knotenpool mit einem einzelnen Host.

Bei den Beispielarbeitslasten wird Folgendes vorausgesetzt:

  • Ihr Cluster wird mit einem TPU v5e-Knotenpool mit 3 Knoten ausgeführt.
  • Der TPU-Knotenpool verwendet den Maschinentyp ct5lp-hightpu-1t.

Wenn Sie eine andere Clusterkonfiguration als zuvor beschrieben verwenden, müssen Sie das Manifest für die Serverbereitstellung bearbeiten.

Für die JAX Stable Diffusion-Demo benötigen Sie einen CPU-Knotenpool mit einem Maschinentyp mit mehr als 16 Gi Arbeitsspeicher (z. B. e2-standard-4). Dies wird im Befehl gcloud container clusters create oder durch Hinzufügen eines zusätzlichen Knotenpools zum vorhandenen Cluster mit folgendem Befehl konfiguriert:

gcloud beta container node-pools create your-pool-name \
  --zone=your-cluster-zone \
  --cluster=your-cluster-name \
  --machine-type=e2-standard-4 \
  --num-nodes=1

Ersetzen Sie Folgendes:

  • your-pool-name: Der Name des Knotenpools, der erstellt werden soll.
  • your-cluster-zone: Die Zone, in der Ihr Cluster erstellt wurde.
  • your-cluster-name: Der Name des Clusters, dem der Knotenpool hinzugefügt werden soll.
  • your-machine-type: Der Maschinentyp der Knoten, die in Ihrem Knotenpool erstellt werden sollen.

Modellspeicher einrichten

Es gibt mehrere Möglichkeiten, Ihr Modell für die Bereitstellung zu speichern. In dieser Anleitung verwenden wir den folgenden Ansatz:

  • Zur Umwandlung des vortrainierten Modells für TPUs verwenden wir eine Virtual Private Cloud mit Persistent Disk mit ReadWriteMany-Zugriff (RWX).
  • Um das Modell auf mehreren TPUs mit einem einzelnen Host bereitzustellen, verwenden wir dieselbe VPC, die vom Cloud Storage-Bucket unterstützt wird.

Führen Sie den folgenden Befehl aus, um einen Cloud Storage-Bucket zu erstellen.

gcloud storage buckets create gs://your-bucket-name \
  --project=your-bucket-project-id \
  --location=your-bucket-location

Ersetzen Sie Folgendes:

  • your-bucket-name: Der Name des Cloud Storage-Bucket.
  • your-bucket-project-id: Die Projekt-ID, in der Sie den Cloud Storage-Bucket erstellt haben.
  • your-bucket-location: Der Standort Ihres Cloud Storage-Bucket. Geben Sie den Standort an, an dem Ihr GKE-Cluster ausgeführt wird, um die Leistung zu verbessern.

Führen Sie die folgenden Schritte aus, um dem GKE-Cluster Zugriff auf den Bucket zu gewähren. In den folgenden Beispielen werden der Standard-Namespace und das Kubernetes-Standarddienstkonto verwendet, um die Einrichtung zu vereinfachen. Weitere Informationen finden Sie unter Zugriff auf Cloud Storage-Buckets mit GKE Workload Identity konfigurieren.

  1. Erstellen Sie ein IAM-Dienstkonto für Ihre Anwendung oder verwenden Sie stattdessen ein vorhandenes IAM-Dienstkonto. Sie können ein beliebiges IAM-Dienstkonto im Projekt Ihres Cloud Storage-Bucket verwenden.

    gcloud iam service-accounts create your-iam-service-acct \
--project=your-bucket-project-id

    Ersetzen Sie Folgendes:

    • your-iam-service-acct ist der Name des neuen IAM-Dienstkontos.
    • your-bucket-project-id: die ID des Projekts, in dem Sie Ihr IAM-Dienstkonto erstellt haben. Das IAM-Dienstkonto muss sich im selben Projekt wie Ihr Cloud Storage-Bucket befinden.

  2. Prüfen Sie, ob Ihr IAM-Dienstkonto die erforderlichen Speicherrollen hat.

    gcloud storage buckets add-iam-policy-binding gs://your-bucket-name \
--member "serviceAccount:your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com" \
--role "roles/storage.objectAdmin"

    Ersetzen Sie Folgendes:

    • your-bucket-name: Der Name Ihres Cloud Storage-Buckets
    • your-iam-service-acct ist der Name des neuen IAM-Dienstkontos.
    • your-bucket-project-id: die ID des Projekts, in dem Sie Ihr IAM-Dienstkonto erstellt haben.

  3. Erlauben Sie dem Kubernetes-Dienstkonto, die Identität des IAM-Dienstkontos anzunehmen. Fügen Sie dazu eine IAM-Richtlinienbindung zwischen den beiden Dienstkonten hinzu. Durch diese Bindung kann das Kubernetes-Dienstkonto als IAM-Dienstkonto fungieren.

    gcloud iam service-accounts add-iam-policy-binding your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:your-project-id.svc.id.goog[default/default]"

    Ersetzen Sie Folgendes:

    • your-iam-service-acct ist der Name des neuen IAM-Dienstkontos.
    • your-bucket-project-id: die ID des Projekts, in dem Sie Ihr IAM-Dienstkonto erstellt haben.
    • your-project-id: die ID des Projekts, in dem Sie den GKE-Cluster erstellt haben. Ihre Cloud Storage-Buckets und Ihr GKE-Cluster können sich im selben Projekt oder in verschiedenen Projekten befinden.

  4. Kennzeichnen Sie das Kubernetes-Dienstkonto mit der E-Mail-Adresse des IAM-Dienstkontos.

    kubectl annotate serviceaccount default \
  --namespace default \
  iam.gke.io/gcp-service-account=your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com

    Ersetzen Sie Folgendes:

    • your-iam-service-acct ist der Name des neuen IAM-Dienstkontos.
    • your-bucket-project-id: die ID des Projekts, in dem Sie Ihr IAM-Dienstkonto erstellt haben.

  5. Führen Sie den folgenden Befehl aus, um den Bucket-Namen in die YAML-Dateien dieser Demo zu übertragen:

    find . -type f -name "*.yaml" | xargs sed -i "s/BUCKET_NAME/your-bucket-name/g"

    Ersetzen Sie your-bucket-name durch den Namen Ihres Cloud Storage-Buckets.

  6. Erstellen Sie mit dem folgenden Befehl die Anforderungen für nichtflüchtige Volumes und nichtflüchtige Volumes:

    kubectl apply -f pvc-pv.yaml

Inferenz und Bereitstellung des JAX-Modells

Installieren Sie Python-Abhängigkeiten zum Ausführen von Python-Skripts der Anleitung, die Anfragen an den JAX-Modelldienst senden.

pip install -r jax/requirements.txt

Demo zur Bereitstellung von JAX BERT E2E ausführen:

In dieser Demo wird ein vortrainiertes BERT-Modell von Hugging Face verwendet.

Der Kubernetes-Pod führt die folgenden Schritte aus:

  1. Es lädt das Python-Skript export_bert_model.py aus den Beispielressourcen herunter und verwendet es, um das vortrainierte bert-Modell in ein temporäres Verzeichnis herunterzuladen.
  2. Verwendet das Cloud TPU Converter-Image, um das vortrainierte Modell von CPU in TPU zu konvertieren, und speichert das Modell in dem Cloud Storage-Bucket, den Sie während der setup erstellt haben.

Dieser Kubernetes-Pod ist für die Ausführung auf der Standardknotenpool-CPU konfiguriert. Führen Sie den Pod mit dem folgenden Befehl aus:

kubectl apply -f jax/bert/install-bert.yaml

Prüfen Sie folgendermaßen, ob das Modell korrekt installiert wurde:

kubectl get pods install-bert

Es kann einige Minuten dauern, bis STATUS Completed gelesen hat.

TF-Modellserver für das Modell starten

Bei den Beispielarbeitslasten in dieser Anleitung wird Folgendes vorausgesetzt:

  • Der Cluster wird mit einem TPU v5-Knotenpool mit drei Knoten ausgeführt
  • Der Knotenpool verwendet den Maschinentyp ct5lp-hightpu-1t, der einen TPU-Chip enthält.

Wenn Sie eine andere Clusterkonfiguration als zuvor beschrieben verwenden, müssen Sie das Manifest für die Serverbereitstellung bearbeiten.

Bereitstellung anwenden
kubectl apply -f jax/bert/serve-bert.yaml

Prüfen Sie so, ob der Server ausgeführt wird:

kubectl get deployment bert-deployment

Es kann einen Moment dauern, bis AVAILABLE 3 gelesen hat.

Load-Balancer-Dienst anwenden
kubectl apply -f jax/bert/loadbalancer.yaml

Prüfen Sie so, ob der Load-Balancer für externen Traffic bereit ist:

kubectl get svc tf-bert-service

Es kann einige Minuten dauern, bis für EXTERNAL_IP eine IP-Adresse aufgeführt wird.

Anfrage an den Modellserver senden

Rufen Sie die externe IP-Adresse vom Load-Balancer-Dienst ab:

EXTERNAL_IP=$(kubectl get services tf-bert-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

Führen Sie ein Skript aus, um eine Anfrage an den Server zu senden:

python3 jax/bert/bert_request.py $EXTERNAL_IP

Erwartete Ausgabe:

For input "The capital of France is [MASK].", the result is ". the capital of france is paris.."
For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."
Bereinigung

Führen Sie kubectl delete in umgekehrter Reihenfolge aus, um Ressourcen zu bereinigen.

kubectl delete -f jax/bert/loadbalancer.yaml
kubectl delete -f jax/bert/serve-bert.yaml
kubectl delete -f jax/bert/install-bert.yaml

Demo für JAX Stable Diffusion E2E-Bereitstellung ausführen

In dieser Demo wird das vortrainierte Stable Diffusion Model von „Gesichtsumarmung“ verwendet.

TPU-kompatibles TF2-Modell aus Flax Stable Diffusion-Modell exportieren

Für den Export der stabilen Diffusionsmodelle muss der Cluster einen CPU-Knotenpool mit einem Maschinentyp mit mehr als 16 Gi Arbeitsspeicher haben, wie unter Cluster einrichten beschrieben.

Der Kubernetes-Pod führt die folgenden Schritte aus:

  1. Es lädt das Python-Skript export_stable_diffusion_model.py aus den Beispielressourcen herunter und verwendet es, um das vortrainierte stabile Diffusionsmodell in ein temporäres Verzeichnis herunterzuladen.
  2. Verwendet das Cloud TPU Converter-Image, um das vortrainierte Modell von CPU in TPU zu konvertieren, und speichert das Modell in dem Cloud Storage-Bucket, den Sie während der Speichereinrichtung erstellt haben.

Dieser Kubernetes-Pod ist für die Ausführung auf dem Standard-CPU-Knotenpool konfiguriert. Führen Sie den Pod mit dem folgenden Befehl aus:

kubectl apply -f jax/stable-diffusion/install-stable-diffusion.yaml

Prüfen Sie folgendermaßen, ob das Modell korrekt installiert wurde:

kubectl get pods install-stable-diffusion

Es kann einige Minuten dauern, bis STATUS Completed gelesen hat.

Servercontainer des TF-Modells für das Modell starten

Die Beispielarbeitslasten wurden mit den folgenden Annahmen konfiguriert:

  • Der Cluster wird mit einem TPU v5-Knotenpool mit drei Knoten ausgeführt
  • Der Knotenpool verwendet den Maschinentyp ct5lp-hightpu-1t, wobei Folgendes gilt:
    • Topologie ist 1 × 1
    • Anzahl der TPU-Chips ist 1

Wenn Sie eine andere Clusterkonfiguration als zuvor beschrieben verwenden, müssen Sie das Manifest für die Serverbereitstellung bearbeiten.

Wenden Sie die Bereitstellung an:

kubectl apply -f jax/stable-diffusion/serve-stable-diffusion.yaml

Prüfen Sie, ob der Server wie erwartet ausgeführt wird:

kubectl get deployment stable-diffusion-deployment

Es kann einen Moment dauern, bis AVAILABLE 3 gelesen hat.

Wenden Sie den Load-Balancer-Dienst an:

kubectl apply -f jax/stable-diffusion/loadbalancer.yaml

Prüfen Sie so, ob der Load-Balancer für externen Traffic bereit ist:

kubectl get svc tf-stable-diffusion-service

Es kann einige Minuten dauern, bis für EXTERNAL_IP eine IP-Adresse aufgeführt wird.

Anfrage an den Modellserver senden

Rufen Sie eine externe IP-Adresse vom Load-Balancer ab:

EXTERNAL_IP=$(kubectl get services tf-stable-diffusion-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

Skript zum Senden einer Anfrage an den Server ausführen

python3 jax/stable-diffusion/stable_diffusion_request.py $EXTERNAL_IP

Erwartete Ausgabe:

Der Prompt ist Painting of a squirrel skating in New York und das Ausgabebild wird als stable_diffusion_images.jpg in Ihrem aktuellen Verzeichnis gespeichert.

Bereinigung

Führen Sie kubectl delete in umgekehrter Reihenfolge aus, um Ressourcen zu bereinigen.

kubectl delete -f jax/stable-diffusion/loadbalancer.yaml
kubectl delete -f jax/stable-diffusion/serve-stable-diffusion.yaml
kubectl delete -f jax/stable-diffusion/install-stable-diffusion.yaml

Demo für TensorFlow ResNet-50 E2E ausführen:

Installieren Sie Python-Abhängigkeiten zum Ausführen von Python-Skripts der Anleitung, die Anfragen an den TF-Modelldienst senden.

pip install -r tf/resnet50/requirements.txt
Schritt 1: Modell konvertieren

Modellkonvertierung anwenden:

kubectl apply -f tf/resnet50/model-conversion.yml

Prüfen Sie folgendermaßen, ob das Modell korrekt installiert wurde:

kubectl get pods resnet-model-conversion

Es kann einige Minuten dauern, bis STATUS Completed gelesen hat.

Schritt 2: Modell mit TensorFlow bereitstellen

Wenden Sie die Bereitstellung der Modellbereitstellung an:

kubectl apply -f tf/resnet50/deployment.yml

Prüfen Sie mit dem folgenden Befehl, ob der Server wie erwartet ausgeführt wird:

kubectl get deployment resnet-deployment

Es kann einen Moment dauern, bis AVAILABLE 3 gelesen hat.

Wenden Sie den Load-Balancer-Dienst an:

kubectl apply -f tf/resnet50/loadbalancer.yml

Prüfen Sie so, ob der Load-Balancer für externen Traffic bereit ist:

kubectl get svc resnet-service

Es kann einige Minuten dauern, bis für EXTERNAL_IP eine IP-Adresse aufgeführt wird.

Schritt 3: Testanfrage an Modellserver senden

Rufen Sie die externe IP-Adresse vom Load-Balancer ab:

EXTERNAL_IP=$(kubectl get services resnet-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

Führen Sie das Testanfrageskript (HTTP) aus, um die Anfrage an den Modellserver zu senden.

python3 tf/resnet50/request.py --host $EXTERNAL_IP

Die Antwort sollte in etwa so aussehen:

Predict result: ['ImageNet ID: n07753592, Label: banana, Confidence: 0.94921875',
'ImageNet ID: n03532672, Label: hook, Confidence: 0.0223388672', 'ImageNet ID: n07749582,
Label: lemon, Confidence: 0.00512695312
Schritt 4: Bereinigen

Führen Sie die folgenden kubectl delete-Befehle aus, um Ressourcen zu bereinigen:

kubectl delete -f tf/resnet50/loadbalancer.yml
kubectl delete -f tf/resnet50/deployment.yml
kubectl delete -f tf/resnet50/model-conversion.yml

Löschen Sie den Knotenpool und den Cluster von GKE, wenn Sie sie nicht mehr benötigen.

PyTorch-Modellinferenz und -bereitstellung

Installieren Sie Python-Abhängigkeiten, um Python-Skripts der Anleitung auszuführen, die Anfragen an den PyTorch-Modelldienst senden:

pip install -r pt/densenet161/requirements.txt

Demo zur Bereitstellung von TorchServe Densenet161 E2E ausführen:

  1. Modellarchiv generieren.

    1. Modellarchiv anwenden:
    kubectl apply -f pt/densenet161/model-archive.yml
    1. Prüfen Sie folgendermaßen, ob das Modell korrekt installiert wurde:
    kubectl get pods densenet161-model-archive

    Es kann einige Minuten dauern, bis STATUS Completed gelesen hat.

  2. Stellen Sie das Modell mit TorchServe bereit:

    1. Bereitstellung der Modellbereitstellung anwenden:

      kubectl apply -f pt/densenet161/deployment.yml

    2. Prüfen Sie mit dem folgenden Befehl, ob der Server wie erwartet ausgeführt wird:

      kubectl get deployment densenet161-deployment

      Es kann einen Moment dauern, bis AVAILABLE 3 gelesen hat.

    3. Wenden Sie den Load-Balancer-Dienst an:

      kubectl apply -f pt/densenet161/loadbalancer.yml

      Prüfen Sie mit dem folgenden Befehl, ob der Load-Balancer für externen Traffic bereit ist:

      kubectl get svc densenet161-service

      Es kann einige Minuten dauern, bis für EXTERNAL_IP eine IP-Adresse aufgeführt wird.

  3. Testanfrage an Modellserver senden:

    1. Rufen Sie die externe IP-Adresse vom Load-Balancer ab:

      EXTERNAL_IP=$(kubectl get services densenet161-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

    2. Führen Sie das Testanfrageskript aus, um die Anfrage (HTTP) an den Modellserver zu senden:

      python3 pt/densenet161/request.py --host $EXTERNAL_IP

      Die Antwort sollte in etwa so aussehen:

      Request successful. Response: {'tabby': 0.47878125309944153, 'lynx': 0.20393909513950348, 'tiger_cat': 0.16572578251361847, 'tiger': 0.061157409101724625, 'Egyptian_cat': 0.04997897148132324

  4. Bereinigen Sie Ressourcen mit den folgenden kubectl delete-Befehlen:

    kubectl delete -f pt/densenet161/loadbalancer.yml
kubectl delete -f pt/densenet161/deployment.yml
kubectl delete -f pt/densenet161/model-archive.yml

    Löschen Sie den Knotenpool und den Cluster von GKE, wenn Sie sie nicht mehr benötigen.

Fehlerbehebung – allgemeine Probleme

Informationen zur GKE-Fehlerbehebung finden Sie unter Fehlerbehebung für TPU in GKE.

TPU-Initialisierung fehlgeschlagen

Wenn der folgende Fehler auftritt, prüfen Sie, ob Sie den TPU-Container entweder im privilegierten Modus ausführen oder die ulimit im Container erhöht haben. Weitere Informationen finden Sie unter Ohne privilegierten Modus ausführen.

TPU platform initialization failed: FAILED_PRECONDITION: Couldn't mmap: Resource
temporarily unavailable.; Unable to create Node RegisterInterface for node 0,
config: device_path:      "/dev/accel0" mode: KERNEL debug_data_directory: ""
dump_anomalies_only: true crash_in_debug_dump: false allow_core_dump: true;
could not create driver instance

Planungs-Deadlock

Angenommen, Sie haben zwei Jobs (Job A und Job B), die beide auf TPU-Slices mit einer bestimmten TPU-Topologie (z. B. v4-32) geplant werden sollen. Angenommen, Sie haben zwei v4-32-TPU-Slices im GKE-Cluster. Diese nennen wir Slice X und Segment Y. Da Ihr Cluster genügend Kapazität hat, um beide Jobs zu planen, sollten beide Jobs theoretisch schnell geplant werden – ein Job auf jedem der beiden TPU-v4-32-Slices.

Allerdings kann es ohne sorgfältige Planung in einen Planungsdeadlock geraten. Angenommen, der Kubernetes-Planer plant einen Kubernetes-Pod aus Job A im Segment X und dann einen Kubernetes-Pod von Job B im Segment X. In diesem Fall versucht der Planer unter Berücksichtigung der Kubernetes-Pod-Affinitätsregeln für Job A, alle verbleibenden Kubernetes-Pods für Job A im Segment X zu planen. Dasselbe für Job B. Somit kann weder Job A noch Job B vollständig für ein einzelnes Segment geplant werden. Das Ergebnis ist ein Planungsdeadlock.

Um das Risiko eines Planungsdeadlocks zu vermeiden, können Sie die Kubernetes-Pod-Anti-Affinität mit cloud.google.com/gke-nodepool als topologyKey verwenden, wie im folgenden Beispiel gezeigt:

apiVersion: batch/v1
kind: Job
metadata:
 name: pi
spec:
 parallelism: 2
 template:
   metadata:
     labels:
       job: pi
   spec:
     affinity:
       podAffinity:
         requiredDuringSchedulingIgnoredDuringExecution:
         - labelSelector:
             matchExpressions:
             - key: job
               operator: In
               values:
               - pi
           topologyKey: cloud.google.com/gke-nodepool
       podAntiAffinity:
         requiredDuringSchedulingIgnoredDuringExecution:
         - labelSelector:
             matchExpressions:
             - key: job
               operator: NotIn
               values:
               - pi
           topologyKey: cloud.google.com/gke-nodepool
           namespaceSelector:
             matchExpressions:
             - key: kubernetes.io/metadata.name
               operator: NotIn
               values:
               - kube-system
     containers:
     - name: pi
       image: perl:5.34.0
       command: ["sleep",  "60"]
     restartPolicy: Never
 backoffLimit: 4

TPU-Knotenpoolressourcen mit Terraform erstellen

Sie können auch Terraform verwenden, um die Cluster- und Knotenpoolressourcen zu verwalten.

TPU-Slice-Knotenpool mit mehreren Hosts in einem vorhandenen GKE-Cluster erstellen

Wenn Sie einen vorhandenen Cluster in einem TPU-Knotenpool mit mehreren Hosts erstellen möchten, können Sie das folgende Terraform-Snippet verwenden:

resource "google_container_cluster" "cluster_multi_host" {
  …
  release_channel {
    channel = "RAPID"
  }
  workload_identity_config {
    workload_pool = "my-gke-project.svc.id.goog"
  }
  addons_config {
    gcs_fuse_csi_driver_config {
      enabled = true
    }
  }
}

resource "google_container_node_pool" "multi_host_tpu" {
  provider           = google-beta
  project            = "${project-id}"
  name               = "${node-pool-name}"
  location           = "${location}"
  node_locations     = ["${node-locations}"]
  cluster            = google_container_cluster.cluster_multi_host.name
  initial_node_count = 2

  node_config {
    machine_type = "ct4p-hightpu-4t"
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = ["${reservation-name}"]
    }
    workload_metadata_config {
      mode = "GKE_METADATA"
    }
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = "2x2x2"
  }
}

Ersetzen Sie die folgenden Werte:

  • your-project: Ihr Google Cloud-Projekt, in dem Sie die Arbeitslast ausführen.
  • your-node-pool: Der Name des Knotenpools, den Sie erstellen.
  • us-central2: Die Region, in der Sie Ihre Arbeitslast ausführen.
  • us-central2-b: Die Zone, in der Sie Ihre Arbeitslast ausführen.
  • your-reservation-name: Der Name Ihrer Reservierung.

TPU-Slice-Knotenpool mit einem einzelnen Host in einem vorhandenen GKE-Cluster erstellen

Verwenden Sie das folgende Terraform-Snippet:

resource "google_container_cluster" "cluster_single_host" {
  …
  cluster_autoscaling {
    autoscaling_profile = "OPTIMIZE_UTILIZATION"
  }
  release_channel {
    channel = "RAPID"
  }
  workload_identity_config {
  workload_pool = "${project-id}.svc.id.goog"
  }
  addons_config {
    gcs_fuse_csi_driver_config {
      enabled = true
    }
  }
}

resource "google_container_node_pool" "single_host_tpu" {
  provider           = google-beta
  project            = "${project-id}"
  name               = "${node-pool-name}"
  location           = "${location}"
  node_locations     = ["${node-locations}"]
  cluster            = google_container_cluster.cluster_single_host.name
  initial_node_count = 0
  autoscaling {
    total_min_node_count = 2
    total_max_node_count = 22
    location_policy      = "ANY"
  }

  node_config {
    machine_type = "ct4p-hightpu-4t"
    workload_metadata_config {
      mode = "GKE_METADATA"
    }
  }
}

Ersetzen Sie die folgenden Werte:

  • your-project: Ihr Google Cloud-Projekt, in dem Sie die Arbeitslast ausführen.
  • your-node-pool: Der Name des Knotenpools, den Sie erstellen.
  • us-central2: Die Region, in der Sie Ihre Arbeitslast ausführen.
  • us-central2-b: Die Zone, in der Sie Ihre Arbeitslast ausführen.