Présentation des TPU dans GKE

Les clients de Google Kubernetes Engine (GKE) peuvent désormais créer des pools de nœuds Kubernetes contenant des pods TPU v4 et v5e. Un pod TPU est un groupe d'appareils TPU connectés par des interconnexions à haut débit. Pour les charges de travail qui ne nécessitent pas de pod TPU complet, vous pouvez utiliser un sous-ensemble d'un pod TPU complet, appelé tranche TPU. Comme pour les pods TPU complets, chaque appareil TPU d'une tranche possède sa propre VM TPU. Nous appelons une VM TPU et son appareil connecté en tant qu'hôte ou nœud TPU. Pour en savoir plus sur les pods TPU, consultez la page Architecture du système.

Étant donné que le terme pod utilisé dans le contexte de GKE désigne généralement un pod Kubernetes, pour éviter toute confusion, nous appelons toujours un ensemble d'un ou plusieurs appareils TPU en tant que tranche.

Lorsque vous travaillez avec GKE, vous devez d'abord créer un cluster GKE. Vous devez ensuite ajouter des pools de nœuds à votre cluster. Les pools de nœuds GKE sont des ensembles de VM qui partagent les mêmes attributs. Pour les charges de travail TPU, les pools de nœuds sont constitués de VM TPU.

Types de pools de nœuds

GKE est compatible avec deux types de pools de nœuds TPU:

• Tranche de TPU à hôtes multiples
• Tranche de TPU à hôte unique

Pool de nœuds de tranche TPU multi-hôtes

Un pool de nœuds de tranches de TPU multi-hôtes est un pool de nœuds qui contient au moins deux VM TPU interconnectées. Un appareil TPU est connecté à chaque VM. Les TPU d'une tranche de plusieurs hôtes sont connectés via une interconnexion à haut débit (ICI). Un pool de nœuds de tranche TPU multi-hôtes est immuable. Une fois qu'un pool de nœuds de tranches multi-hôtes est créé, vous ne pouvez plus y ajouter de nœuds. Par exemple, vous ne pouvez pas créer un pool de nœuds v4-32, puis y ajouter ultérieurement un nœud Kubernetes (VM TPU) supplémentaire. Pour ajouter une tranche de TPU supplémentaire à un cluster GKE, vous devez créer un nouveau pool de nœuds.

Les hôtes d'un pool de nœuds de tranche TPU à hôtes multiples sont traités comme une seule unité atomique. Si GKE ne parvient pas à déployer un nœud du segment, tous les nœuds du segment ne seront pas déployés.

Si un nœud d'une tranche de TPU à hôtes multiples doit être réparé, GKE arrête toutes les VM (nœuds) TPU de la tranche, forçant ainsi l'éviction de tous les pods GKE de la charge de travail. Une fois que toutes les VM TPU de la tranche sont opérationnelles, les pods GKE peuvent être planifiés sur les VM TPU de la nouvelle tranche.

Le schéma suivant montre un exemple de tranche multi-hôte, un TPU v4-32 dont la topologie est 2x2x4. Cette tranche comporte quatre VM TPU. Chaque VM TPU comporte quatre puces TPU v4 connectées à une interconnexion à haut débit (ICI) et chaque puce TPU v4 possède deux Tensor Cores.

Le schéma suivant montre un cluster GKE contenant une tranche TPU v5litepod-16 (v5e) (topologie: 4x4) et une tranche TPU v5litepod-8 (v5e) (topologie: 2x4):

Pools de nœuds de tranches TPU à hôte unique

Un pool de nœuds de tranche à hôte unique est un pool de nœuds contenant une ou plusieurs VM TPU indépendantes. Chacune de ces VM est connectée à un appareil TPU. Bien que les VM d'un pool de nœuds en tranche à hôte unique puissent communiquer via le réseau de centre de données (DCN), les TPU associés aux VM ne sont pas interconnectés. Le schéma suivant montre un exemple de tranche de TPU à hôte unique avec neuf machines v4-8:

Types de machines TPU pour les pools de nœuds GKE

Avant de créer des pools de nœuds, vous devez choisir la version et la taille de la tranche d'emploi requise par votre charge de travail. GKE est compatible avec TPU v4 dans les versions 1.26.1.gke-1500 et ultérieures de GKE, ainsi que dans TPU v5e dans les versions 1.27.2-gke.1500 et ultérieures de GKE.

Pour en savoir plus sur les caractéristiques matérielles des différentes versions de TPU, consultez la section Architecture système. Lorsque vous créez un pool de nœuds TPU, sélectionnez une taille de tranche TPU (topologie de TPU) en fonction de la taille de votre modèle et de la quantité de mémoire requise. Le type de machine que vous spécifiez lors de la création des pools de nœuds dépend de la version et de la taille de vos segments d'application.

v4

Lorsque vous créez un pod ou une tranche TPU v4, utilisez le type de machine ct4p-hightpu-4t, qui comporte un hôte et quatre puces. Pour en savoir plus, consultez les pages Topologies v4 et Architecture du système TPU. Les types de machines de pod TPU v4 sont disponibles dans us-central2-b. Vos clusters GKE doivent exécuter le plan de contrôle 1.26.1-gke.1500 ou une version ultérieure.

V5E

Les types de machines et topologies TPU v5e compatibles avec les cas d'utilisation d'entraînement et d'inférence sont les suivants:

Type de machine	Topologie	Nombre de puces TPU	Nombre de VM	Cas d'utilisation recommandé
ct5lp-hightpu-1t	1x1	1	1	Entraînement, inférence à hôte unique
ct5lp-hightpu-4t	2x2	4	1	Entraînement, inférence à hôte unique
ct5lp-hightpu-8t	2x4	8	1	Entraînement, inférence à hôte unique
ct5lp-hightpu-4t	2x4	8	2	Entraînement, inférence multi-hôte
ct5lp-hightpu-4t	4x4	16	4	Entraînement à grande échelle, inférence multi-hôte
ct5lp-hightpu-4t	4x8	32	8	Entraînement à grande échelle, inférence multi-hôte
ct5lp-hightpu-4t	8x8	64	16	Entraînement à grande échelle, inférence multi-hôte
ct5lp-hightpu-4t	8x16	128	32	Entraînement à grande échelle, inférence multi-hôte
ct5lp-hightpu-4t	16x16	256	64	Entraînement à grande échelle, inférence multi-hôte

