TPU-Arbeitslasten in GKE Autopilot bereitstellen

Auf dieser Seite wird beschrieben, wie Sie Arbeitslasten für maschinelles Lernen (ML) beschleunigen, indem Sie Cloud TPU-Beschleuniger (TPUs) in Google Kubernetes Engine (GKE) Autopilot-Clustern verwenden. Diese Anleitung kann Ihnen helfen, die richtigen Bibliotheken für Ihre ML-Anwendungsframeworks auszuwählen, Ihre TPU-Arbeitslasten für die optimale Ausführung in GKE einzurichten und Ihre Arbeitslasten nach der Bereitstellung zu überwachen.

Diese Seite richtet sich an Plattformadministratoren und ‑operatoren, Daten- und KI-Spezialisten sowie Anwendungsentwickler, die ML-Arbeitslasten auf TPUs vorbereiten und ausführen möchten. Weitere Informationen zu den gängigen Rollen, Verantwortlichkeiten und Beispielaufgaben, auf die wir in Google Cloud -Inhalten verweisen, finden Sie unter Häufig verwendete GKE-Nutzerrollen und -Aufgaben.

Machen Sie sich vor dem Lesen dieser Seite mit den folgenden Ressourcen vertraut:

Funktionsweise von TPUs in Autopilot

Wenn Sie TPUs in Autopilot-Arbeitslasten verwenden möchten, geben Sie Folgendes in Ihrem Arbeitslastmanifest an:

  • Die TPU-Version im Feld spec.nodeSelector.
  • Die TPU-Topologie im Feld spec.nodeSelector. Die Topologie muss von der angegebenen TPU-Version unterstützt werden.
  • Die Anzahl der TPU-Chips in den Feldern spec.containers.resources.requests und spec.containers.resources.limits.

Wenn Sie die Arbeitslast bereitstellen, stellt GKE Knoten mit der angeforderten TPU-Konfiguration bereit und plant Ihre Pods auf den Knoten. GKE platziert die einzelnen Arbeitslasten auf je eigenen Knoten, damit jeder Pod bei minimalem Störungsrisiko auf die vollständigen Ressourcen des Knotens zugreifen kann.

TPUs in Autopilot sind mit den folgenden Funktionen kompatibel:

  1. Spot-Pods
  2. Bestimmte Kapazitätsreservierungen
  3. Pods mit verlängerter Laufzeit
  4. Flex-Start

TPU-Konfiguration planen

Bevor Sie diese Anleitung zum Bereitstellen von TPU-Arbeitslasten verwenden, sollten Sie Ihre TPU-Konfiguration basierend auf Ihrem Modell und dem erforderlichen Arbeitsspeicher planen. Weitere Informationen finden Sie unter TPU-Konfiguration planen.

Preise

Preisinformationen finden Sie unter Autopilot-Preise.

Hinweise

Führen Sie die folgenden Aufgaben aus, bevor Sie beginnen:

  • Aktivieren Sie die Google Kubernetes Engine API.
    • Google Kubernetes Engine API aktivieren
  • Wenn Sie die Google Cloud CLI für diesen Task 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 dem Befehl gcloud components update ab. In früheren gcloud CLI-Versionen werden die Befehle in diesem Dokument möglicherweise nicht unterstützt.
  • Sie benötigen einen Autopilot-Cluster mit GKE-Version 1.32.3-gke.1927000 oder höher. Eine Anleitung finden Sie unter Autopilot-Cluster erstellen.
  • Wenn Sie reservierte TPUs verwenden möchten, benötigen Sie eine spezifische Kapazitätsreservierung. Eine Anleitung finden Sie unter Reservierung nutzen.

Kontingent für TPUs und andere GKE-Ressourcen sicherstellen

In den folgenden Abschnitten erfahren Sie, wie Sie sicherstellen können, dass Sie bei der Verwendung von TPUs in GKE über ausreichendes Kontingent verfügen.

Zum Erstellen von TPU-Slice-Knoten benötigen Sie TPU-Kontingent, sofern Sie keine vorhandene Kapazitätsreservierung verwenden. Wenn Sie reservierte TPUs verwenden, überspringen Sie diesen Abschnitt.

Zum Erstellen von TPU-Slice-Knoten in GKE ist ein Compute Engine API-Kontingent (compute.googleapis.com) und kein Cloud TPU API-Kontingent (tpu.googleapis.com) erforderlich. Der Name des Kontingents ist in regulären Autopilot-Pods und in Spot-Pods unterschiedlich.

