Cette page a été traduite par l'API Cloud Translation.

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 Google Cloudspécifiquement ciblées. Cela permet de s'assurer que vous avez accès aux accélérateurs disponibles et que les workers démarrent rapidement. Les pipelines qui consomment une réservation TPU ne nécessitent pas de quota TPU supplémentaire.

Si vous ne faites pas de réservation et que vous 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 avant de le faire. 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

Accédez à la page Quotas de la console Google Cloud :

Accéder à la section "Quotas"

Dans la zone Filtre, procédez comme suit :

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 envisagez 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

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 les appareils XLA dans l'environnement d'exécution de votre pipeline. Pour ce faire, vous devez installer les bibliothèques TPU en fonction des besoins de votre pipeline et configurer les 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 disposant des 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é, utilisez l'option de pipeline sdk_container_image. Pour en savoir plus, consultez 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

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

type	topologie	Dataflow requis `worker_machine_type`	des 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é :

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 job avec des TPU

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

Les conteneurs 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 processus peut réserver un appareil TPU pour 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 Dataflow TPU. Pour limiter le parallélisme des nœuds de calcul, suivez les instructions suivantes :

Si votre cas d'utilisation implique l'exécution d'inférences sur un modèle, utilisez l'API RunInference de Beam. Pour en savoir plus, consultez Inférence de grands modèles linguistiques dans Beam.
Si vous ne pouvez pas utiliser l'API Beam RunInference, utilisez les objets partagés multiprocessus de Beam pour limiter certaines opérations à un seul processus.
Si vous ne pouvez pas utiliser les recommandations précédentes et que vous préférez lancer un seul processus Python par nœud de calcul, définissez l'option de pipeline --experiments=no_use_multiple_sdk_containers.
Pour les nœuds de calcul disposant de plus de 100 processeurs virtuels, réduisez le nombre de threads à l'aide de l'option de pipeline --number_of_worker_harness_threads. Consultez le tableau suivant pour savoir si votre type de TPU utilise plus de 100 vCPU.

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