Exécuter un pipeline avec des TPU

Cette page explique comment exécuter un pipeline Apache Beam sur Dataflow avec des TPU. Les tâches utilisant des TPU sont facturées comme indiqué sur la page des tarifs de Dataflow.

Pour en savoir plus sur l'utilisation des TPU avec Dataflow, consultez la page Compatibilité de Dataflow avec les TPU.

Facultatif : Effectuer une réservation spécifique pour utiliser des accélérateurs

Bien que vous puissiez utiliser des TPU à la demande, nous vous recommandons vivement d'utiliser des TPU Dataflow avec des réservations spécifiquement ciblées Google Cloud. Cela vous permet de vous assurer d'avoir accès aux accélérateurs disponibles et de bénéficier de temps de démarrage rapides des nœuds de calcul. Les pipelines qui consomment une réservation de TPU ne nécessitent pas de quota de TPU supplémentaire.

Si vous n'effectuez pas de réservation et choisissez d'utiliser des TPU à la demande, provisionnez un quota de TPU avant d'exécuter votre pipeline.

Facultatif : Provisionner un quota de TPU

Vous pouvez utiliser des TPU avec une capacité à la demande ou à l'aide d'une réservation. Si vous souhaitez utiliser des TPU à la demande, vous devez provisionner un quota de TPU au préalable. Si vous utilisez une réservation spécifiquement ciblée, vous pouvez ignorer cette section.

Pour utiliser des TPU à la demande sans réservation, vérifiez la limite et l'utilisation actuelle de votre quota d'API Compute Engine pour les TPU comme suit :

Console

  1. Accédez à la page Quotas dans la Google Cloud console :

    Accéder à la section "Quotas"

  2. Dans la zone Filter (Filtrer), procédez comme suit :

    1. Utilisez le tableau suivant pour sélectionner et copier la propriété du quota en fonction de la version du TPU et du type de machine. Par exemple, si vous prévoyez de créer des nœuds TPU v5e à la demande dont le type de machine commence par ct5lp-, saisissez Name: TPU v5 Lite PodSlice chips.

      Version du TPU, type de machine commençant par Propriété et nom du quota pour les instances à la demande
      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. Sélectionnez la propriété Dimensions (par exemple, emplacement) et saisissez region: suivi du nom de la région dans laquelle vous prévoyez de démarrer votre pipeline. Par exemple, saisissez region:us-west4 si vous prévoyez d'utiliser la zone us-west4-a. Le quota de TPU est régional. Par conséquent, toutes les zones d'une même région consomment le même quota de TPU.

Configurer une image de conteneur personnalisée

Pour interagir avec les TPU dans les pipelines Dataflow, vous devez fournir un logiciel capable de fonctionner sur des appareils XLA dans votre environnement d'exécution de pipeline. Pour cela, vous devez installer des bibliothèques de TPU en fonction des besoins de votre pipeline et configurer des variables d'environnement en fonction de l'appareil TPU que vous utilisez.

Pour personnaliser l'image de conteneur, installez Apache Beam dans une image de base prête à l'emploi qui contient les bibliothèques de TPU nécessaires. Vous pouvez également installer le logiciel TPU dans les images publiées avec les versions du SDK Apache Beam.

Pour fournir une image de conteneur personnalisée, utilisez l'option de pipeline sdk_container_image. Pour en savoir plus, consultez la page Utiliser des conteneurs personnalisés dans Dataflow.

Lorsque vous utilisez un accélérateur TPU, vous devez définir les variables d'environnement suivantes dans l'image de conteneur.

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

En fonction de l'accélérateur que vous utilisez, les variables du tableau suivant doivent également être définies.

type topologie worker_machine_type Dataflow requis variables d'environnement supplémentaires
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

Voici un exemple de fichier Dockerfile pour l'image de conteneur personnalisée :

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

Exécuter votre tâche avec des TPU

Les points à prendre en compte pour exécuter une tâche Dataflow à l'aide de TPU sont les suivants :

  • Les conteneurs de TPU pouvant être volumineux, pour éviter de manquer d'espace disque, augmentez la taille du disque de démarrage par défaut à 50 gigaoctets ou à une taille appropriée, selon les besoins de votre image de conteneur, à l'aide de l'option de pipeline --disk_size_gb.
  • Limitez le parallélisme intra-nœud de calcul.

TPU et parallélisme des nœuds de calcul

Dans la configuration par défaut, les pipelines Python Dataflow lancent un processus du SDK Apache Beam par cœur de VM. Les types de machines TPU disposent d'un grand nombre de cœurs de processeur virtuel, mais un seul processus peut effectuer des calculs sur un appareil TPU. De plus, un appareil TPU peut être réservé par un processus pendant toute sa durée de vie. Par conséquent, vous devez limiter le parallélisme intra-nœud de calcul lorsque vous exécutez un pipeline TPU Dataflow. Pour limiter le parallélisme des nœuds de calcul, suivez les conseils ci-dessous :

Le tableau suivant répertorie le total des ressources de calcul par nœud de calcul pour chaque configuration de TPU.

Type de TPU topologie type de machine Puces TPU vCPU Mémoire RAM (Go)
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

Exécuter un pipeline avec des TPU

Pour exécuter une tâche Dataflow à l'aide de TPU, utilisez la commande suivante.

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

Remplacez les éléments suivants :

  • PIPELINE : fichier de code source de votre pipeline.
  • PROJECT : nom du Google Cloud projet.
  • BUCKET : bucket Cloud Storage.
  • REGION : région Dataflow, par exemple us-central1.
  • TPU_TYPE : type de TPU compatible, par exemple tpu-v5-lite-podslice. Pour obtenir la liste complète des types et des topologies, consultez la section Accélérateurs TPU compatibles.
  • TPU_TOPOLOGY : topologie du TPU, par exemple 1x1.
  • MACHINE_TYPE : type de machine correspondant, par exemple ct5lp-hightpu-1t.
  • DISK_SIZE_GB : taille du disque de démarrage pour chaque VM de nœud de calcul, par exemple 100.
  • IMAGE : chemin d'accès à Artifact Registry pour votre image Docker.
  • NUMBER_OF_THREADS : facultatif. Nombre de threads de faisceau de nœud de calcul.

Vérifier la tâche Dataflow

Pour vérifier que la tâche utilise des VM de nœud de calcul à l'aide de TPU, procédez comme suit :

  1. Dans la Google Cloud console, accédez à la page Dataflow > Jobs.

    Accéder aux tâches

  2. Sélectionnez une tâche.

  3. Cliquez sur l'onglet Job metrics (Métriques de job).

  4. Dans la section Autoscaling (Autoscaling), vérifiez qu'il existe au moins une VM Current workers (Nœuds de calcul actuels).

  5. Dans le panneau latéral Job info (Informations sur la tâche), vérifiez que machine_type commence par ct. Par exemple, ct6e-standard-1t. Cela indique l'utilisation du TPU.

Résoudre les problèmes liés à la tâche Dataflow

Si vous rencontrez des problèmes lors de l'exécution de votre tâche Dataflow à l'aide de TPU, consultez la page Résoudre les problèmes liés à la tâche TPU Dataflow.

Étape suivante