Führen Sie die folgenden Schritte aus, um das Limit und die aktuelle Nutzung Ihres Compute Engine API-Kontingents für TPUs zu prüfen:

  1. Rufen Sie in der Google Cloud Console die Seite Kontingente auf:

    Kontingente aufrufen

  2. Gehen Sie im Feld Filter folgendermaßen vor:

    1. Verwenden Sie die folgende Tabelle, um das Attribut des Kontingents basierend auf der TPU-Version und dem Wert im Knotenselektor cloud.google.com/gke-tpu-accelerator auszuwählen und zu kopieren. Wenn Sie beispielsweise On-Demand-TPU v5e-Knoten erstellen möchten, deren Wert im Knotenselektor cloud.google.com/gke-tpu-accelerator tpu-v5-lite-podslice ist, geben Sie Name: TPU v5 Lite PodSlice chips ein.

      TPU-Version, cloud.google.com/gke-tpu-accelerator Property und Name des Kontingents für On-Demand-Instanzen Property und Name des Kontingents für Spot2-Instanzen
      TPU v3,
      tpu-v3-device 
      Dimensions (e.g. location):
      tpu_family:CT3
      		 Nicht zutreffend
      TPU v3,
      tpu-v3-slice 
      Dimensions (e.g. location):
      tpu_family:CT3P
      		 Nicht zutreffend
      TPU v4,
      tpu-v4-podslice 
      Name:
      TPU v4 PodSlice chips
      		 
      Name:
      Preemptible TPU v4 PodSlice chips
      TPU v5e,
      tpu-v5-lite-podslice 
      Name:
      TPU v5 Lite PodSlice chips
      		 
      Name:
      Preemptible TPU v5 Lite Podslice
      chips
      TPU v5p,
      tpu-v5p-slice 
      Name:
      TPU v5p chips
      		 
      Name:
      Preemptible TPU v5p chips
      TPU Trillium,
      tpu-v6e-slice 
      Dimensions (e.g. location):
      tpu_family:CT6E
      		 
      Name:
      Preemptible TPU slices v6e

    2. Wählen Sie das Attribut Dimensionen (z.B. Standorte) aus und geben Sie region: gefolgt vom Namen der Region ein, in der Sie TPUs in GKE erstellen möchten. Geben Sie beispielsweise region:us-west4 ein, wenn Sie TPU-Slice-Knoten in der Zone us-west4-a erstellen möchten. Das TPU-Kontingent ist regional, d. h. alle Zonen innerhalb derselben Region nutzen dasselbe TPU-Kontingent.

Wenn keine Kontingente vorhanden sind, die Sie eingegeben haben, wurde dem Projekt keines der angegebenen Kontingente für die gewünschte Region zugewiesen und Sie müssen eine TPU-Kontingentanpassung anfordern.

Wenn eine TPU-Reservierung erstellt wird, erhöhen sich sowohl der Limitwert als auch der Wert der aktuellen Nutzung für das entsprechende Kontingent um die Anzahl der Chips in der TPU-Reservierung. Wenn beispielsweise eine Reservierung für 16 TPU v5e-Chips erstellt wird, deren Wert im Knotenselektor cloud.google.com/gke-tpu-accelerator tpu-v5-lite-podslice ist, werden sowohl die Begrenzung als auch die aktuelle Nutzung für das Kontingent TPU v5 Lite PodSlice chips in der entsprechenden Region um 16 erhöht.

Kontingente für zusätzliche GKE-Ressourcen

Möglicherweise müssen Sie die folgenden GKE-bezogenen Kontingente in den Regionen erhöhen, in denen GKE Ihre Ressourcen erstellt.

  • Kontingent für Persistent Disk SSD (GB): Das Bootlaufwerk jedes Kubernetes-Knotens benötigt standardmäßig 100 GB. Daher sollte dieses Kontingent mindestens so hoch wie das Produkt aus der maximalen Anzahl der GKE-Knoten, die Sie voraussichtlich erstellen werden, und 100 GB (Knoten × 100 GB) sein.
  • Kontingent für verwendete IP-Adressen: Jeder Kubernetes-Knoten verbraucht eine IP-Adresse. Daher sollte dieses Kontingent mindestens so hoch sein wie die maximale Anzahl von GKE-Knoten, die Sie voraussichtlich erstellen werden.
  • Achten Sie darauf, dass max-pods-per-node mit dem Subnetzbereich übereinstimmt: Jeder Kubernetes-Knoten verwendet sekundäre IP-Bereiche für Pods. Für max-pods-per-node von 32 sind beispielsweise 64 IP-Adressen erforderlich, was einem /26-Subnetz pro Knoten entspricht. Dieser Bereich darf nicht für andere Cluster freigegeben werden. Verwenden Sie das Flag --max-pods-per-node, um die Anzahl der Pods zu begrenzen, die auf einem Knoten geplant werden dürfen, damit der IP-Adressbereich nicht ausgeschöpft wird. Das Kontingent für max-pods-per-node sollte mindestens so hoch sein wie die maximale Anzahl von GKE-Knoten, die Sie voraussichtlich erstellen werden.

Informationen zum Anfordern einer Kontingenterhöhung finden Sie unter Kontingentanpassung anfordern.

Optionen für die Bereitstellung von TPUs in GKE

Mit GKE Autopilot können Sie TPUs direkt in einzelnen Arbeitslasten verwenden, indem Sie Kubernetes-nodeSelectors verwenden.

Alternativ können Sie TPUs mit benutzerdefinierten Compute-Klassen anfordern. Mit benutzerdefinierten Compute-Klassen können Plattformadministratoren eine Hierarchie von Knotenkonfigurationen für GKE definieren, die bei Entscheidungen zur Knotenskalierung priorisiert werden sollen, damit Arbeitslasten auf der ausgewählten Hardware ausgeführt werden.

Eine Anleitung finden Sie im Abschnitt TPUs mit benutzerdefinierten Compute-Klassen zentral bereitstellen.

TPU-Anwendung vorbereiten

