Esegui una pipeline con le TPU

Questa pagina spiega come eseguire una pipeline Apache Beam su Dataflow con le TPU. I job che utilizzano le TPU comportano addebiti come specificato nella pagina dei prezzi di Dataflow.

Per ulteriori informazioni sull'utilizzo delle TPU con Dataflow, consulta Supporto di Dataflow per le TPU.

(Facoltativo) Effettuare una prenotazione specifica per utilizzare gli acceleratori

Sebbene tu possa utilizzare le TPU on demand, ti consigliamo vivamente di utilizzare le TPU Dataflow con riservazioni Google Cloud specificamente mirate. In questo modo, hai accesso agli acceleratori disponibili e i tempi di avvio dei worker sono rapidi. Le pipeline che utilizzano una prenotazione TPU non richiedono una quota TPU aggiuntiva.

Se non effettui una prenotazione e scegli di utilizzare le TPU on demand, esegui il provisioning della quota TPU prima di eseguire la pipeline.

(Facoltativo) Esegui il provisioning della quota TPU

Puoi utilizzare le TPU con una capacità on demand o utilizzando una prenotazione. Se vuoi utilizzare le TPU on demand, devi eseguire il provisioning della quota TPU prima di farlo. Se utilizzi una prenotazione con target specifico, puoi saltare questa sezione.

Per utilizzare le TPU on demand senza prenotazione, controlla il limite e l'utilizzo attuale della quota API Compute Engine per le TPU nel seguente modo:

Console

  1. Vai alla pagina Quote nella console Google Cloud :

    Vai a Quote

  2. Nella casella Filtro , procedi nel seguente modo:

    1. Utilizza la seguente tabella per selezionare e copiare la proprietà della quota in base alla versione della TPU e al tipo di macchina. Ad esempio, se prevedi di creare nodi TPU v5e on demand il cui tipo di macchina inizia con ct5lp-, inserisci Name: TPU v5 Lite PodSlice chips.

      Versione TPU, tipo di macchina inizia con Proprietà e nome della quota per le istanze on demand
      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. Seleziona la proprietà Dimensioni (ad es. località) e inserisci region: seguito dal nome della regione in cui prevedi di avviare la pipeline. Ad esempio, inserisci region:us-west4 se prevedi di utilizzare la zona us-west4-a. La quota TPU è regionale, quindi tutte le zone all'interno della stessa regione consumano la stessa quota TPU.

Configurare un'immagine container personalizzata

Per interagire con le TPU nelle pipeline Dataflow, devi fornire software in grado di operare su dispositivi XLA nell'ambiente di runtime della pipeline. Ciò richiede l'installazione di librerie TPU in base alle esigenze della pipeline e la configurazione delle variabili di ambiente in base al dispositivo TPU che utilizzi.

Per personalizzare l'immagine container, installa Apache Beam in un'immagine di base pronta all'uso che contiene le librerie TPU necessarie. In alternativa, installa il software TPU nelle immagini pubblicate con le release dell'SDK Apache Beam.

Per fornire un'immagine container personalizzata, utilizza l'opzione della pipeline sdk_container_image. Per saperne di più, consulta Utilizzare container personalizzati in Dataflow.

Quando utilizzi un acceleratore TPU, devi impostare le seguenti variabili di ambiente nell'immagine container.

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

A seconda dell'acceleratore utilizzato, è necessario impostare anche le variabili nella tabella seguente.

tipo topologia Dataflow richiesto worker_machine_type variabili di ambiente aggiuntive
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

Un Dockerfile di esempio per l'immagine container personalizzata potrebbe avere il seguente aspetto:

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"]

Esegui il job con le TPU

Le considerazioni per l'esecuzione di un job Dataflow con le TPU includono quanto segue:

  • Poiché i container TPU possono essere di grandi dimensioni, per evitare di esaurire lo spazio su disco, aumenta la dimensione predefinita del disco di avvio a 50 gigabyte o a una dimensione appropriata come richiesto dall'immagine container utilizzando l'opzione della pipeline --disk_size_gb.
  • Limita il parallelismo intra-worker.

TPU e parallelismo dei worker

Nella configurazione predefinita, le pipeline Dataflow Python avviano un processo SDK Apache Beam per core VM. I tipi di macchine TPU hanno un numero elevato di core vCPU, ma solo un processo può eseguire calcoli su un dispositivo TPU. Inoltre, un dispositivo TPU potrebbe essere riservato da un processo per l'intera durata del processo. Pertanto, devi limitare il parallelismo intra-worker quando esegui una pipeline Dataflow TPU. Per limitare il parallelismo dei worker, segui queste indicazioni:

  • Se il tuo caso d'uso prevede l'esecuzione di inferenze su un modello, utilizza l'API Beam RunInference. Per ulteriori informazioni, consulta Inferenza del modello linguistico di grandi dimensioni in Beam.
  • Se non puoi utilizzare l'API Beam RunInference, utilizza gli oggetti condivisi multiprocesso di Beam per limitare determinate operazioni a un singolo processo.
  • Se non puoi utilizzare i suggerimenti precedenti e preferisci avviare un solo processo Python per worker, imposta l'opzione della pipeline --experiments=no_use_multiple_sdk_containers.
  • Per i worker con più di 100 vCPU, riduci il numero di thread utilizzando l'opzione della pipeline --number_of_worker_harness_threads. Utilizza la seguente tabella per verificare se il tuo tipo di TPU utilizza più di 100 vCPU.

La tabella seguente elenca le risorse di calcolo totali per worker per ogni configurazione TPU.

Tipo di TPU topologia tipo di macchina Chip TPU 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

Esegui una pipeline con le TPU

Per eseguire un job Dataflow con le TPU, utilizza il comando seguente.

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

Sostituisci quanto segue:

  • PIPELINE: il file del codice sorgente della pipeline.
  • PROJECT: Il Google Cloud nome del progetto.
  • BUCKET: il bucket Cloud Storage.
  • REGION: una regione Dataflow, ad esempio us-central1.
  • TPU_TYPE: un tipo di TPU supportato, ad esempio tpu-v5-lite-podslice. Per un elenco completo di tipi e topologie, consulta Acceleratori TPU supportati.
  • TPU_TOPOLOGY: la topologia TPU, ad esempio 1x1.
  • MACHINE_TYPE: il tipo di macchina corrispondente, ad esempio ct5lp-hightpu-1t.
  • DISK_SIZE_GB: le dimensioni del disco di avvio per ogni VM worker, ad esempio 100.
  • IMAGE: il percorso Artifact Registry per l'immagine Docker.
  • NUMBER_OF_THREADS: (Facoltativo) Il numero di thread del worker harness.

Verifica il job Dataflow

Per verificare che il job utilizzi VM worker con TPU:

  1. Nella console Google Cloud , vai alla pagina Dataflow > Job.

    Vai a Job

  2. Seleziona un lavoro.

  3. Fai clic sulla scheda Metriche dei job.

  4. Nella sezione Scalabilità automatica, verifica che esista almeno una VM Current workers.

  5. Nel riquadro laterale Informazioni sul job, verifica che machine_type inizi con ct. Ad esempio: ct6e-standard-1t. Indica l'utilizzo della TPU.

Risolvere i problemi relativi al job Dataflow

Se riscontri problemi durante l'esecuzione del job Dataflow con le TPU, consulta Risolvere i problemi relativi al job Dataflow TPU.

Passaggi successivi