Cloud TPU v5e est un produit combiné d'entraînement et d'inférence. Les tâches d'entraînement sont optimisées pour le débit et la disponibilité, tandis que les tâches d'inférence sont optimisées pour la latence. Pour en savoir plus, consultez les pages Types d'accélérateurs d'entraînement v5e et Types d'accélérateurs d'inférence v5e.

Les machines TPU v5e sont disponibles dans us-west4-a, us-east5-b et us-east1-c. Vos clusters GKE doivent exécuter le plan de contrôle version 1.27.2-gke.2100 ou ultérieure. Pour en savoir plus sur la version v5e, consultez la page Entraînement sur Cloud TPU v5e.

Comparaison des types de machines:

Type de machine	CT5lp-hightpu-1t	CT5lp-hightpu-4t	CT5lp-hightpu-8t
Nombre de puces v5e	1	4	8
Nombre de processeurs virtuels	24	112	224
Mémoire RAM (Go)	48	192	384
Nombre de nœuds NUMA	1	1	2
Probabilité de préemption	Élevée	Moyenne	Faible

Pour libérer de l'espace pour les VM comportant plus de puces, le programmeur GKE peut préempter et reprogrammer les VM comportant moins de puces. Les VM à 8 puces sont donc susceptibles de préempter les VM à 1 ou 4 puces.

Limites et problèmes connus

• Réservations spécifiques uniquement: lorsque vous utilisez des TPU dans GKE, SPECIFIC est la seule valeur acceptée pour l'option --reservation-affinity de la commande gcloud container node-pools create.
• Seule la variante de VM Spot des TPU préemptifs est compatible: les VM Spot sont semblables aux VM préemptives et sont soumises aux mêmes limites de disponibilité, mais leur durée maximale n'est pas de 24 heures.
• Les clusters Autopilot ne sont pas compatibles: les TPU ne sont pas disponibles dans les clusters GKE Autopilot.
• Pas de répartition des coûts disponible: la répartition des coûts GKE et la mesure de l'utilisation n'incluent aucune donnée sur l'utilisation ou les coûts des TPU.
• L'autoscaler peut calculer la capacité: l'autoscaler de cluster peut calculer la capacité de manière incorrecte pour les nouveaux nœuds TPU avant que ces nœuds ne soient disponibles. L'autoscaler de cluster peut ensuite effectuer un scaling à la hausse supplémentaire et, par conséquent, créer plus de nœuds que nécessaire. L'autoscaler de cluster réduira les nœuds supplémentaires, s'ils ne sont pas nécessaires, après une opération de scaling à la baisse régulière.
• L'autoscaler annule le scaling à la hausse: l'autoscaler de cluster annule le scaling à la hausse des pools de nœuds TPU qui restent en attente pendant plus de 15 minutes. L'autoscaler de cluster relancera ces opérations de scaling à la hausse ultérieurement. Ce comportement peut réduire l'obtention de TPU pour les clients qui n'utilisent pas les réservations.
• Le rejet peut empêcher le scaling à la baisse: les charges de travail non-TPU qui tolèrent le rejet TPU peuvent empêcher le scaling à la baisse du pool de nœuds si elles sont recréées lors du drainage du pool de nœuds TPU.

Garantir des quotas TPU et GKE suffisants

Quota lié aux pods TPU v4

Étant donné que la capacité de TPU v4 est située dans la région us-central2-b, vous devrez peut-être augmenter certains quotas liés à GKE dans la région us-central2.

Les quotas suivants comportent des valeurs par défaut qui devront probablement être augmentées:

• Quota de disques persistants SSD (Go) dans la région us-central2: le disque de démarrage de chaque nœud Kubernetes nécessite 100 Go par défaut. Par conséquent, ce quota doit être défini au moins égal au (nombre maximal de nœuds GKE que vous prévoyez de créer en us-central2) x 100 Go.
• Quota d'adresses IP en cours d'utilisation dans la région us-central2: chaque nœud Kubernetes utilise une adresse IP. Par conséquent, ce quota doit être défini au moins égal au nombre maximal de nœuds GKE que vous prévoyez de créer dans us-central2.

Pour demander une augmentation de quota, consultez Demander un quota supérieur. Pour en savoir plus sur les types de quotas de TPU, consultez la section Quotas de TPU.

Quota associé à v5e

Suivez la même procédure que pour la version 4 et pour les demandes d'augmentation de quota de fichiers pour les disques et les adresses IP dans les régions souhaitées (par exemple, us-west4, us-east1, us-east5) si nécessaire.

L'approbation de vos demandes d'augmentation de quota peut prendre quelques jours. Si vous rencontrez des difficultés pour faire approuver vos demandes d'augmentation de quota en quelques jours, contactez l'équipe chargée de votre compte Google.

Migrer votre réservation TPU

Si vous ne prévoyez pas d'utiliser une réservation TPU existante avec des TPU dans GKE, ignorez cette section et accédez à la section Créer un cluster Google Kubernetes Engine.

Pour utiliser des TPU réservés avec GKE, vous devez d'abord migrer votre réservation de TPU vers un nouveau système de réservation basé sur Compute Engine.

Plusieurs points importants sont à noter sur cette migration:

• Réservation incompatible avec l'API Queued Resource: la capacité de TPU migrée vers le nouveau système de réservation basé sur Compute Engine ne peut pas être utilisée avec l'API Queued Resource de Cloud TPU. Si vous avez l'intention d'utiliser des ressources en file d'attente TPU avec votre réservation, vous ne devez migrer qu'une partie de votre réservation TPU vers le nouveau système de réservation basé sur Compute Engine.
• Aucune charge de travail en cours d'exécution pendant la migration: aucune charge de travail ne peut être en cours d'exécution active sur les TPU au moment de leur migration vers le nouveau système de réservation basé sur Compute Engine. Par conséquent, vous devez planifier à l'avance le temps de migration qui vous permettra d'arrêter toutes les charges de travail avant la migration.
• Vous devez sélectionner une heure de migration: vous devez sélectionner une heure à laquelle vous souhaitez effectuer la migration et planifier la migration avec l'équipe chargée de votre compte Google Cloud. La période de migration doit être comprise dans les heures de bureau (du lundi au vendredi, de 9h à 17h, heure du Pacifique).

Créer un cluster Google Kubernetes Engine

Consultez la section Créer un cluster dans la documentation de Google Kubernetes Engine.