Für TPU-Arbeitslasten gelten die folgenden Vorbereitungsanforderungen.

  1. Frameworks wie JAX, PyTorch und TensorFlow greifen auf TPU-VMs mit der freigegebenen Bibliothek libtpu zu. libtpu enthält den XLA-Compiler, die TPU-Laufzeitsoftware und den TPU-Treiber. Jede Version von PyTorch und JAX erfordert eine bestimmte Version von libtpu.so. Um Konflikte bei Paketversionen zu vermeiden, empfehlen wir die Verwendung eines JAX AI-Images. Für die Verwendung von TPUs in GKE benötigen Sie die folgenden Versionen:
    TPU-Typ libtpu.so-Version
    TPU Trillium (v6e)
    tpu-v6e-slice
    TPU v5e
    tpu-v5-lite-podslice
    TPU v5p
    tpu-v5p-slice
    TPU v4
    tpu-v4-podslice
    TPU v3
    tpu-v3-slice
    tpu-v3-device
  2. Legen Sie die folgenden Umgebungsvariablen für den Container fest, der die TPU-Ressourcen anfordert:
    • TPU_WORKER_ID: Eine eindeutige Ganzzahl für jeden Pod. Diese ID bezeichnet eine eindeutige Worker-ID im TPU-Slice. Die unterstützten Werte für dieses Feld reichen von null bis zur Anzahl der Pods minus eins.
    • TPU_WORKER_HOSTNAMES: Eine durch Kommas getrennte Liste von TPU-VM-Hostnamen oder IP-Adressen, die innerhalb des Slices miteinander kommunizieren müssen. Für jede TPU-VM im Slice sollte ein Hostname oder eine IP-Adresse vorhanden sein. Die Liste der IP-Adressen oder Hostnamen ist geordnet und ab null nach der TPU_WORKER_ID indexiert.

      • GKE fügt diese Umgebungsvariablen automatisch ein. Dazu wird ein mutierender Webhook verwendet, wenn ein Job mit completionMode: Indexed, subdomain, parallelism > 1 erstellt wird, und es werden google.com/tpu-Attribute angefordert. GKE fügt einen monitorlosen Service hinzu, sodass die DNS-Einträge für die Pods, die den Service unterstützen, hinzugefügt werden.

Nachdem Sie die Arbeitslastvorbereitung abgeschlossen haben, können Sie einen Job ausführen, der TPUs verwendet.

TPUs in einer Arbeitslast anfordern

In diesem Abschnitt erfahren Sie, wie Sie einen Job erstellen, in dem TPUs in Autopilot angefordert werden. In jeder Arbeitslast, für die TPUs erforderlich sind, müssen Sie Folgendes angeben:

  • Knotenselektoren für die TPU-Version und -Topologie
  • Die Anzahl der TPU-Chips für einen Container in Ihrer Arbeitslast

Eine Liste der unterstützten TPU-Versionen, Topologien und der entsprechenden Anzahl von TPU-Chips und Knoten in einem Slice finden Sie unter TPU-Version auswählen.

Hinweise zu TPU-Anfragen in Arbeitslasten

Nur ein Container in einem Pod kann TPUs verwenden. Die Anzahl der TPU-Chips, die ein Container anfordert, muss der Anzahl der TPU-Chips entsprechen, die an einen Knoten im Slice angehängt sind. Wenn Sie beispielsweise TPU v5e (tpu-v5-lite-podslice) mit einer 2x4-Topologie anfragen, können Sie eine der folgenden Optionen anfordern:

  • 4-Chips, die zwei Knoten mit mehreren Hosts und jeweils 4 Chips erstellen
  • 8-Chips, die einen Knoten mit einem Host und 8 TPU-Chips erstellen

Als Best Practice zur Maximierung Ihrer Kosteneffizienz sollten Sie immer alle TPUs in dem von Ihnen angeforderten Slice nutzen. Wenn Sie ein Slice mit mehreren Hosts und zwei Knoten mit jeweils 4 TPU-Chips anfordern, sollten Sie eine Arbeitslast bereitstellen, die auf beiden Knoten ausgeführt wird und alle 8 TPU-Chips im Slice verbraucht.

Arbeitslast erstellen, die TPUs anfordert

