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 auf der Preisseite für Dataflow 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 erstellen

Sie können TPUs zwar bei Bedarf verwenden, wir empfehlen jedoch dringend, Dataflow-TPUs mit speziell ausgerichteten Google CloudReservierungen 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 On-Demand verwenden möchten, müssen Sie TPU-Kontingent bereitstellen, bevor Sie Ihre Pipeline ausführen.

Optional: TPU-Kontingent bereitstellen

Sie können TPUs mit On-Demand-Kapazität oder mit einer Reservierung verwenden. Wenn Sie TPUs bei Bedarf verwenden möchten, müssen Sie vorab ein TPU-Kontingent bereitstellen. Wenn Sie eine speziell ausgerichtete Reservierung verwenden, können Sie diesen Abschnitt überspringen.

Wenn Sie On-Demand-TPUs ohne Reservierung verwenden möchten, prüfen Sie das Limit und die aktuelle Nutzung Ihres Compute Engine API-Kontingents für TPUs so:

Console

  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 Maschinentyp auszuwählen und zu kopieren. 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 Property 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 das Attribut Dimensionen (z.B. Standort) aus und geben Sie region: gefolgt vom Namen der Region ein, in der Sie Ihre 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 in Dataflow-Pipelines mit TPUs interagieren möchten, müssen Sie Software bereitstellen, die in Ihrer Pipeline-Laufzeitumgebung auf XLA-Geräten ausgeführt werden kann. Dazu müssen Sie TPU-Bibliotheken basierend auf den Anforderungen Ihrer 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 die mit Apache Beam SDK-Releases veröffentlichten Images installieren.

Wenn Sie ein benutzerdefiniertes Container-Image bereitstellen möchten, verwenden Sie die Pipeline-Option sdk_container_image. 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 könnte so aussehen:

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 mit der Pipeline-Option --disk_size_gb auf 50 GB oder eine andere geeignete Größe erhöhen, die für Ihr Container-Image erforderlich ist, um zu vermeiden, dass der Speicherplatz aufgebraucht wird.
  • Beschränken Sie die Parallelität innerhalb von Workern.

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 von Workern einschränken, wenn Sie eine Dataflow-TPU-Pipeline ausführen. So beschränken Sie die Worker-Parallelität:

  • Wenn Sie Inferenzvorgänge für ein Modell ausführen möchten, verwenden Sie die Beam-RunInference-API. Weitere Informationen finden Sie unter LLM-Inferenz in Beam.
  • Wenn Sie die Beam RunInference API nicht verwenden können, verwenden Sie die gemeinsamen Objekte von Beam für mehrere Prozesse, um bestimmte Vorgänge auf einen einzelnen Prozess zu beschränken.
  • Wenn Sie die oben genannten Empfehlungen nicht verwenden können und nur einen Python-Prozess pro Worker starten möchten, legen Sie die Pipeline-Option --experiments=no_use_multiple_sdk_containers fest.
  • Bei Workern mit mehr als 100 vCPUs können Sie die Anzahl der Threads mit der Pipelineoption --number_of_worker_harness_threads reduzieren. In der folgenden Tabelle sehen Sie, ob Ihr TPU-Typ mehr als 100 vCPUs verwendet.

In der folgenden Tabelle sind die gesamten Rechenressourcen 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 Name des Google Cloud Projekts.
  • 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 Ihr 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 VM für Aktuelle Worker vorhanden ist.

  5. Prüfen Sie in der Seitenleiste Jobinfo, ob machine_type mit ct beginnt. Beispiel: ct6e-standard-1t Dies gibt die TPU-Nutzung an.

Fehlerbehebung bei Dataflow-Jobs

Wenn bei der Ausführung Ihres Dataflow-Jobs mit TPUs Probleme auftreten, lesen Sie die Informationen unter Fehlerbehebung bei Dataflow-TPU-Jobs.

Nächste Schritte