Créer un pool de nœuds TPU

Consultez la section Créer un pool de nœuds dans la documentation de Google Kubernetes Engine.

Exécution sans mode privilégié

Si vous souhaitez réduire le champ d'application des autorisations sur votre conteneur, consultez la section Mode privilège TPU.

Exécuter des charges de travail sur des nœuds TPU

Consultez la section Exécuter vos charges de travail sur des nœuds TPU dans la documentation Google Kubernetes Engine.

Sélecteurs de nœuds

Pour que Kubernetes puisse planifier votre charge de travail sur les nœuds TPU, vous devez spécifier deux sélecteurs pour chaque nœud TPU dans votre fichier manifeste Google Kubernetes Engine:

• Définissez cloud.google.com/gke-accelerator-type sur tpu-v5-lite-podslice ou tpu-v4-podslice.
• Définissez cloud.google.com/gke-tpu-topology sur la topologie TPU du nœud TPU.

Important :Vos pods GKE doivent obligatoirement spécifier ces deux étiquettes sous le nœud nodeSelector ou nodeAffinity dans le fichier manifeste de votre charge de travail.

Les sections Charges de travail d'entraînement et Charges de travail d'inférence contiennent des exemples de fichiers manifestes illustrant l'utilisation de ces sélecteurs de nœuds.

Hôtes multiples

Pour obtenir un exemple d'exécution d'une charge de travail sur une tranche de TPU multi-hôte, consultez la section Exécuter la charge de travail sur une tranche de TPU multi-hôte.

Un seul hôte

Pour obtenir un exemple d'exécution d'une charge de travail sur une tranche de TPU à hôte unique, consultez la page Exécuter vos charges de travail sur des nœuds TPU.

Remarques concernant la planification des charges de travail

Les TPU présentent des caractéristiques uniques qui nécessitent une planification et une gestion spéciales des charges de travail dans Kubernetes. Pour en savoir plus, consultez la section Remarques concernant la planification des charges de travail dans la documentation GKE.

Réparation de nœuds TPU

Si un nœud TPU d'un pool de nœuds de tranches TPU à plusieurs hôtes n'est pas opérationnel, l'intégralité du pool de nœuds est recréée. Pour en savoir plus, consultez la page Réparation automatique des nœuds dans la documentation GKE.

Multitranche : au-delà d'une seule tranche

Vous pouvez agréger des tranches plus petites au sein d'une multitranche pour gérer des charges de travail d'entraînement plus importantes. Pour en savoir plus, consultez la section Cloud TPU Multislice.

Tutoriels sur les charges de travail d'entraînement

Ces tutoriels se concentrent sur les charges de travail d'entraînement sur une tranche de TPU multi-hôte (par exemple, quatre machines v5e). Elles couvrent les modèles suivants:

• Modèles FLAX pour les visages: Train Diffusion sur Pokémon
• PyTorch/XLA: GPT2 sur WikiText

Conditions d'accès

Vous devez avoir accès au bucket Cloud Storage gs://gke-tpu-examples/.

Télécharger les ressources du tutoriel

Téléchargez les scripts Python et les spécifications YAML du tutoriel pour chaque modèle pré-entraîné à l'aide de la commande suivante:

gsutil -m cp -r "gs://gke-tpu-examples/training"

Créer un cluster et s'y connecter

Créez un cluster GKE standard régional pour que le plan de contrôle Kubernetes soit répliqué dans trois zones, ce qui offre une plus grande disponibilité. Créez votre cluster dans us-west4, us-east1 ou us-central2 selon la version de TPU que vous utilisez. Pour en savoir plus sur les TPU et les zones, consultez la page Régions et zones de Cloud TPU.

La commande suivante crée un cluster régional GKE abonné au version disponible rapide avec un pool de nœuds contenant initialement un nœud par zone. La commande active également les fonctionnalités de pilote CSI Workload Identity et Cloud Storage FUSE sur votre cluster, car les exemples de charges de travail d'inférence présentés dans ce guide utilisent des buckets Cloud Storage pour stocker des modèles pré-entraînés.

gcloud container clusters create cluster-name \
  --region your-region \
  --release-channel rapid \
  --num-nodes=1 \
  --workload-pool=project-id.svc.id.goog \
  --addons GcsFuseCsiDriver

Pour activer les fonctionnalités de pilote CSI Workload Identity et Cloud Storage FUSE pour les clusters existants, exécutez la commande suivante:

gcloud container clusters update cluster-name \
  --region your-region \
  --update-addons GcsFuseCsiDriver=ENABLED \
  --workload-pool=project-id.svc.id.goog

Les exemples de charges de travail ont été configurés selon les hypothèses suivantes:

• le pool de nœuds utilise tpu-topology=4x4 avec quatre nœuds
• le pool de nœuds utilise machine-type ct5lp-hightpu-4t

Exécutez la commande suivante pour vous connecter au cluster que vous venez de créer:

gcloud container clusters get-credentials cluster-name \
--location=cluster-region

Modèles FLAX pour les visages: Train Diffusion sur Pokémon

Cet exemple entraîne le modèle Stable Diffusion à partir de HuggingFace à l'aide de l'ensemble de données Pokémon.

Le modèle de diffusion stable est un modèle latent texte-image qui génère des images photoréalistes à partir de n'importe quelle entrée de texte. Pour en savoir plus sur Stable Diffusion, consultez les pages suivantes:

• Présentation de la diffusion stable
• Source du code de diffusion stable

Créer une image Docker

Ce fichier se trouve dans le dossier gs://gke-tpu-examples/training/diffusion/. Exécutez les commandes suivantes pour compiler et transférer l'image Docker.

cd diffusion/
docker build -t gcr.io/project-id/diffusion:latest .
docker push gcr.io/project-id/diffusion:latest

Déployer la charge de travail

Créez un fichier avec le contenu suivant et nommez-le tpu_job_diffusion.yaml. Remplissez le champ "Image" avec l'image que vous venez de créer.

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-diffusion
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-diffusion
spec:
  backoffLimit: 0
  # Completions and parallelism should be the number of chips divided by 4.
  # (e.g. 4 for a v5litepod-16)
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job-diffusion
        image: gcr.io/project-name/diffusion:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU usage metrics, if supported
        securityContext:
          privileged: true
        command:
        - bash
        - -c
        - |
          cd examples/text_to_image
          python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