Mit den folgenden Schritten wird ein Job erstellt, der TPUs anfordert. Wenn Sie Arbeitslasten haben, die auf TPU-Slices mit mehreren Hosts ausgeführt werden, müssen Sie auch einen monitorlosen Dienst erstellen, der Ihre Arbeitslast nach Namen auswählt. Mit diesem monitorlosen Service können Pods auf verschiedenen Knoten im Slice mit mehreren Hosts miteinander kommunizieren. Dazu wird die Kubernetes-DNS-Konfiguration so aktualisiert, dass sie auf die Pods in der Arbeitslast verweist.

  1. Speichern Sie das folgende Manifest als tpu-autopilot.yaml:

    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      # Optional: Run in GKE Sandbox
      # runtimeClassName: gvisor
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_TYPE
        cloud.google.com/gke-tpu-topology: TOPOLOGY
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS

    Ersetzen Sie Folgendes:

    • TPU_TYPE: Der zu verwendende TPU-Typ, z. B. tpu-v4-podslice. Muss ein von GKE unterstützter Wert sein.
    • TOPOLOGY: die Anordnung der TPU-Chips im Slice, z. B. 2x2x4. Muss eine unterstützte Topologie für den ausgewählten TPU-Typ sein.
    • NUMBER_OF_CHIPS: Die Anzahl der TPU-Chips, die für den Container verwendet werden sollen. Muss für limits und requests derselbe Wert sein.
    • MEMORY_SIZE: Die maximale Menge an Arbeitsspeicher, die von der TPU verwendet wird. Die Arbeitsspeicherlimits hängen von der verwendeten TPU-Version und ‑Topologie ab. Weitere Informationen finden Sie unter Mindest- und Höchstwerte für Beschleuniger.

    Optional können Sie auch die folgenden Felder ändern:

    • image: Das zu verwendende JAX AI-Image. Im Beispielmanifest ist dieses Feld auf das neueste JAX AI-Image festgelegt. Wenn Sie eine andere Version festlegen möchten, sehen Sie sich die Liste der aktuellen JAX AI-Images an.
    • runtimeClassname: gvisor: Die Einstellung, mit der Sie diesen Pod in GKE Sandbox ausführen können. Entfernen Sie die Kommentarzeichen für diese Zeile, um sie zu verwenden. GKE Sandbox unterstützt TPUs ab Version v4. Weitere Informationen finden Sie unter GKE Sandbox.

  2. Stellen Sie den Job bereit:

    kubectl create -f tpu-autopilot.yaml

    Wenn Sie diesen Job erstellen, führt GKE automatisch die folgenden Schritte aus:

    1. Stellt Knoten bereit, auf denen die Pods ausgeführt werden. Je nach TPU-Typ, Topologie und Ressourcenanforderungen, die Sie angegeben haben, sind diese Knoten entweder Slices mit einem Host oder Slices mit mehreren Hosts.
    2. Fügt den Pods Markierungen und den Knoten Toleranzen hinzu, um zu verhindern, dass andere Arbeitslasten auf denselben Knoten wie TPU-Arbeitslasten ausgeführt werden.

  3. Nach Abschluss dieses Abschnitts können Sie weitere Kosten vermeiden, indem Sie den erstellten Arbeitslast löschen:

    kubectl delete -f tpu-autopilot.yaml

Arbeitslast erstellen, die TPUs und die Planung von Sammlungen anfordert

In TPU Trillium können Sie die Planung von Sammlungen verwenden, um TPU-Slice-Knoten zu gruppieren. Durch das Gruppieren dieser TPU-Slice-Knoten lässt sich die Anzahl der Replikate leichter an die Arbeitslastanforderungen anpassen. Google Cloud steuert Softwareupdates, um dafür zu sorgen, dass immer genügend Slices in der Sammlung verfügbar sind, um Traffic zu verarbeiten.

TPU Trillium unterstützt die Planung von Sammlungen für Knotenpools mit einem einzelnen Host und mehreren Hosts, auf denen Inferenz-Arbeitslasten ausgeführt werden. Im Folgenden wird beschrieben, wie sich die Planung von Sammlungen je nach verwendetem TPU-Slice-Typ verhält:

  • TPU-Slice mit mehreren Hosts:GKE gruppiert TPU-Slices mit mehreren Hosts zu einer Sammlung. Jeder GKE-Knotenpool ist ein Replikat in dieser Sammlung. Wenn Sie eine Sammlung definieren möchten, erstellen Sie einen TPU-Slice mit mehreren Hosts und weisen Sie der Sammlung einen eindeutigen Namen zu. Wenn Sie der Sammlung weitere TPU-Slices hinzufügen möchten, erstellen Sie einen weiteren TPU-Slice-Knotenpool mit mehreren Hosts mit demselben Sammlungsnamen und Arbeitslasttyp.
  • TPU-Slice mit einem Host:GKE betrachtet den gesamten TPU-Slice-Knotenpool mit einem Host als Sammlung. Wenn Sie der Sammlung weitere TPU-Slices hinzufügen möchten, können Sie die Größe des TPU-Slice-Knotenpools mit einem Host ändern.

Informationen zu den Einschränkungen der Planung von Sammlungen finden Sie unter Funktionsweise der Planung von Sammlungen.

TPU-Slice mit mehreren Hosts verwenden

Die Planung der Erfassung in Knoten mit TPU-Slice mit mehreren Hosts ist für Autopilot-Cluster in Version 1.31.2-gke.1537000 und höher verfügbar. TPU-Slice-Knoten mit mehreren Hosts mit einer 2x4-Topologie werden nur in Version 1.31.2-gke.1115000 oder höher unterstützt. Wenn Sie TPU-Slice-Knoten mit mehreren Hosts erstellen und als Sammlung gruppieren möchten, fügen Sie Ihrer Arbeitslastspezifikation die folgenden Kubernetes-Labels hinzu:

  • cloud.google.com/gke-nodepool-group-name: Jede Sammlung muss auf Clusterebene einen eindeutigen Namen haben. Der Wert des Labels cloud.google.com/gke-nodepool-group-name muss den Anforderungen für Clusterlabels entsprechen.

  • cloud.google.com/gke-workload-type: HIGH_AVAILABILITY

    Der folgende Codeblock definiert beispielsweise eine Sammlung mit einem TPU-Slice mit mehreren Hosts:

      nodeSelector:
    cloud.google.com/gke-nodepool-group-name: ${COLLECTION_NAME}
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 4x4
...

