Execute um pipeline com TPUs

Esta página explica como executar um pipeline do Apache Beam no Dataflow com TPUs. As tarefas que usam TPUs incorrem em cobranças, conforme especificado na página de preços do Dataflow.

Para mais informações sobre a utilização de TPUs com o Dataflow, consulte o artigo Apoio técnico do Dataflow para TPUs.

Opcional: faça uma reserva específica para usar aceleradores

Embora possa usar TPUs a pedido, recomendamos vivamente que use TPUs do Dataflow com reservas da Google Cloud Platform especificamente segmentadas. Isto ajuda a garantir que tem acesso aos aceleradores disponíveis e a tempos de arranque rápidos dos trabalhadores. Os pipelines que consomem uma reserva de TPU não requerem uma quota de TPU adicional.

Se não fizer uma reserva e optar por usar as UTPs a pedido, disponibilize a quota de UTPs antes de executar o pipeline.

Opcional: aprovisione a quota de TPUs

Pode usar as TPUs numa capacidade a pedido ou através de uma reserva. Se quiser usar as TPUs a pedido, tem de aprovisionar a quota de TPUs antes de o fazer. Se usar uma reserva especificamente segmentada, pode ignorar esta secção.

Para usar TPUs a pedido sem reserva, verifique o limite e a utilização atual da sua quota da API Compute Engine para TPUs da seguinte forma:

Consola

  1. Aceda à página Quotas na Google Cloud consola:

    Aceder a Quotas

  2. Na caixa Filtrar, faça o seguinte:

    1. Use a tabela seguinte para selecionar e copiar a propriedade da quota com base na versão da TPU e no tipo de máquina. Por exemplo, se planear criar nós da TPU v5e a pedido cujo tipo de máquina comece por ct5lp-, introduza Name: TPU v5 Lite PodSlice chips.

      Versão da TPU, o tipo de máquina começa com Propriedade e nome da quota para instâncias a pedido
      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. Selecione a propriedade Dimensões (por exemplo, localização) e introduza region: seguido do nome da região na qual planeia iniciar o seu pipeline. Por exemplo, introduza region:us-west4 se planear usar a zona us-west4-a. A quota de TPUs é regional, pelo que todas as zonas na mesma região consomem a mesma quota de TPUs.

Configure uma imagem de contentor personalizada

Para interagir com TPUs em pipelines do Dataflow, tem de fornecer software que possa funcionar em dispositivos XLA no ambiente de execução do pipeline. Isto requer a instalação de bibliotecas de TPUs com base nas necessidades do seu pipeline e a configuração de variáveis de ambiente com base no dispositivo de TPU que usa.

Para personalizar a imagem do contentor, instale o Apache Beam numa imagem base pronta a usar que tenha as bibliotecas de TPU necessárias. Em alternativa, instale o software da TPU nas imagens publicadas com as versões do SDK do Apache Beam.

Para fornecer uma imagem de contentor personalizada, use a opção sdk_container_imagepipeline. Para mais informações, consulte o artigo Use contentores personalizados no Dataflow.

Quando usa um acelerador de TPU, tem de definir as seguintes variáveis de ambiente na imagem do contentor.

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

Consoante o acelerador que usar, as variáveis na tabela seguinte também têm de ser definidas.

escrever topologia Dataflow obrigatório worker_machine_type variáveis de ambiente adicionais
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

Um Dockerfile de exemplo para a imagem do contentor personalizado pode ter o seguinte aspeto:

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

Execute o seu trabalho com TPUs

As considerações para executar uma tarefa do Dataflow com TPUs incluem o seguinte:

  • Uma vez que os contentores de TPUs podem ser grandes, para evitar ficar sem espaço em disco, aumente o tamanho predefinido do disco de arranque para 50 gigabytes ou um tamanho adequado conforme necessário para a imagem do contentor através da opção de pipeline --disk_size_gb.
  • Limite o paralelismo entre trabalhadores.

TPUs e paralelismo de trabalhadores

Na configuração predefinida, os pipelines Python do Dataflow iniciam um processo do SDK do Apache Beam por núcleo da VM. Os tipos de máquinas de TPU têm um grande número de núcleos de vCPU, mas apenas um processo pode realizar cálculos num dispositivo de TPU. Além disso, um dispositivo TPU pode ser reservado por um processo durante a duração do processo. Por conseguinte, tem de limitar o paralelismo intra-trabalhador quando executa um pipeline de TPU do Dataflow. Para limitar o paralelismo dos trabalhadores, use as seguintes orientações:

  • Se o seu exemplo de utilização envolver a execução de inferências num modelo, use a API Beam RunInference. Para mais informações, consulte o artigo Inferência do modelo de linguagem (conteúdo extenso) no Beam.
  • Se não puder usar a API Beam RunInference, use os objetos partilhados de vários processos do Beam para restringir determinadas operações a um único processo.
  • Se não puder usar as recomendações anteriores e preferir iniciar apenas um processo Python por trabalhador, defina a opção de pipeline --experiments=no_use_multiple_sdk_containers.
  • Para trabalhadores com mais de 100 vCPUs, reduza o número de threads usando a opção de pipeline --number_of_worker_harness_threads. Use a tabela seguinte para ver se o seu tipo de TPU usa mais de 100 vCPUs.

A tabela seguinte indica o total de recursos de computação por trabalhador para cada configuração de TPU.

Tipo de TPU topologia tipo de máquina Chips de 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

Execute um pipeline com TPUs

Para executar uma tarefa do Dataflow com TPUs, use o seguinte comando.

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

Substitua o seguinte:

  • PIPELINE: o ficheiro de código-fonte do pipeline.
  • PROJECT: o Google Cloud nome do projeto.
  • BUCKET: o contentor do Cloud Storage.
  • REGION: Uma região do Dataflow, por exemplo, us-central1.
  • TPU_TYPE: um tipo de TPU suportado, por exemplo, tpu-v5-lite-podslice. Para ver uma lista completa de tipos e topologias, consulte o artigo Aceleradores de TPU suportados.
  • TPU_TOPOLOGY: a topologia da TPU, por exemplo, 1x1.
  • MACHINE_TYPE: o tipo de máquina correspondente, por exemplo, ct5lp-hightpu-1t.
  • DISK_SIZE_GB: o tamanho do disco de arranque para cada VM de trabalho, por exemplo, 100.
  • IMAGE: o caminho do Artifact Registry para a sua imagem Docker.
  • NUMBER_OF_THREADS: opcional. O número de threads de worker harness.

Valide a tarefa do Dataflow

Para confirmar que a tarefa usa VMs de trabalho com TPUs, siga estes passos:

  1. Na Google Cloud consola, aceda à página Dataflow > Tarefas.

    Aceda a Empregos

  2. Selecione uma tarefa.

  3. Clique no separador Métricas de tarefas.

  4. Na secção Ajuste automático de escala, confirme que existe, pelo menos, uma VM Current workers.

  5. No painel lateral Informações do trabalho, verifique se o machine_type começa com ct. Por exemplo, ct6e-standard-1t. Isto indica a utilização da TPU.

Resolva problemas com a tarefa do Dataflow

Se tiver problemas ao executar a tarefa do Dataflow com TPUs, consulte o artigo Resolva problemas da tarefa do Dataflow TPU.

O que se segue?