Déployez-la ensuite à l'aide de la commande suivante:

kubectl apply -f tpu_job_diffusion.yaml

Effectuer un nettoyage

Une fois votre tâche exécutée, vous pouvez la supprimer à l'aide de la commande suivante:

kubectl delete -f tpu_job_diffusion.yaml

PyTorch/XLA: GPT2 sur WikiText

Ce tutoriel explique comment exécuter GPT2 sur des TPU v5e à l'aide de HuggingFace sur PyTorch/XLA à l'aide de l'ensemble de données wikitext.

Créer une image Docker

Ce Dockerfile se trouve dans le dossier gs://gke-tpu-examples/training/gpt/. Exécutez les commandes suivantes pour créer et transférer l'image Docker.

cd gpt/
docker build -t gcr.io/project-id/gpt:latest .
docker push gcr.io/project-id/gpt:latest

Déployer la charge de travail

Copiez le fichier YAML suivant et enregistrez-le dans un fichier appelé tpu_job_gpt.yaml. Remplissez le champ "Image" avec l'image que vous venez de créer.

kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-gpt
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-gpt
spec:
  backoffLimit: 0
  # Completions and parallelism should be the number of chips divided by 4.
  # (for example, 4 for a v5litepod-16)
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job-gpt
        image: gcr.io/project-name/gpt:latest
        ports:
        - containerPort: 8479
        - containerPort: 8478
        - containerPort: 8477
        - containerPort: 8476
        - containerPort: 8431 # Port to export TPU usage metrics, if supported.
        securityContext:
          privileged: true
        env:
        - name: PJRT_DEVICE
          value: 'TPU_C_API'
        - name: XLA_USE_BF16
          value: '1'
        command:
        - bash
        - -c
        - |
          numactl --cpunodebind=0 python3 -u examples/pytorch/xla_spawn.py   --num_cores 4 examples/pytorch/language-modeling/run_clm.py    --num_train_epochs 3 --dataset_name wikitext     --dataset_config_name wikitext-2-raw-v1 --per_device_train_batch_size 16    --per_device_eval_batch_size 16 --do_train --do_eval  --output_dir /tmp/test-clm     --overwrite_output_dir --config_name my_config_2.json --cache_dir /tmp --tokenizer_name gpt2  --block_size 1024 --optim adafactor --adafactor true --save_strategy no --logging_strategy no --fsdp "full_shard" --fsdp_config fsdp_config.json
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

Déployez le workflow à l'aide de la commande suivante:

kubectl apply -f tpu_job_gpt.yaml

Effectuer un nettoyage

Une fois l'exécution de votre job terminée, vous pouvez la supprimer à l'aide de la commande suivante:

kubectl delete -f tpu_job_gpt.yaml

Tutoriel: Charges de travail d'inférence à hôte unique

Ce tutoriel explique comment exécuter une charge de travail d'inférence à hôte unique sur des TPU GKE v5e pour des modèles pré-entraînés avec JAX, TensorFlow et PyTorch. De manière générale, vous devez effectuer quatre étapes distinctes sur le cluster GKE:

1. créer un bucket Cloud Storage et configurer l'accès à celui-ci ; Un bucket Cloud Storage est utilisé pour stocker le modèle pré-entraîné.

2. télécharger et convertir un modèle pré-entraîné en modèle compatible avec le TPU ; Appliquer un pod GKE qui télécharge le modèle pré-entraîné, utilise le convertisseur Cloud TPU et stocke les modèles convertis dans un bucket Cloud Storage à l'aide du pilote CSI FUSE de Cloud Storage. Le convertisseur Cloud TPU ne nécessite pas de matériel spécialisé. Ce tutoriel explique comment télécharger le modèle et exécuter le convertisseur Cloud TPU dans le pool de nœuds de processeur.

3. Lancez le serveur du modèle converti. Appliquez un déploiement qui diffuse le modèle à l'aide d'un framework de serveur reposant sur le volume stocké dans le volume persistant ReadOnlyMany (ROX). Les instances répliquées de déploiement doivent être exécutées sur un nœud TPU de pod v5e avec un pod Kubernetes par nœud.

4. Déployez un équilibreur de charge pour tester le serveur de modèles. Le serveur est exposé aux requêtes externes à l'aide du service LoadBalancer.
   Un script Python a été fourni avec un exemple de requête pour tester le serveur de modèles.

Le schéma suivant montre comment les requêtes sont acheminées par l'équilibreur de charge.

Remarque :Ces modèles ne sont que des exemples de référence et ne sont pas optimisés pour les performances. Leur utilisation avec des services d'analyse comparative ou de production n'est pas recommandée.

Exemples de déploiement de serveur

Ces exemples de charges de travail sont configurés selon les hypothèses suivantes:

• Le cluster s'exécute avec un pool de nœuds TPU v5 doté de trois nœuds
• Le pool de nœuds utilise le type de machine ct5lp-hightpu-1t où :
  • la topologie est de 1 x 1
  • le nombre de puces TPU est de 1

Le fichier manifeste GKE suivant définit un seul déploiement sur le serveur hôte.

apiVersion: apps/v1
kind: Deployment
metadata:
  name: bert-deployment
spec:
  selector:
    matchLabels:
      app: tf-bert-server
    replicas: 3 # number of nodes in node pool
    template:
      metadata:
        annotations:
          gke-gcsfuse/volumes: "true"
        labels:
          app: tf-bert-server
      spec:
        nodeSelector:
          cloud.google.com/gke-tpu-topology: 1x1 # target topology
          cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        containers:
          name: serve-bert
          image: us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
          securityContext:
            privileged: true
          env:
            name: MODEL_NAME
            value: 'bert'
          volumeMounts:
            mountPath: "/models/"
            name: bert-external-storage
          ports:
            containerPort: 8500
            containerPort: 8501
            containerPort: 8431 # Port to export TPU usage metrics, if supported.
          resources:
            requests:
              google.com/tpu: 1 # TPU chip request
            limits:
              google.com/tpu: 1 # TPU chip request
        volumes:
          - name: bert-external-storage
            persistentVolumeClaim:
              claimName: external-storage-pvc

Si vous utilisez un nombre différent de nœuds dans votre pool de nœuds TPU, remplacez le champ replicas par le nombre de nœuds.

Si vous utilisez un autre type de machine:

• Définissez cloud.google.com/gke-tpu-topology sur la topologie correspondant au type de machine que vous utilisez.
• Définissez les deux champs google.com/tpu sous resources pour qu'ils correspondent au nombre de puces pour le type de machine correspondant.

Prérequis

Télécharger les ressources du tutoriel

Téléchargez les scripts Python et les fichiers manifestes YAML du tutoriel à l'aide de la commande suivante:

gsutil -m cp -r "gs://gke-tpu-examples/single-host-inference" .

Accédez au répertoire single-host-inference :

cd single-host-inference/

Configurer l'environnement Python

Les scripts Python que vous utilisez dans ce tutoriel nécessitent Python version 3.9 ou ultérieure. N'oubliez pas d'installer requirements.txt pour chaque tutoriel avant d'exécuter les scripts de test Python.

Si vous n'avez pas la configuration Python appropriée dans votre environnement local, vous pouvez utiliser Google Cloud Shell pour télécharger et exécuter les scripts Python de ce tutoriel.

Configurer le cluster

1. Créez un cluster à l'aide du type de machine e2-standard-4.

   gcloud container clusters create cluster-name \
   --region your-region \
   --release-channel rapid \
   --num-nodes=1 \
   --machine-type=e2-standard-4 \
   --workload-pool=project-id.svc.id.goog \
   --addons GcsFuseCsiDriver
   

2. Créez le pool de nœuds TPU à hôte unique.

Les exemples de charges de travail supposent que:

• Votre cluster s'exécute avec un pool de nœuds TPU v5e disposant de trois nœuds.
• Le pool de nœuds TPU utilise le type de machine ct5lp-hightpu-1t.

Si vous utilisez une configuration de cluster différente de celle décrite précédemment, vous devez modifier le fichier manifeste de déploiement du serveur.

Pour la démonstration de diffusion stable de JAX, vous aurez besoin d'un pool de nœuds de processeur avec un type de machine disposant de 16 Gio de mémoire disponible (par exemple, e2-standard-4). Pour ce faire, utilisez la commande gcloud container clusters create ou ajoutez un pool de nœuds supplémentaire au cluster existant à l'aide de la commande suivante:

gcloud beta container node-pools create your-pool-name \
  --zone=your-cluster-zone \
  --cluster=your-cluster-name \
  --machine-type=e2-standard-4 \
  --num-nodes=1

Remplacez les éléments suivants :

• your-pool-name : nom du pool de nœuds à créer
• your-cluster-zone: zone dans laquelle votre cluster a été créé.
• your-cluster-name: nom du cluster dans lequel ajouter le pool de nœuds.
• your-machine-type: type de machine des nœuds à créer dans votre pool de nœuds

Configurer le stockage des modèles

Il existe plusieurs façons de stocker votre modèle en vue de l'inférence. Dans ce tutoriel, nous allons utiliser l'approche suivante:

• Pour convertir le modèle pré-entraîné afin qu'il fonctionne sur des TPU, nous allons utiliser un cloud privé virtuel sauvegardé sur un disque persistant avec un accès RWX (ReadWriteMany).
• Pour diffuser le modèle sur plusieurs TPU à hôte unique, nous utiliserons le même VPC reposant sur le bucket Cloud Storage.

Exécutez la commande suivante pour créer un bucket Cloud Storage.

gcloud storage buckets create gs://your-bucket-name \
  --project=your-bucket-project-id \
  --location=your-bucket-location

Remplacez les éléments suivants :

• your-bucket-name: nom du bucket Cloud Storage.
• your-bucket-project-id: ID du projet dans lequel vous avez créé le bucket Cloud Storage.
• your-bucket-location: emplacement de votre bucket Cloud Storage. Pour améliorer les performances, spécifiez l'emplacement d'exécution de votre cluster GKE.

Copiez les scripts Python du tutoriel dans votre propre bucket Cloud Storage à l'aide de la commande suivante:

gsutil -m cp -r "gs://gke-tpu-examples/*" "gs://your-bucket-name"

Remplacez les éléments suivants :

• your-bucket-name: nom du bucket Cloud Storage.

Procédez comme suit pour autoriser votre cluster GKE à accéder au bucket. Pour simplifier la configuration, les exemples suivants utilisent l'espace de noms par défaut et le compte de service Kubernetes par défaut. Pour en savoir plus, consultez la page Configurer l'accès aux buckets Cloud Storage à l'aide de GKE Workload Identity.

1. Créez un compte de service IAM pour votre application ou utilisez un compte de service IAM existant. Vous pouvez utiliser n'importe quel compte de service IAM dans le projet de votre bucket Cloud Storage.

   Remarque :Si vous utilisez un compte de service IAM existant, ignorez cette étape.

   gcloud iam service-accounts create your-iam-service-acct \
   --project=your-bucket-project-id
   

   Remplacez les éléments suivants :

   • your-iam-service-acct: nom du nouveau compte de service IAM.
   • your-bucket-project-id: ID du projet dans lequel vous avez créé votre compte de service IAM. Le compte de service IAM doit se trouver dans le même projet que votre bucket Cloud Storage.

2. Assurez-vous que votre compte de service IAM dispose des rôles de stockage dont vous avez besoin.

   gcloud storage buckets add-iam-policy-binding gs://your-bucket-name \
   --member "serviceAccount:[e-mail]" \
   --role "roles/storage.objectAdmin"
   

   Remplacez les éléments suivants :

   • your-bucket-name : nom du bucket Cloud Storage.
   • your-iam-service-acct: nom de votre nouveau compte de service IAM.
   • your-bucket-project-id: ID du projet dans lequel vous avez créé votre compte de service IAM.

3. Autorisez le compte de service Kubernetes à emprunter l'identité du compte de service IAM en ajoutant une liaison de stratégie IAM entre les deux comptes de service. Cette liaison permet au compte de service Kubernetes d'agir en tant que compte de service IAM.

   gcloud iam service-accounts add-iam-policy-binding [e-mail] \
     --role roles/iam.workloadIdentityUser \
     --member "serviceAccount:your-project-id.svc.id.goog[default/default]"
   

   Remplacez les éléments suivants :

   • your-iam-service-acct : nom de votre compte de service IAM.
   • your-bucket-project-id: ID du projet dans lequel vous avez créé votre bucket Cloud Storage.
   • your-project-id: ID du projet dans lequel vous avez créé votre cluster GKE. Vos buckets Cloud Storage et votre cluster GKE peuvent se trouver dans le même projet ou dans des projets différents.