TPU-Slice mit einem Host verwenden

Die Planung der Erfassung in Knoten mit Single-Host-TPU-Slice ist für Autopilot-Cluster in Version 1.31.2-gke.1088000 und höher verfügbar. Wenn Sie Knoten für einzelne Hosts für TPU-Slices erstellen und als Sammlung gruppieren möchten, fügen Sie das Label cloud.google.com/gke-workload-type:HIGH_AVAILABILITY in Ihre Arbeitslastspezifikation ein.

Der folgende Codeblock definiert beispielsweise eine Sammlung mit einem TPU-Slice mit einem einzelnen Host:

  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 2x2
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
  ...

Benutzerdefinierte Compute-Klassen zum Bereitstellen einer Sammlung verwenden

Weitere Informationen zum Bereitstellen einer Arbeitslast, die TPU-Arbeitslast- und Erfassungsplanung mit benutzerdefinierten Compute-Klassen anfordert, finden Sie unter TPU-Erfassung mit mehreren Hosts und Arbeitslasttyp für TPU-SLO definieren.

TPUs mit benutzerdefinierten Compute-Klassen zentral bereitstellen

So stellen Sie TPUs mit einer benutzerdefinierten Compute-Klasse bereit, die den TPU-Regeln entspricht, und stellen die Arbeitslast bereit:

  1. Speichern Sie das folgende Manifest als tpu-compute-class.yaml:

    apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: tpu-class
spec:
  priorities:
  - tpu:
      type: tpu-v5-lite-podslice
      count: 4
      topology: 2x4
  - spot: true
    tpu:
      type: tpu-v5-lite-podslice
      count: 4
      topology: 2x4
  - flexStart:
      enabled: true
    tpu:
      type: tpu-v6e-slice
      count: 4
      topology: 2x4
  nodePoolAutoCreation:
    enabled: true

  2. Stellen Sie die Compute-Klasse bereit:

    kubectl apply -f tpu-compute-class.yaml

    Weitere Informationen zu benutzerdefinierten Compute-Klassen und TPUs finden Sie unter TPU-Konfiguration.

  3. Speichern Sie das folgende Manifest als tpu-job.yaml:

    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/compute-class: tpu-class
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS

    Ersetzen Sie Folgendes:

    • NUMBER_OF_CHIPS: Die Anzahl der TPU-Chips, die für den Container verwendet werden sollen. Muss für limits und requests derselbe Wert sein, der dem Wert im Feld tpu.count in der ausgewählten benutzerdefinierten Compute-Klasse entspricht.
    • MEMORY_SIZE: Die maximale Menge an Speicher, die von der TPU verwendet wird. Die Arbeitsspeicherlimits hängen von der verwendeten TPU-Version und ‑Topologie ab. Weitere Informationen finden Sie unter Mindest- und Höchstwerte für Beschleuniger.
    • NUMBER_OF_CHIPS: Die Anzahl der TPU-Chips, die für den Container verwendet werden sollen. Muss für limits und requests derselbe Wert sein.

  4. Stellen Sie den Job bereit:

    kubectl create -f tpu-job.yaml

    Wenn Sie diesen Job erstellen, führt GKE automatisch die folgenden Schritte aus:

    • Stellt Knoten bereit, auf denen die Pods ausgeführt werden. Je nach TPU-Typ, Topologie und Ressourcenanforderungen, die Sie angegeben haben, sind diese Knoten entweder Slices mit einem Host oder Slices mit mehreren Hosts. Je nach Verfügbarkeit von TPU-Ressourcen mit der höchsten Priorität kann GKE auf niedrigere Prioritäten zurückgreifen, um die Verfügbarkeit zu maximieren.
    • Fügt den Pods Markierungen und den Knoten Toleranzen hinzu, um zu verhindern, dass andere Arbeitslasten auf denselben Knoten wie TPU-Arbeitslasten ausgeführt werden.

    Weitere Informationen zu benutzerdefinierten Compute-Klassen

  5. Nach Abschluss dieses Abschnitts können Sie weitere Kosten vermeiden, indem Sie die erstellten Ressourcen löschen:

    kubectl delete -f tpu-job.yaml

Beispiel: Gesamtzahl der TPU-Chips in einem Slice mit mehreren Hosts anzeigen

Die folgende Arbeitslast gibt die Anzahl der TPU-Chips für alle Knoten in einem TPU-Slice mit mehreren Hosts zurück. Zum Erstellen eines Slices mit mehreren Hosts hat die Arbeitslast die folgenden Parameter:

  • TPU-Version: TPU v4
  • Topologie: 2x2x4

