Pipeline mit TPUs ausführen

Auf dieser Seite wird erläutert, wie eine Apache Beam-Pipeline in Dataflow mit TPUs ausgeführt wird. Für Jobs, die TPUs verwenden, fallen Gebühren an, wie in der Dataflow-Preisseite beschrieben .

Weitere Informationen zur Verwendung von TPUs mit Dataflow finden Sie unter Dataflow-Unterstützung für TPUs.

Optional: Bestimmte Reservierung für die Verwendung von Beschleunigern vornehmen

Sie können TPUs zwar nach Bedarf verwenden, wir empfehlen jedoch dringend, Dataflow-TPUs mit speziell ausgerichteten Google Cloud Reservierungen zu verwenden. So haben Sie Zugriff auf verfügbare Beschleuniger und schnelle Worker-Startzeiten. Für Pipelines, die eine TPU-Reservierung nutzen, ist kein zusätzliches TPU-Kontingent erforderlich.

Wenn Sie keine Reservierung vornehmen und TPUs nach Bedarf verwenden möchten, stellen Sie das TPU-Kontingent bereit, bevor Sie die Pipeline ausführen.

Optional: TPU-Kontingent bereitstellen

Sie können TPUs nach Bedarf oder mit einer Reservierung verwenden. Wenn Sie TPUs nach Bedarf verwenden möchten, müssen Sie das TPU-Kontingent vorher bereitstellen. Wenn Sie eine speziell ausgerichtete Reservierung verwenden, können Sie diesen Abschnitt überspringen.

So prüfen Sie das Limit und die aktuelle Nutzung Ihres Compute Engine API-Kontingents für TPUs, wenn Sie TPUs nach Bedarf ohne Reservierung verwenden möchten:

Console

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

    Kontingente aufrufen

  2. Gehen Sie im Feld Filter so vor:

    1. Wählen Sie in der folgenden Tabelle die Eigenschaft des Kontingents basierend auf der TPU-Version und dem Maschinentyp aus und kopieren Sie sie. Wenn Sie beispielsweise On-Demand-TPU v5e-Knoten erstellen möchten, deren Maschinentyp mit ct5lp- beginnt, geben Sie Name: TPU v5 Lite PodSlice chips ein.

      TPU-Version, Maschinentyp beginnt mit Eigenschaft und Name des Kontingents für On-Demand-Instanzen
      TPU v5e,
      ct5lp- 
      Name:
      TPU v5 Lite PodSlice chips
      TPU v5p,
      ct5p- 
      Name:
      TPU v5p chips
      TPU v6e,
      ct6e- 
      Dimensions (e.g. location):
      tpu_family:CT6E

    2. Wählen Sie die Eigenschaft Dimensionen (z.B. Standort) aus und geben Sie region: gefolgt vom Namen der Region ein, in der Sie die Pipeline starten möchten. Geben Sie beispielsweise region:us-west4 ein, wenn Sie die Zone us-west4-a verwenden möchten. Das TPU-Kontingent ist regional, d. h. alle Zonen innerhalb derselben Region nutzen dasselbe TPU-Kontingent.

Benutzerdefiniertes Container-Image konfigurieren

Wenn Sie mit TPUs in Dataflow-Pipelines interagieren möchten, müssen Sie Software bereitstellen, die in der Pipeline-Laufzeitumgebung auf XLA-Geräten ausgeführt werden kann. Dazu müssen Sie TPU-Bibliotheken basierend auf den Anforderungen der Pipeline installieren und Umgebungsvariablen basierend auf dem verwendeten TPU-Gerät konfigurieren.

Wenn Sie das Container-Image anpassen möchten, installieren Sie Apache Beam in einem Standard-Basis-Image , das die erforderlichen TPU-Bibliotheken enthält. Alternativ können Sie die TPU-Software in den Images installieren, die mit Apache Beam SDK-Releases veröffentlicht wurden.

Verwenden Sie die Pipeline-Option sdk_container_image, um ein benutzerdefiniertes Container-Image bereitzustellen. Weitere Informationen finden Sie unter Benutzerdefinierte Container in Dataflow verwenden.

Wenn Sie einen TPU-Beschleuniger verwenden, müssen Sie die folgenden Umgebungsvariablen im Container-Image festlegen.

ENV TPU_SKIP_MDS_QUERY=1 # Don't query metadata
ENV TPU_HOST_BOUNDS=1,1,1 # There's only one host
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0 # Always 0 for single-host TPUs

Je nach verwendetem Beschleuniger müssen auch die Variablen in der folgenden Tabelle festgelegt werden.

Typ Topologie Erforderlicher Dataflow-worker_machine_type Zusätzliche Umgebungsvariablen
tpu-v5-lite-podslice 1x1 ct5lp-hightpu-1t 
TPU_ACCELERATOR_TYPE=v5litepod-1
TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v5-lite-podslice 2x2 ct5lp-hightpu-4t 
TPU_ACCELERATOR_TYPE=v5litepod-4
TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v5-lite-podslice 2x4 ct5lp-hightpu-8t 
TPU_ACCELERATOR_TYPE=v5litepod-8
TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v6e-slice 1x1 ct6e-standard-1t 
TPU_ACCELERATOR_TYPE=v6e-1
TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v6e-slice 2x2 ct6e-standard-4t 
TPU_ACCELERATOR_TYPE=v6e-4
TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v6e-slice 2x4 ct6e-standard-8t 
TPU_ACCELERATOR_TYPE=v6e-8
TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v5p-slice 2x2x1 ct5p-hightpu-4t 
TPU_ACCELERATOR_TYPE=v5p-8
TPU_CHIPS_PER_HOST_BOUNDS=2,2,1