4. Annotez le compte de service Kubernetes avec l'adresse e-mail du compte de service IAM.

   kubectl annotate serviceaccount default \
     --namespace default \
     iam.gke.io/gcp-service-account=[e-mail]
   

   Remplacez les éléments suivants :

   • your-iam-service-acct: nom du nouveau compte de service IAM.
   • your-bucket-project: ID du projet dans lequel vous avez créé votre bucket Cloud Storage.

5. Exécutez la commande suivante pour renseigner le nom de votre bucket dans les fichiers YAML de cette démonstration:

   find . -type f -name "*.yaml" | xargs sed -i "s/BUCKET_NAME/your-bucket-name/g"
   

   Remplacez your-bucket-name par le nom de votre bucket Cloud Storage.

6. Créez le volume persistant et la revendication de volume persistant à l'aide de la commande suivante:

   kubectl apply -f pvc-pv.yaml
   

Inférence et inférence de modèle JAX

Installer les exigences de JAX pour Python

Installez des dépendances Python pour exécuter les scripts Python du tutoriel qui envoient des requêtes au service de modèles JAX.

pip install -r jax/requirements.txt

Exécutez la démonstration de diffusion JAX BERT E2E:

Cette démonstration utilise un modèle BERT pré-entraîné de Hugging Face.

Remarque :Vous devez exécuter les instructions de configuration avant de lancer la démonstration de diffusion E2E.

Le pod Kubernetes effectue les étapes suivantes:

1. Il télécharge et utilise le script Python export_bert_model.py à partir des exemples de ressources pour télécharger le modèle BERT pré-entraîné dans un répertoire temporaire.
2. Utilise l'image du convertisseur Cloud TPU pour convertir le modèle pré-entraîné du processeur en TPU et stocke le modèle dans le bucket Cloud Storage que vous avez créé lors de la configuration.

Ce pod Kubernetes est configuré pour s'exécuter sur le processeur du pool de nœuds par défaut. Exécutez le pod à l'aide de la commande suivante:

kubectl apply -f jax/bert/install-bert.yaml

Vérifiez que le modèle a été installé correctement en procédant comme suit:

kubectl get pods install-bert

STATUS peut mettre quelques minutes à lire Completed.

Lancer le serveur de modèles TF pour le modèle

Les exemples de charges de travail de ce tutoriel supposent que:

• Le cluster s'exécute avec un pool de nœuds TPU v5 comportant trois nœuds
• Le pool de nœuds utilise le type de machine ct5lp-hightpu-1t contenant une puce TPU.

Si vous utilisez une configuration de cluster différente de celle décrite précédemment, vous devez modifier le fichier manifeste de déploiement du serveur.

Appliquer le déploiement

kubectl apply -f jax/bert/serve-bert.yaml

Vérifiez que le serveur est en cours d'exécution à l'aide des éléments suivants:

kubectl get deployment bert-deployment

AVAILABLE peut mettre une minute à lire 3.

Appliquer le service d'équilibrage de charge

kubectl apply -f jax/bert/loadbalancer.yaml

Vérifiez que l'équilibreur de charge est prêt pour le trafic externe à l'aide des éléments suivants:

kubectl get svc tf-bert-service<

L'enregistrement d'une adresse IP pour EXTERNAL_IP peut prendre quelques minutes.

Envoyer la requête au serveur de modèles

Obtenez une adresse IP externe à partir du service d'équilibrage de charge:

EXTERNAL_IP=$(kubectl get services tf-bert-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

Exécutez un script pour envoyer une requête au serveur:

python3 jax/bert/bert_request.py $EXTERNAL_IP

Résultat attendu :

For input "The capital of France is [MASK].", the result is ". the capital of france is paris.."
For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."

Effectuer un nettoyage

Pour nettoyer les ressources, exécutez kubectl delete dans l'ordre inverse.

kubectl delete -f jax/bert/loadbalancer.yaml
kubectl delete -f jax/bert/serve-bert.yaml
kubectl delete -f jax/bert/install-bert.yaml

Lancer la démonstration de diffusion de bout en bout de JAX Stable Diffusion

Cette démonstration utilise le modèle de diffusion stable pré-entraîné de Hugging Face.

Remarque :Vous devez exécuter les instructions de configuration avant de lancer la démonstration de diffusion E2E.

Exportation du modèle enregistré TF2 compatible avec le TPU à partir du modèle de diffusion stable Flax

L'exportation des modèles de diffusion stables nécessite que le cluster dispose d'un pool de nœuds de processeur avec un type de machine disposant de plus de 16 Gio de mémoire disponible, comme décrit dans la section Cluster de configuration.

Le pod Kubernetes exécute les étapes suivantes:

1. Télécharge et utilise le script Python export_stable_diffusion_model.py des exemples de ressources pour télécharger le modèle de diffusion stable pré-entraîné dans un répertoire temporaire.
2. Utilise l'image du convertisseur Cloud TPU pour convertir le modèle pré-entraîné du processeur en TPU et stocke le modèle dans le bucket Cloud Storage que vous avez créé lors de la configuration du stockage.

Ce pod Kubernetes est configuré pour s'exécuter sur le pool de nœuds de processeur par défaut. Exécutez le pod à l'aide de la commande suivante:

kubectl apply -f jax/stable-diffusion/install-stable-diffusion.yaml

Vérifiez que le modèle a été installé correctement en procédant comme suit:

kubectl get pods install-stable-diffusion

STATUS peut mettre quelques minutes à lire Completed.

Lancer le conteneur de serveur TF Model pour le modèle

Les exemples de charges de travail ont été configurés selon les hypothèses suivantes:

• Le cluster s'exécute avec un pool de nœuds TPU v5 avec trois nœuds
• Le pool de nœuds utilise le type de machine ct5lp-hightpu-1t avec :
  • la topologie est de 1 x 1
  • le nombre de puces TPU est de 1

Si vous utilisez une configuration de cluster différente de celle décrite précédemment, vous devez modifier le fichier manifeste de déploiement du serveur.

Appliquez le déploiement:

kubectl apply -f jax/stable-diffusion/serve-stable-diffusion.yaml

Vérifiez que le serveur fonctionne comme prévu:

kubectl get deployment stable-diffusion-deployment

AVAILABLE peut mettre une minute à lire 3.

Appliquez le service d'équilibrage de charge:

kubectl apply -f jax/stable-diffusion/loadbalancer.yaml

Vérifiez que l'équilibreur de charge est prêt pour le trafic externe à l'aide des éléments suivants:

kubectl get svc tf-stable-diffusion-service

L'enregistrement d'une adresse IP pour EXTERNAL_IP peut prendre quelques minutes.

Envoyer la requête au serveur de modèles

Obtenez une adresse IP externe à partir de l'équilibreur de charge:

EXTERNAL_IP=$(kubectl get services tf-stable-diffusion-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

Exécutez le script pour envoyer une requête au serveur

python3 jax/stable-diffusion/stable_diffusion_request.py $EXTERNAL_IP

Résultat attendu :

L'invite est Painting of a squirrel skating in New York et l'image de sortie sera enregistrée sous le nom stable_diffusion_images.jpg dans votre répertoire actuel.

Remarque :Le modèle pré-entraîné étant relativement volumineux, le chargement du modèle par le service d'inférence prend environ deux minutes. Si des erreurs surviennent lorsque vous envoyez une requête au serveur de modèle, veuillez réessayer après deux minutes.

Effectuer un nettoyage

Pour nettoyer les ressources, exécutez kubectl delete dans l'ordre inverse.

kubectl delete -f jax/stable-diffusion/loadbalancer.yaml
kubectl delete -f jax/stable-diffusion/serve-stable-diffusion.yaml
kubectl delete -f jax/stable-diffusion/install-stable-diffusion.yaml

Inférence et inférence de modèles TensorFlow

Installer les exigences de TensorFlow pour Python

Installez des dépendances Python pour exécuter les scripts Python du tutoriel qui envoient des requêtes au service de modèle TF.

pip install -r tf/resnet50/requirements.txt

Exécutez la démonstration de diffusion de bout en bout TensorFlow ResNet-50:

Étape 1: Convertissez le modèle

Appliquer la conversion du modèle:

kubectl apply -f tf/resnet50/model-conversion.yml

Vérifiez que le modèle a été installé correctement en procédant comme suit:

kubectl get pods resnet-model-conversion

STATUS peut mettre quelques minutes à lire Completed.

Étape 2: Diffuser le modèle avec TensorFlow Serving

Appliquez le déploiement d'inférence du modèle:

kubectl apply -f tf/resnet50/deployment.yml

Vérifiez que le serveur fonctionne comme prévu à l'aide de la commande suivante:

kubectl get deployment resnet-deployment

AVAILABLE peut mettre une minute à lire 3.

Appliquez le service d'équilibrage de charge:

kubectl apply -f tf/resnet50/loadbalancer.yml

Vérifiez que l'équilibreur de charge est prêt pour le trafic externe à l'aide des éléments suivants:

kubectl get svc resnet-service

L'enregistrement d'une adresse IP pour EXTERNAL_IP peut prendre quelques minutes.

Étape 3: Envoyer une requête de test au serveur de modèles

Obtenez l'adresse IP externe à partir de l'équilibreur de charge:

EXTERNAL_IP=$(kubectl get services resnet-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

Exécutez le script de requête de test (HTTP) pour envoyer la requête au serveur de modèles.

python3 tf/resnet50/request.py --host $EXTERNAL_IP

La réponse doit se présenter comme suit:

Predict result: ['ImageNet ID: n07753592, Label: banana, Confidence: 0.94921875',
'ImageNet ID: n03532672, Label: hook, Confidence: 0.0223388672', 'ImageNet ID: n07749582,
Label: lemon, Confidence: 0.00512695312

Étape 4: Nettoyage

Pour nettoyer les ressources, exécutez les commandes kubectl delete suivantes:

kubectl delete -f tf/resnet50/loadbalancer.yml
kubectl delete -f tf/resnet50/deployment.yml
kubectl delete -f tf/resnet50/model-conversion.yml

Assurez-vous de supprimer le pool de nœuds et le cluster GKE lorsque vous n'en avez plus besoin.

Inférence et inférence de modèles PyTorch

Configuration requise pour PyTorch

Installez les dépendances Python pour exécuter les scripts Python du tutoriel qui envoient des requêtes au service de modèles PyTorch:

pip install -r pt/densenet161/requirements.txt

Exécutez la démonstration de diffusion E2E de TorchServe Densenet161:

1. Générer une archive de modèle.

   1. Appliquer l'archive du modèle:

   kubectl apply -f pt/densenet161/model-archive.yml
   

   1. Vérifiez que le modèle a été installé correctement en procédant comme suit:

   kubectl get pods densenet161-model-archive
   

   STATUS peut mettre quelques minutes à lire Completed.

2. Diffuser le modèle avec TorchServe:

   1. Appliquez le déploiement de l'inférence du modèle:

      kubectl apply -f pt/densenet161/deployment.yml
      

   2. Vérifiez que le serveur fonctionne comme prévu à l'aide de la commande suivante:

      kubectl get deployment densenet161-deployment
      

      AVAILABLE peut mettre une minute à lire 3.

   3. Appliquez le service d'équilibrage de charge:

      kubectl apply -f pt/densenet161/loadbalancer.yml
      

      Vérifiez que l'équilibreur de charge est prêt pour le trafic externe à l'aide de la commande suivante:

      kubectl get svc densenet161-service
      

      L'enregistrement d'une adresse IP pour EXTERNAL_IP peut prendre quelques minutes.

3. Envoyez une requête de test au serveur de modèles:

   1. Obtenez une adresse IP externe à partir de l'équilibreur de charge:

      EXTERNAL_IP=$(kubectl get services densenet161-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')
      

   2. Exécutez le script de requête de test pour envoyer la requête (HTTP) au serveur de modèles :

      python3 pt/densenet161/request.py --host $EXTERNAL_IP
      

      Une réponse semblable à celle-ci doit s'afficher:

      Request successful. Response: {'tabby': 0.47878125309944153, 'lynx': 0.20393909513950348, 'tiger_cat': 0.16572578251361847, 'tiger': 0.061157409101724625, 'Egyptian_cat': 0.04997897148132324
      

4. Nettoyez les ressources en exécutant les commandes kubectl delete suivantes:

   kubectl delete -f pt/densenet161/loadbalancer.yml
   kubectl delete -f pt/densenet161/deployment.yml
   kubectl delete -f pt/densenet161/model-archive.yml
   

   Assurez-vous de supprimer le pool de nœuds et le cluster GKE lorsque vous n'en avez plus besoin.

Résoudre les problèmes courants

Dépannage de GKE

Vous trouverez des informations de dépannage pour GKE sur la page Dépannage de GKE.

Échec de l'initialisation du TPU

Si vous rencontrez l'erreur suivante, vérifiez que vous exécutez votre conteneur TPU en mode privilégié ou que vous avez augmenté la valeur ulimit dans votre conteneur. Pour en savoir plus, consultez la section Exécution sans mode privilégié.

TPU platform initialization failed: FAILED_PRECONDITION: Couldn't mmap: Resource
temporarily unavailable.; Unable to create Node RegisterInterface for node 0,
config: device_path:      "/dev/accel0" mode: KERNEL debug_data_directory: ""
dump_anomalies_only: true crash_in_debug_dump: false allow_core_dump: true;
could not create driver instance

Interblocage de la programmation

Supposons que vous ayez deux tâches (tâche A et tâche B) et qu'elles doivent toutes deux être planifiées sur des tranches de TPU avec une topologie de TPU donnée (par exemple, v4-32). Supposons également que vous ayez deux tranches de TPU v4-32 dans le cluster GKE. Nous appellerons ces tranches X et Y. Étant donné que votre cluster dispose d'une grande capacité pour planifier les deux tâches, en théorie, les deux tâches doivent être planifiées rapidement : une tâche sur chacune des deux tranches de TPU v4-32.

Cependant, sans planification minutieuse, il est possible de rencontrer un interblocage dans la planification. Supposons que le programmeur Kubernetes planifie un pod de la tâche A sur la tranche X, puis planifie un pod de la tâche B sur la tranche X. Dans ce cas, compte tenu des règles d'affinité des pods pour la tâche A, le programmeur tentera de planifier tous les pods restants pour la tâche A sur la tranche X. Identique pour la tâche B. Ainsi, ni la tâche A, ni la tâche B ne pourront être entièrement planifiées sur une seule tranche. Il en résulte un interblocage de la planification.

Pour éviter le risque d'interblocage de la planification, vous pouvez utiliser l'anti-affinité de pod avec cloud.google.com/gke-nodepool comme topologyKey, comme illustré dans l'exemple ci-dessous:

apiVersion: batch/v1
kind: Job
metadata:
 name: pi
spec:
 parallelism: 2
 template:
   metadata:
     labels:
       job: pi
   spec:
     affinity:
       podAffinity:
         requiredDuringSchedulingIgnoredDuringExecution:
         - labelSelector:
             matchExpressions:
             - key: job
               operator: In
               values:
               - pi
           topologyKey: cloud.google.com/gke-nodepool
       podAntiAffinity:
         requiredDuringSchedulingIgnoredDuringExecution:
         - labelSelector:
             matchExpressions:
             - key: job
               operator: NotIn
               values:
               - pi
           topologyKey: cloud.google.com/gke-nodepool
           namespaceSelector:
             matchExpressions:
             - key: kubernetes.io/metadata.name
               operator: NotIn
               values:
               - kube-system
     containers:
     - name: pi
       image: perl:5.34.0
       command: ["sleep",  "60"]
     restartPolicy: Never
 backoffLimit: 4

Créer des ressources de pool de nœuds TPU avec Terraform

En plus du SDK gcloud, vous pouvez également utiliser Terraform pour gérer les ressources de votre cluster et de votre pool de nœuds.

Créer un pool de nœuds de tranche TPU multi-hôtes dans un cluster GKE existant

Si vous disposez d'un cluster existant dans lequel vous souhaitez créer un pool de nœuds TPU multi-hôtes, vous pouvez utiliser l'extrait Terraform suivant:

resource "google_container_cluster" "cluster_multi_host" {
  …
  release_channel {
    channel = "RAPID"
  }
  workload_identity_config {
    workload_pool = "my-gke-project.svc.id.goog"
  }
  addons_config {
    gcs_fuse_csi_driver_config {
      enabled = true
    }
  }
}

resource "google_container_node_pool" "multi_host_tpu" {
  provider           = google-beta
  project            = "your-project"
  name               = "your-node-pool"
  location           = "us-central2"
  node_locations     = ["us-central2-b"]
  cluster            = google_container_cluster.cluster_multi_host.name
  initial_node_count = 2

  node_config {
    machine_type = "ct4p-hightpu-4t"
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = ["your-reservation-name"]
    }
    workload_metadata_config {
      mode = "GKE_METADATA"
    }
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = "2x2x2"
  }
}

Remplacez les valeurs suivantes :

• your-project: projet Google Cloud dans lequel vous exécutez votre charge de travail.
• your-node-pool: nom du pool de nœuds que vous créez.
• us-central2: région dans laquelle vous exécutez votre charge de travail.
• us-central2-b: zone dans laquelle vous exécutez votre charge de travail.
• your-reservation-name : nom de votre réservation.

Créer un pool de nœuds de tranche TPU à hôte unique dans un cluster GKE existant

Utilisez l'extrait Terraform suivant:

resource "google_container_cluster" "cluster_single_host" {
  …
  cluster_autoscaling {
    autoscaling_profile = "OPTIMIZE_UTILIZATION"
  }
  release_channel {
    channel = "RAPID"
  }
  workload_identity_config {
  workload_pool = "your-project-id.svc.id.goog"
  }
  addons_config {
    gcs_fuse_csi_driver_config {
      enabled = true
    }
  }
}

resource "google_container_node_pool" "single_host_tpu" {
  provider           = google-beta
  project            = "your-project"
  name               = "your-node-pool"
  location           = "us-central2"
  node_locations     = ["us-central2-b"]
  cluster            = google_container_cluster.cluster_single_host.name
  initial_node_count = 0
  autoscaling {
    total_min_node_count = 2
    total_max_node_count = 22
    location_policy      = "ANY"
  }

  node_config {
    machine_type = "ct4p-hightpu-4t"
    workload_metadata_config {
      mode = "GKE_METADATA"
    }
  }
}

Remplacez les valeurs suivantes :

• your-project: projet Google Cloud dans lequel vous exécutez votre charge de travail.
• your-node-pool: nom du pool de nœuds que vous créez.
• us-central2: région dans laquelle vous exécutez votre charge de travail.
• us-central2-b: zone dans laquelle vous exécutez votre charge de travail.