Diese Versions- und Topologieauswahl führt zu einem Slice mit mehreren Hosts.

  1. Speichern Sie das folgende Manifest als available-chips-multihost.yaml: 
    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice # Node selector to target TPU v4 slice nodes.
        cloud.google.com/gke-tpu-topology: 2x2x4 # Specifies the physical topology for the TPU slice.
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())' # Python command to count available TPU chips.
        resources:
          requests:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4 # Request 4 TPU chips for this workload.
          limits:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4 # Limit to 4 TPU chips for this workload.
  2. Stellen Sie das Manifest bereit: 
    kubectl create -f available-chips-multihost.yaml

    GKE führt ein TPU v4-Slice mit vier VMs aus (TPU-Slice mit mehreren Hosts). Das Slice hat 16 miteinander verbundene TPU-Chips.

  3. Prüfen Sie, ob der Job vier Pods erstellt hat: 
    kubectl get pods

    Die Ausgabe sieht in etwa so aus:

    NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
  4. Rufen Sie die Logs eines der Pods ab: 
    kubectl logs POD_NAME

    Ersetzen Sie POD_NAME durch den Namen eines der erstellten Pods. Beispiel: tpu-job-podslice-0-5cd8r.

    Die Ausgabe sieht etwa so aus: 

    TPU cores: 16
  5. Optional: Arbeitslast entfernen: 
    kubectl delete -f available-chips-multihost.yaml

Beispiel: TPU-Chips in einem einzelnen Knoten anzeigen

Die folgende Arbeitslast ist ein statischer Pod, der die Anzahl der TPU-Chips anzeigt, die mit einem bestimmten Knoten verbunden sind. Zum Erstellen eines Knotens mit einem einzelnen Host hat die Arbeitslast folgende Parameter:

  • TPU-Version: TPU v5e
  • Topologie: 2x4

Diese Versions- und Topologieauswahl führt zu einem Slice mit einem einzelnen Host.

  1. Speichern Sie das folgende Manifest als available-chips-singlehost.yaml: 
    apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice # Node selector to target TPU v5e slice nodes.
    cloud.google.com/gke-tpu-topology: 2x4 # Specify the physical topology for the TPU slice.
  containers:
  - name: tpu-job
    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8 # Request 8 TPU chips for this container.
      limits:
        google.com/tpu: 8 # Limit to 8 TPU chips for this container.
  2. Stellen Sie das Manifest bereit: 
    kubectl create -f available-chips-singlehost.yaml

    GKE stellt Knoten mit acht TPU-Slices mit einem einzelnen Host bereit, die TPU v5e verwenden. Jeder TPU-Knoten hat acht TPU-Chips (TPU-Slice mit einzelnem Host).

  3. Rufen Sie die Logs des Pods ab: 
    kubectl logs tpu-job-jax-v5

    Die Ausgabe sieht etwa so aus: 

    Total TPU chips: 8
  4. Optional: Arbeitslast entfernen: 
      kubectl delete -f available-chips-singlehost.yaml

TPUs beobachten und überwachen

Dashboard

Die Beobachtbarkeit von Knotenpools in der Google Cloud Console ist allgemein verfügbar. Wenn Sie den Status Ihrer TPU-Knotenpools mit mehreren Hosts in GKE aufrufen möchten, rufen Sie das von Cloud Monitoring bereitgestellte Dashboard GKE TPU Node Pool Status auf:

Zum GKE TPU-Knotenpoolstatus

Dieses Dashboard bietet umfassende Informationen zum Zustand Ihrer TPU-Knotenpools mit mehreren Hosts. Weitere Informationen finden Sie unter Zustandsmesswerte für TPU-Knoten und ‑Knotenpools überwachen.

Auf der Seite Kubernetes-Cluster in derGoogle Cloud Console werden auf dem Tab Beobachtbarkeit auch TPU-Beobachtbarkeitsmesswerte wie die TPU-Nutzung unter der Überschrift Beschleuniger > TPU angezeigt. Weitere Informationen finden Sie unter Beobachtbarkeitsmesswerte aufrufen.

Das TPU-Dashboard wird nur ausgefüllt, wenn in Ihrem GKE-Cluster Systemmesswerte aktiviert sind.

Laufzeitmesswerte

In GKE-Version 1.27.4-gke.900 oder höher exportieren TPU-Arbeitslasten, die sowohl JAX-Version 0.4.14 oder höher verwenden als auch containerPort: 8431 angeben, TPU-Auslastungsmesswerte als GKE-Systemmesswerte. Die folgenden Messwerte sind in Cloud Monitoring verfügbar, um die Laufzeitleistung Ihrer TPU-Arbeitslast zu überwachen:

  • Arbeitszyklus: Prozentsatz der Zeit im vergangenen Stichprobenzeitraum (60 Sekunden), während der TensorCores auf einem TPU-Chip aktiv mit Verarbeitung beschäftigt waren. Ein höherer Prozentsatz bedeutet eine bessere TPU-Auslastung.
  • Arbeitsspeichernutzung: Menge des dem Beschleuniger zugeteilten Arbeitsspeichers in Byte. Alle 60 Sekunden wird eine Stichprobe erstellt.
  • Speicherkapazität: Gesamter Arbeitsspeicher des Beschleunigers in Byte. Alle 60 Sekunden wird eine Stichprobe erstellt.

Diese Messwerte befinden sich im Kubernetes-Knotenschema (k8s_node) und im Kubernetes-Containerschema (k8s_container).

Kubernetes-Container:

  • kubernetes.io/container/accelerator/duty_cycle
  • kubernetes.io/container/accelerator/memory_used
  • kubernetes.io/container/accelerator/memory_total