Ein Beispiel für ein Dockerfile für das benutzerdefinierte Container-Image:

FROM python:3.11-slim

COPY --from=apache/beam_python3.11_sdk:2.66.0 /opt/apache/beam /opt/apache/beam

# Configure the environment to access TPU device

ENV TPU_SKIP_MDS_QUERY=1
ENV TPU_HOST_BOUNDS=1,1,1
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0

# Configure the environment for the chosen accelerator.
# Adjust according to the accelerator you use.
ENV TPU_ACCELERATOR_TYPE=v5litepod-1
ENV TPU_CHIPS_PER_HOST_BOUNDS=1,1,1

# Install TPU software stack.
RUN pip install jax[tpu] apache-beam[gcp]==2.66.0 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

ENTRYPOINT ["/opt/apache/beam/boot"]

Job mit TPUs ausführen

Folgendes ist für die Ausführung eines Dataflow-Jobs mit TPUs zu beachten:

  • Da TPU-Container groß sein können, sollten Sie die Standardgröße des Bootlaufwerks auf mindestens 50 Gigabyte oder eine geeignete Größe erhöhen, die für Ihr Container-Image erforderlich ist. Verwenden Sie dazu die Pipeline-Option --disk_size_gb.
  • Beschränken Sie die Parallelität innerhalb des Workers.

TPUs und Parallelität der Worker

In der Standardkonfiguration starten Dataflow-Python-Pipelines einen Apache Beam SDK-Prozess pro VM Kern. TPU-Maschinentypen haben eine große Anzahl von vCPU-Kernen, aber nur ein Prozess kann Berechnungen auf einem TPU-Gerät ausführen. Außerdem kann ein TPU-Gerät von einem Prozess für die gesamte Lebensdauer des Prozesses reserviert werden. Daher müssen Sie die Parallelität innerhalb des Workers einschränken, wenn Sie eine Dataflow-TPU-Pipeline ausführen. So beschränken Sie die Parallelität der Worker:

In der folgenden Tabelle sind die gesamten Compute-Ressourcen pro Worker für jede TPU-Konfiguration aufgeführt.

TPU-Typ Topologie Maschinentyp TPU-Chips vCPU RAM (GB)
tpu-v5-lite-podslice 1x1 ct5lp-hightpu-1t 1 24 48
tpu-v5-lite-podslice 2x2 ct5lp-hightpu-4t 4 112 192
tpu-v5-lite-podslice 2x4 ct5lp-hightpu-8t 8 224 384
tpu-v6e-slice 1x1 ct6e-standard-1t 1 44 176
tpu-v6e-slice 2x2 ct6e-standard-4t 4 180 720
tpu-v6e-slice 2x4 ct6e-standard-8t 8 360 1440
tpu-v5p-slice 2x2x1 ct5p-hightpu-4t 4 208 448

Pipeline mit TPUs ausführen

Führen Sie den folgenden Befehl aus, um einen Dataflow-Job mit TPUs auszuführen.

python PIPELINE \
  --runner "DataflowRunner" \
  --project "PROJECT" \
  --temp_location "gs://BUCKET/tmp" \
  --region "REGION" \
  --dataflow_service_options "worker_accelerator=type:TPU_TYPE;topology:TPU_TOPOLOGY" \
  --worker_machine_type "MACHINE_TYPE" \
  --disk_size_gb "DISK_SIZE_GB" \
  --sdk_container_image "IMAGE" \
  --number_of_worker_harness_threads NUMBER_OF_THREADS

Ersetzen Sie Folgendes:

  • PIPELINE: Ihre Quellcode-Pipelinedatei
  • PROJECT: Der Google Cloud Projektname.
  • BUCKET: Der Cloud Storage-Bucket
  • REGION: Eine Dataflow-Region, z. B. us-central1
  • TPU_TYPE: Ein unterstützter TPU-Typ, z. B. tpu-v5-lite-podslice Eine vollständige Liste der Typen und Topologien finden Sie unter Unterstützte TPU-Beschleuniger.
  • TPU_TOPOLOGY: Die TPU-Topologie, z. B. 1x1
  • MACHINE_TYPE: Der entsprechende Maschinentyp, z. B. ct5lp-hightpu-1t.
  • DISK_SIZE_GB: Die Größe des Bootlaufwerks für jede Worker-VM, z. B. 100.
  • IMAGE: Der Artifact Registry-Pfad für das Docker-Image
  • NUMBER_OF_THREADS: Optional. Die Anzahl der Worker-Harness-Threads.

Dataflow-Job prüfen

So prüfen Sie, ob der Job Worker-VMs mit TPUs verwendet:

  1. Rufen Sie in der Google Cloud Console die Seite Dataflow > Jobs auf.

    Zu Jobs

  2. Wählen Sie einen Job aus.

  3. Klicken Sie auf den Tab Jobmesswerte.

  4. Prüfen Sie im Bereich Autoscaling, ob mindestens eine Aktuelle Worker-VM vorhanden ist.

  5. Prüfen Sie im Seitenbereich Jobinformationen, ob machine_type mit ct beginnt. Beispiel: ct6e-standard-1t. Dies weist auf die TPU-Nutzung hin.

Fehlerbehebung bei Dataflow-Jobs

Wenn bei der Ausführung Ihres Dataflow-Jobs mit TPUs Probleme auftreten, siehe Fehlerbehebung bei Dataflow-TPU Jobs.

Nächste Schritte