Ejecuta una canalización con TPUs

En esta página, se explica cómo ejecutar una canalización de Apache Beam en Dataflow con TPU. Los trabajos que usan TPU generan cargos, como se especifica en la página de precios de Dataflow.

Para obtener más información sobre el uso de TPU con Dataflow, consulta la compatibilidad de Dataflow con las TPU.

Opcional: Haz una reserva específica para usar aceleradores

Si bien puedes usar TPUs a pedido, te recomendamos que uses TPUs de Dataflow con reservas específicamente segmentadas Google Cloud. Esto ayuda a garantizar que tengas acceso a los aceleradores disponibles y a tiempos de inicio rápidos de los trabajadores. Las canalizaciones que consumen una reserva de TPU no requieren cuota de TPU adicional.

Si no haces una reserva y decides usar las TPU a pedido, aprovisiona la cuota de TPU antes de ejecutar tu canalización.

Opcional: Asigna una cuota de TPU

Puedes usar las TPU con capacidad a pedido o con una reserva. Si deseas usar TPUs según demanda, debes aprovisionar la cuota de TPU antes de hacerlo. Si usas una reserva específicamente segmentada, puedes omitir esta sección.

Para usar las TPU bajo demanda sin una reserva, verifica el límite y el uso actual de la cuota de la API de Compute Engine para las TPU de la siguiente manera:

Console

  1. Ve a la página Cuotas en la consola de Google Cloud :

    Ir a Cuotas

  2. En la casilla Filtro, haz lo siguiente:

    1. Usa la siguiente tabla para seleccionar y copiar la propiedad de la cuota según la versión de TPU y el tipo de máquina. Por ejemplo, si planeas crear nodos TPU v5e a pedido cuyo tipo de máquina comience con ct5lp-, ingresa Name: TPU v5 Lite PodSlice chips.

      Versión de TPU, el tipo de máquina comienza con Propiedad y nombre de la cuota para instancias bajo demanda
      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. Selecciona la propiedad Dimensiones (p.ej., ubicación) y, luego, ingresa region: seguido del nombre de la región en la que planeas iniciar tu canalización. Por ejemplo, ingresa region:us-west4 si planeas usar la zona us-west4-a. La cuota de TPU es regional, por lo que todas las zonas dentro de la misma región consumen la misma cuota de TPU.

Configura una imagen de contenedor personalizada

Para interactuar con las TPU en las canalizaciones de Dataflow, debes proporcionar software que pueda operar en dispositivos XLA en tu entorno de ejecución de canalizaciones. Esto requiere que instales bibliotecas de TPU según las necesidades de tu canalización y que configures variables de entorno según el dispositivo TPU que uses.

Para personalizar la imagen del contenedor, instala Apache Beam en una imagen base lista para usar que tenga las bibliotecas de TPU necesarias. Como alternativa, instala el software de TPU en las imágenes publicadas con los lanzamientos del SDK de Apache Beam.

Para proporcionar una imagen de contenedor personalizada, usa la opción de canalización sdk_container_image. Para obtener más información, consulta Usa contenedores personalizados en Dataflow.

Cuando usas un acelerador de TPU, debes establecer las siguientes variables de entorno en la imagen del contenedor.

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

Según el acelerador que uses, también deberás establecer las variables de la siguiente tabla.

tipo topology worker_machine_type de Dataflow requerida Variables de entorno adicionales
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 ejemplo de Dockerfile para la imagen de contenedor personalizada podría verse de la siguiente manera:

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

Ejecuta tu trabajo con TPU

Las consideraciones para ejecutar un trabajo de Dataflow con TPU incluyen las siguientes:

  • Debido a que los contenedores de TPU pueden ser grandes, para evitar quedarte sin espacio en el disco, aumenta el tamaño predeterminado del disco de arranque a 50 gigabytes o a un tamaño adecuado según lo requiera tu imagen de contenedor con la opción de canalización --disk_size_gb.
  • Limita el paralelismo dentro del trabajador.

TPUs y paralelismo de trabajadores

En la configuración predeterminada, las canalizaciones de Python de Dataflow inician un proceso del SDK de Apache Beam por núcleo de VM. Los tipos de máquinas TPU tienen una gran cantidad de núcleos de CPU virtuales, pero solo un proceso puede realizar cálculos en un dispositivo TPU. Además, un proceso puede reservar un dispositivo TPU durante toda su vida útil. Por lo tanto, debes limitar el paralelismo dentro del trabajador cuando ejecutes una canalización de TPU de Dataflow. Para limitar el paralelismo de trabajadores, sigue estas instrucciones:

  • Si tu caso de uso implica ejecutar inferencias en un modelo, usa la API de RunInference de Beam. Para obtener más información, consulta Inferencia de modelos de lenguaje grandes en Beam.
  • Si no puedes usar la API de Beam RunInference, usa los objetos compartidos de varios procesos de Beam para restringir ciertas operaciones a un solo proceso.
  • Si no puedes usar las recomendaciones anteriores y prefieres iniciar solo un proceso de Python por trabajador, configura la opción de canalización --experiments=no_use_multiple_sdk_containers.
  • Para los trabajadores con más de 100 CPU virtuales, reduce la cantidad de subprocesos con la opción de canalización --number_of_worker_harness_threads. Usa la siguiente tabla para ver si tu tipo de TPU usa más de 100 CPU virtuales.

En la siguiente tabla, se enumeran los recursos de procesamiento totales por trabajador para cada configuración de TPU.

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

Ejecuta una canalización con TPUs

Para ejecutar un trabajo de Dataflow con TPU, usa el siguiente 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

Reemplaza lo siguiente:

  • PIPELINE: Tu archivo de código fuente de canalización
  • PROJECT: El nombre del proyecto de Google Cloud
  • BUCKET: El bucket de Cloud Storage.
  • REGION: Una región de Dataflow, por ejemplo, us-central1.
  • TPU_TYPE: Es un tipo de TPU compatible, por ejemplo, tpu-v5-lite-podslice. Para obtener una lista completa de los tipos y las topologías, consulta Aceleradores de TPU compatibles.
  • TPU_TOPOLOGY: Es la topología de TPU, por ejemplo, 1x1.
  • MACHINE_TYPE: Es el tipo de máquina correspondiente, por ejemplo, ct5lp-hightpu-1t.
  • DISK_SIZE_GB: Es el tamaño del disco de arranque para cada VM de trabajador, por ejemplo, 100.
  • IMAGE: Es la ruta de acceso de Artifact Registry para tu imagen de Docker.
  • NUMBER_OF_THREADS: Opcional Es la cantidad de subprocesos del agente de trabajo.

Verifica tu trabajo de Dataflow

Para confirmar que el trabajo use VMs de trabajador con TPU, sigue estos pasos:

  1. En la consola de Google Cloud , ve a la página Trabajos de Dataflow >.

    Ir a Trabajos

  2. Selecciona un trabajo.

  3. Haz clic en la pestaña Métricas del trabajo.

  4. En la sección Ajuste de escala automático, confirma que haya al menos una VM de Trabajadores actuales.

  5. En el panel lateral Información del trabajo, verifica que machine_type comience con ct. Por ejemplo, ct6e-standard-1t. Esto indica el uso de la TPU.

Soluciona problemas de tu trabajo de Dataflow

Si tienes problemas para ejecutar tu trabajo de Dataflow con TPU, consulta Soluciona problemas de tu trabajo de TPU de Dataflow.

¿Qué sigue?