Kubernetes-Knoten:

  • kubernetes.io/node/accelerator/duty_cycle
  • kubernetes.io/node/accelerator/memory_used
  • kubernetes.io/node/accelerator/memory_total

Messwerte zum Zustand von TPU-Knoten und ‑Knotenpools überwachen

Wenn bei einem Trainingsjob ein Fehler auftritt oder er fehlschlägt, können Sie die Messwerte der zugrunde liegenden Infrastruktur prüfen, um festzustellen, ob die Unterbrechung durch ein Problem mit dem zugrunde liegenden Knoten oder Knotenpool verursacht wurde.

Knotenstatus

In GKE-Version 1.32.1-gke.1357001 oder höher wird der Zustand eines GKE-Knotens durch den folgenden GKE-Systemmesswert dargestellt:

  • kubernetes.io/node/status_condition

Im Feld condition werden Bedingungen für den Knoten gemeldet, z. B. Ready, DiskPressure und MemoryPressure. Im Feld status wird der gemeldete Status der Bedingung angezeigt, der True, False oder Unknown sein kann. Dies ist ein Messwert mit dem überwachten Ressourcentyp k8s_node.

Diese PromQL-Abfrage zeigt, ob ein bestimmter Knoten Ready ist:

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    node_name="NODE_NAME",
    condition="Ready",
    status="True"}

Um Probleme in einem Cluster zu beheben, können Sie sich Knoten ansehen, die andere Bedingungen aufweisen:

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition!="Ready",
    status="True"}

Möglicherweise möchten Sie sich speziell Knoten ansehen, die nicht Ready sind:

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition="Ready",
    status="False"}

Wenn keine Daten vorhanden sind, sind die Knoten bereit. Der Status wird alle 60 Sekunden erfasst.

Mit der folgenden Abfrage können Sie den Knotenstatus in der gesamten Flotte ermitteln:

avg by (condition,status)(
  avg_over_time(
    kubernetes_io:node_status_condition{monitored_resource="k8s_node"}[${__interval}]))

Knotenpoolstatus

Der folgende GKE-Systemmesswert für die überwachte Ressource k8s_node_pool gibt den Status eines GKE-Knotenpools an:

  • kubernetes.io/node_pool/status

Dieser Messwert wird nur für Multi-Host-TPU-Knotenpools angegeben.

Im Feld status wird der Status des Knotenpools angegeben, z. B. Provisioning, Running, Error, Reconciling oder Stopping. Statusaktualisierungen erfolgen nach Abschluss der GKE API-Vorgänge.

Mit der folgenden PromQL-Abfrage können Sie prüfen, ob ein bestimmter Knotenpool den Status Running hat:

kubernetes_io:node_pool_status{
    monitored_resource="k8s_node_pool",
    cluster_name="CLUSTER_NAME",
    node_pool_name="NODE_POOL_NAME",
    status="Running"}

Verwenden Sie die folgende PromQL-Abfrage, um die Anzahl der Knotenpools in Ihrem Projekt nach Status gruppiert zu überwachen:

count by (status)(
  count_over_time(
    kubernetes_io:node_pool_status{monitored_resource="k8s_node_pool"}[${__interval}]))

Knotenpoolverfügbarkeit

Der folgende GKE-Systemmesswert gibt an, ob ein TPU-Knotenpool mit mehreren Hosts verfügbar ist:

  • kubernetes.io/node_pool/multi_host/available

Der Messwert hat den Wert True, wenn alle Knoten im Knotenpool verfügbar sind, andernfalls False. Der Messwert wird alle 60 Sekunden abgerufen.

Verwenden Sie die folgende PromQL-Abfrage, um die Verfügbarkeit von TPU-Knotenpools mit mehreren Hosts in Ihrem Projekt zu prüfen:

avg by (node_pool_name)(
  avg_over_time(
    kubernetes_io:node_pool_multi_host_available{
      monitored_resource="k8s_node_pool",
      cluster_name="CLUSTER_NAME"}[${__interval}]))

Anzahl der Knotenunterbrechungen

Der folgende GKE-Systemmesswert gibt die Anzahl der Unterbrechungen für einen GKE-Knoten seit der letzten Stichprobe an (der Messwert wird alle 60 Sekunden abgerufen):

  • kubernetes.io/node/interruption_count

Die Felder interruption_type (z. B. TerminationEvent, MaintenanceEvent oder PreemptionEvent) und interruption_reason (z. B. HostError, Eviction oder AutoRepair) können Aufschluss darüber geben, warum ein Knoten unterbrochen wurde.

Wenn Sie eine Aufschlüsselung der Unterbrechungen und ihrer Ursachen in TPU-Knoten in den Clustern in Ihrem Projekt erhalten möchten, verwenden Sie die folgende PromQL-Abfrage:

  sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node"}[${__interval}]))

Wenn Sie nur die Hostwartungsereignisse sehen möchten, aktualisieren Sie die Abfrage, um den HW/SW Maintenance-Wert für interruption_reason zu filtern. Verwenden Sie die folgende PromQL-Abfrage:

  sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node", interruption_reason="HW/SW Maintenance"}[${__interval}]))

Verwenden Sie die folgende PromQL-Abfrage, um die Anzahl der Unterbrechungen nach Knotenpool zusammenzufassen:

  sum by (node_pool_name,interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_pool_interruption_count{monitored_resource="k8s_node_pool", interruption_reason="HW/SW Maintenance", node_pool_name=NODE_POOL_NAME }[${__interval}]))

Wiederherstellungszeiten (Time To Recover, TTR) des Knotenpools

Der folgende GKE-Systemmesswert gibt die Verteilung der Dauer des Wiederherstellungszeitraums für GKE-TPU-Knotenpools mit mehreren Hosts an:

  • kubernetes.io/node_pool/accelerator/times_to_recover

Jede in diesem Messwert aufgezeichnete Stichprobe gibt ein einzelnes Wiederherstellungsereignis für den Knotenpool nach einer Ausfallzeit an.

Dieser Messwert ist nützlich, um die Wiederherstellungszeit und die Zeit zwischen Unterbrechungen von Multi-Host-TPU-Knotenpools zu verfolgen.

Mit der folgenden PromQL-Abfrage können Sie die mittlere Zeit bis zur Wiederherstellung (Mean Time to Recovery, MTTR) für die letzten 7 Tage in Ihrem Cluster berechnen:

sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_sum{
    monitored_resource="k8s_node_pool", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_count{
    monitored_resource="k8s_node_pool",cluster_name="CLUSTER_NAME"}[7d]))

Zeiten zwischen Unterbrechungen (Times Between Interruptions, TBI) des Knotenpools

Die Zeit zwischen Unterbrechungen des Knotenpools gibt an, wie lange Ihre Infrastruktur läuft, bevor es zu einer Unterbrechung kommt. Er wird als Durchschnitt über einen bestimmten Zeitraum berechnet. Der Zähler gibt die Gesamtzeit an, in der Ihre Infrastruktur verfügbar war, und der Nenner die Gesamtzahl der Unterbrechungen Ihrer Infrastruktur.

Das folgende PromQL-Beispiel zeigt die durchschnittliche Zeit zwischen Unterbrechungen (Mean Time Between Interruptions, MTBI) für den angegebenen Cluster über einen Zeitraum von sieben Tagen:

sum(count_over_time(
  kubernetes_io:node_memory_total_bytes{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_interruption_count{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))

Hostmesswerte

In GKE-Version 1.28.1-gke.1066000 oder höher exportieren VMs in einem TPU-Slice TPU-Auslastungsmesswerte als GKE-Systemmesswerte. Die folgenden Messwerte sind in Cloud Monitoring verfügbar, um die Leistung Ihres TPU-Hosts zu überwachen:

  • TensorCore-Auslastung: Aktueller Prozentsatz des verwendeten TensorCore. Der TensorCore-Wert entspricht der Summe der Matrixmultiplikationseinheiten (MXUs) plus der Vektoreinheit. Der TensorCore-Auslastungswert ist der Quotient aus den TensorCore-Vorgängen, die im letzten Stichprobenzeitraum (60 Sekunden) ausgeführt wurden, und der unterstützten Anzahl von TensorCore-Vorgängen im selben Zeitraum. Ein größerer Wert bedeutet eine bessere Auslastung.
  • Auslastung der Arbeitsspeicherbandbreite: Aktueller Prozentsatz der verwendeten Arbeitsspeicherbandbreite des Beschleunigers. Wird berechnet, indem die in einem Stichprobenzeitraum (60 s) verwendete Arbeitsspeicherbandbreite durch die maximal unterstützte Bandbreite im selben Stichprobenzeitraum geteilt wird.

Diese Messwerte befinden sich im Kubernetes-Knotenschema (k8s_node) und im Kubernetes-Containerschema (k8s_container).

Kubernetes-Container:

  • kubernetes.io/container/accelerator/tensorcore_utilization
  • kubernetes.io/container/accelerator/memory_bandwidth_utilization

Kubernetes-Knoten:

  • kubernetes.io/node/accelerator/tensorcore_utilization
  • kubernetes.io/node/accelerator/memory_bandwidth_utilization

Weitere Informationen finden Sie unter Kubernetes-Messwerte und GKE-Systemmesswerte.

Logging

Logs, die von Containern ausgegeben werden, die auf GKE-Knoten ausgeführt werden, einschließlich TPU-VMs, werden vom GKE-Logging-Agent erfasst und an Logging gesendet und sind in Logging sichtbar.

Empfehlungen für TPU-Arbeitslasten in Autopilot

Die folgenden Empfehlungen können die Effizienz Ihrer TPU-Arbeitslasten verbessern:

  • Verwenden Sie Pods mit verlängerter Laufzeit für einen Kulanzzeitraum von bis zu sieben Tagen, bevor GKE Ihre Pods wegen Herunterskalierungen oder Knotenupgrades beendet. Sie können Wartungsfenster und ‑ausschlüsse mit Pods mit langer Laufzeit verwenden, um automatische Knotenupgrades weiter zu verzögern.
  • Verwenden Sie Kapazitätsreservierungen, damit Ihre Arbeitslasten angeforderte TPUs erhalten, ohne in Verfügbarkeits-Warteschlangen gestellt zu werden.

Informationen zum Einrichten von Cloud TPU in GKE finden Sie in den folgenden Google Cloud Ressourcen: