Exécuter des applications Cloud TPU sur GKE

Ce guide vous explique comment :

Pour en savoir plus sur les architectures de VM TPU, consultez la page Architecture du système. Ce guide ne peut être utilisé qu'avec l'architecture de nœuds TPU.

Avantages de l'exécution d'applications Cloud TPU sur GKE

Les applications d'entraînement Cloud TPU peuvent être configurées pour s'exécuter dans des conteneurs GKE au sein de pods GKE. Opter pour cette solution procure les avantages suivants :

  • Configuration et gestion simplifiées : GKE gère le cycle de vie des TPU. Une fois que l'initialisation et l'entraînement Cloud TPU sont configurés avec GKE, l'entraînement est plus facilement répété et géré par GKE, y compris pour la reprise après échec d'une tâche.

  • Coût optimisé : vous n'êtes facturé pour le TPU que lorsque la tâche est active. GKE crée et supprime automatiquement les TPU en fonction des besoins en ressources du pod.

  • Flexibilité d'utilisation : vous n'avez que très peu de modifications à apporter dans la spécification de pod pour demander un accélérateur matériel différent (processeur, GPU ou TPU) :

    kind: Pod
    metadata:
      name: example-tpu
      annotations:
        # The Cloud TPUs that will be created for this Job will support
        # TensorFlow 2.6.0. This version MUST match the
        # TensorFlow version that your model is built on.
        tf-version.cloud-tpus.google.com: "2.6.0"
    spec:
      containers:
      - name: example-container
        resources:
          limits:
            cloud-tpus.google.com/v2: 8
            # See the line above for TPU, or below for CPU / GPU.
            # cpu: 2
            # nvidia.com/gpu: 1
    
  • Évolutivité : GKE fournit des API (Job et Deployment) qui peuvent facilement évoluer jusqu'à plusieurs centaines de pods GKE et de nœuds TPU.

  • Tolérance aux pannes : l'API Job de GKE, ainsi que le mécanisme de point de contrôle TensorFlow, fournissent la sémantique d'exécution jusqu'à la fin. Si des échecs surviennent sur les instances de VM ou les nœuds Cloud TPU, les tâches d'entraînement sont automatiquement réexécutées avec le dernier état lu à partir du point de contrôle.

Configuration requise et limites de Cloud TPU et de GKE

Tenez compte des points suivants lors de la définition de la configuration GKE :

  • Cloud TPU n'est pas compatible avec les pools de nœuds Windows Server.
  • Vous devez créer les pools de nœuds et le cluster GKE dans une zone où Cloud TPU est disponible. Vous devez également créer des buckets Cloud Storage pour stocker vos données et modèles d'entraînement dans la même région que votre cluster GKE. Pour obtenir la liste des zones disponibles, consultez le document types et zones.
  • Chaque conteneur ne peut demander qu'une ressource Cloud TPU. Toutefois, plusieurs conteneurs d'un pod peuvent en demander chacun une.

Avant de commencer

  1. Connectez-vous à votre compte Google Cloud. Si vous débutez sur Google Cloud, créez un compte pour évaluer les performances de nos produits en conditions réelles. Les nouveaux clients bénéficient également de 300 $ de crédits gratuits pour exécuter, tester et déployer des charges de travail.
  2. Dans Google Cloud Console, sur la page de sélection du projet, sélectionnez ou créez un projet Google Cloud.

    Accéder au sélecteur de projet

  3. Assurez-vous que la facturation est activée pour votre projet Cloud. Découvrez comment vérifier si la facturation est activée sur un projet.

  4. Dans Google Cloud Console, sur la page de sélection du projet, sélectionnez ou créez un projet Google Cloud.

    Accéder au sélecteur de projet

  5. Assurez-vous que la facturation est activée pour votre projet Cloud. Découvrez comment vérifier si la facturation est activée sur un projet.

  6. Lorsque vous exploitez Cloud TPU avec GKE, votre projet utilise des composants facturables de Google Cloud. Consultez la grille tarifaire de Cloud TPU et celle de GKE pour estimer vos coûts, et suivez les instructions permettant de nettoyer les ressources lorsque vous n'en avez plus besoin.

  7. Activez les API suivantes dans la console:

Créer un cluster compatible avec Cloud TPU

Suivez les instructions ci-dessous pour configurer votre environnement et créer un cluster GKE compatible avec Cloud TPU, à l'aide de la CLI gcloud:

  1. Installez les composants gcloud dont vous avez besoin pour exécuter GKE avec Cloud TPU :

    $ gcloud components install kubectl 
  2. Configurez gcloud avec l'ID de votre projet Google Cloud :

    $ gcloud config set project project-name
    

    Remplacez project-name par le nom de votre projet Google Cloud.

    La première fois que vous exécutez cette commande dans une nouvelle VM Cloud Shell, une page Authorize Cloud Shell s'affiche. Cliquez sur Authorize en bas de la page pour autoriser gcloud à effectuer des appels d'API GCP avec vos identifiants.

  3. Configurez gcloud sur la zone dans laquelle vous prévoyez d'utiliser une ressource Cloud TPU. Cet exemple utilise us-central1-b, mais vous pouvez utiliser un TPU dans n'importe quelle zone compatible.

    $ gcloud config set compute/zone us-central1-b
    
  4. Exécutez la commande gcloud container clusters create pour créer un cluster sur GKE compatible avec Cloud TPU.

    $ gcloud container clusters create cluster-name \
      --release-channel=stable \
      --scopes=cloud-platform \
      --enable-ip-alias \
      --enable-tpu
    

    Description des options de commande

    release-channel
    Les versions disponibles permettent de gérer les mises à niveau automatiques de votre cluster. Lorsque vous créez un cluster, vous pouvez choisir sa version disponible. Votre cluster ne sera mis à niveau que vers les versions proposées dans ce canal.
    scopes
    Cette option garantit que tous les nœuds du cluster ont accès au bucket Cloud Storage. Le cluster et le bucket de stockage doivent toutefois appartenir au même projet. Notez que les pods Kubernetes héritent par défaut des champs d'application des nœuds sur lesquels ils sont déployés. Par conséquent, scopes=cloud-platform donne le champ d'application cloud-platform à tous les pods Kubernetes exécutés dans le cluster. Si vous souhaitez limiter l'accès pod par pod, consultez le guide GKE sur l'authentification avec des comptes de service.
    enable-ip-alias
    Indique que le cluster utilise des plages d'adresses IP d'alias. Ceci est obligatoire pour pouvoir utiliser Cloud TPU sur GKE.
    enable-tpu
    Indique que le cluster doit être compatible avec Cloud TPU.
    tpu-ipv4-cidr (facultatif, non spécifié ci-dessus)
    Indique la plage CIDR à utiliser pour Cloud TPU. Spécifiez la valeur IP_RANGE au format IP/20, par exemple 10.100.0.0/20. Si vous ne définissez pas cette option, une plage CIDR de taille /20 est automatiquement allouée et attribuée.

Une fois le cluster créé, un message de ce type doit s'afficher :

NAME             LOCATION       MASTER_VERSION    MASTER_IP     MACHINE_TYPE   NODE_VERSION      NUM_NODES  STATUS
cluster-name  us-central1-b  1.16.15-gke.4901  34.71.245.25  n1-standard-1  1.16.15-gke.4901  3          RUNNING

Demander une ressource Cloud TPU dans la spécification de pod Kubernetes

Dans la spécification de pod Kubernetes :

  • Vous devez générer vos modèles dans des conteneurs avec la même version de TensorFlow. Consultez les versions compatibles.

  • Spécifiez la ressource Cloud TPU dans la section limits située sous le champ resource dans les spécifications du conteneur.

    Notez que l'unité de la ressource Cloud TPU est le nombre de cœurs Cloud TPU. Le tableau suivant répertorie des exemples de demandes de ressources valides. Consultez la section Types et zones des TPU pour obtenir une liste complète des ressources TPU valides.

    Si la ressource que vous prévoyez d'utiliser est un pod Cloud TPU, veuillez effectuer une demande de quota, car le quota par défaut des pods Cloud TPU est de zéro.

    Demande de ressource Type de ressource Cloud TPU
    cloud-tpus.google.com/v2: 8 Appareil Cloud TPU v2 (8 cœurs)
    cloud-tpus.google.com/preemptible-v2: 8 Appareil Cloud TPU préemptif v2 (8 cœurs)
    cloud-tpus.google.com/v3: 8 Appareil Cloud TPU v3 (8 cœurs)
    cloud-tpus.google.com/preemptible-v3: 8 Appareil Cloud TPU préemptif v3 (8 cœurs)
    cloud-tpus.google.com/v2: 32 Pod Cloud TPU v2-32 (32 cœurs)
    cloud-tpus.google.com/v3: 32 Un pod Cloud TPU v3-32 (32 cœurs)

    Pour savoir comment indiquer des ressources et des limites dans la spécification de pod, consultez la documentation Kubernetes.

L'exemple de spécification de pod présenté ci-dessous demande un TPU Cloud TPU v2-8 préemptif doté de TensorFlow 2.6.0.

La durée de vie des nœuds Cloud TPU dépend des pods Kubernetes qui les demandent. La ressource Cloud TPU est créée à la demande lors de la programmation du pod Kubernetes et elle est recyclée lorsque celui-ci est supprimé.

apiVersion: v1
kind: Pod
metadata:
  name: gke-tpu-pod
  annotations:
     # The Cloud TPUs that will be created for this Job will support
     # TensorFlow 2.6.0. This version MUST match the
     # TensorFlow version that your model is built on.
     tf-version.cloud-tpus.google.com: "2.6.0"
spec:
  restartPolicy: Never
  containers:
  - name: gke-tpu-container
    # The official TensorFlow 2.6.0 image.
    # https://hub.docker.com/r/tensorflow/tensorflow
    image: tensorflow/tensorflow:2.6.0
    command:
    - python
    - -c
    - |
      import tensorflow as tf
      print("Tensorflow version " + tf.__version__)

      tpu = tf.distribute.cluster_resolver.TPUClusterResolver('$(KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS)')
      print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

      tf.config.experimental_connect_to_cluster(tpu)
      tf.tpu.experimental.initialize_tpu_system(tpu)
      strategy = tf.distribute.TPUStrategy(tpu)

      @tf.function
      def add_fn(x,y):
          z = x + y
          return z

      x = tf.constant(1.)
      y = tf.constant(1.)
      z = strategy.run(add_fn, args=(x,y))
      print(z)
    resources:
      limits:
        # Request a single Preemptible v2-8 Cloud TPU device to train the model.
        cloud-tpus.google.com/preemptible-v2: 8

Créer la tâche

Pour créer la tâche dans le cluster GKE et installer kubectl, procédez comme suit :

  1. À l'aide d'un éditeur de texte, créez la spécification de pod example-job.yaml, puis copiez-collez dans ce fichier la spécification de pod présentée ci-dessus.

  2. Exécutez la tâche :

    $ kubectl create -f example-job.yaml
    
    pod "gke-tpu-pod" created

    Cette commande crée la tâche qui planifie automatiquement le pod.

  3. Vérifiez que le pod GKE a été planifié et que les nœuds Cloud TPU ont été provisionnés. Un pod GKE demandant des nœuds Cloud TPU peut donc se retrouver en attente pendant 5 minutes avant d'être exécuté. Vous verrez une sortie semblable à la suivante jusqu'à ce que le pod GKE soit planifié.

    $ kubectl get pods -w
    
    NAME          READY     STATUS    RESTARTS   AGE
    gke-tpu-pod   0/1       Pending   0          1m
    

    Après cinq minutes environ, vous devriez obtenir un message semblable à celui-ci :

    NAME          READY     STATUS              RESTARTS   AGE
    gke-tpu-pod   0/1       Pending             0          21s
    gke-tpu-pod   0/1       Pending             0          2m18s
    gke-tpu-pod   0/1       Pending             0          2m18s
    gke-tpu-pod   0/1       ContainerCreating   0          2m18s
    gke-tpu-pod   1/1       Running             0          2m48s
    gke-tpu-pod   0/1       Completed           0          3m8s
    

    Utilisez la combinaison de touches "Ctrl-C" pour quitter la commande "kubectl get".

    Vous pouvez afficher des informations de journal et récupérer des informations plus détaillées sur chaque pod GKE à l'aide des commandes kubectl suivantes. Par exemple, pour afficher la sortie du journal de votre pod GKE, utilisez :

    $ kubectl logs gke-tpu-pod
    La sortie obtenue doit ressembler à ceci :
    2021-09-24 18:55:25.400699: I tensorflow/core/platform/cpu_feature_guard.cc:142]
    This TensorFlow binary is optimized with oneAPI Deep Neural Network Library
    (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
    To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
    2021-09-24 18:55:25.405947: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
    Initialize GrpcChannelCache for job worker -> {0 -> 10.0.16.2:8470}
    2021-09-24 18:55:25.406058: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
    Initialize GrpcChannelCache for job localhost -> {0 -> localhost:32769}
    2021-09-24 18:55:28.091729: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
    Initialize GrpcChannelCache for job worker -> {0 -> 10.0.16.2:8470}
    2021-09-24 18:55:28.091896: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
    Initialize GrpcChannelCache for job localhost -> {0 -> localhost:32769}
    2021-09-24 18:55:28.092579: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:427]
    Started server with target: grpc://localhost:32769
    Tensorflow version 2.6.0
    Running on TPU  ['10.0.16.2:8470']
    PerReplica:{
      0: tf.Tensor(2.0, shape=(), dtype=float32),
      1: tf.Tensor(2.0, shape=(), dtype=float32),
      2: tf.Tensor(2.0, shape=(), dtype=float32),
      3: tf.Tensor(2.0, shape=(), dtype=float32),
      4: tf.Tensor(2.0, shape=(), dtype=float32),
      5: tf.Tensor(2.0, shape=(), dtype=float32),
      6: tf.Tensor(2.0, shape=(), dtype=float32),
      7: tf.Tensor(2.0, shape=(), dtype=float32)
    }
    

    Pour afficher une description complète du pod GKE, utilisez :

    $ kubectl describe pod gke-tpu-pod
    

    Pour en savoir plus, consultez la page Présentation et débogage des applications.

Créer et intégrer en conteneur votre modèle dans une image Docker

Consultez la section Créer et intégrer en conteneur votre modèle pour en savoir plus sur ce processus.

Activer la compatibilité avec Cloud TPU sur un cluster existant

Pour activer la compatibilité Cloud TPU sur un cluster GKE existant, procédez comme suit dans la CLI Google Cloud:

  1. Activez la compatibilité avec Cloud TPU :

    gcloud beta container clusters update cluster-name --enable-tpu
    

    Remplacez cluster-name par le nom de votre cluster.

  2. Mettez à jour l'entrée kubeconfig :

    gcloud container clusters get-credentials cluster-name
    

Définir une plage CIDR personnalisée

Par défaut, GKE alloue un bloc CIDR de taille /20 pour les TPU provisionnés par le cluster. Vous pouvez spécifier une plage CIDR personnalisée pour la ressource Cloud TPU en exécutant la commande suivante :

gcloud beta container clusters update cluster-name \
  --enable-tpu \
  --tpu-ipv4-cidr 10.100.0.0/20

Remplacez les éléments suivants :

  • cluster-name : nom de votre cluster existant.
  • 10.100.0.0/20 : votre plage CIDR personnalisée.

Utiliser une plage CIDR existante avec le VPC partagé

Cette section est destinée aux utilisateurs qui déploient un cluster dans un réseau VPC partagé qui doivent utiliser une plage CIDR existante. Si ce n'est pas le cas, vous pouvez ignorer cette section.

La commande suivante montre comment créer un cluster à l'aide d'une plage CIDR existante:

gcloud beta container clusters create cluster-name \
  --enable-tpu \
  --project TENANT_PROJECT \
  --zone=us-ZONE \
  --enable-ip-alias \
  --network projects/HOST_PROJECT/global/networks/HOST_NETWORK \
  --subnetwork projects/HOST_PROJECT/regions/REGION/subnetworks/SUBNETWORK \
  --services-secondary-range-name SERVICES_SECONDARY_RANGE \
  --enable-tpu \
  --enable-tpu-service-networking

Remplacez les éléments suivants :

  • cluster-name : nom de votre cluster existant.
  • TENANT_PROJECT : nom de votre projet.
  • HOST_PROJECT : nom du projet hôte.
  • HOST_NETWORK : nom complet du réseau hôte.
  • HOST_PROJECT: nom complet du sous-réseau.
  • SERVICES_SECONDARY_RANGE : plage de mise en réseau de services existante.

Consultez le Guide de référence des commandes gcloud beta container clusters create pour en savoir plus sur ces indicateurs.

Désactiver Cloud TPU dans un cluster

Pour désactiver la compatibilité Cloud TPU sur un cluster GKE existant, procédez comme suit dans la CLI Google Cloud:

  1. Vérifiez qu'aucune de vos charges de travail n'utilise Cloud TPU :

    $ kubectl get tpu
    
  2. Désactivez la prise en charge de Cloud TPU dans votre cluster :

    $ gcloud beta container clusters update cluster-name --no-enable-tpu
    

    Remplacez cluster-name par le nom de votre cluster.

    Pour les clusters zonaux, cette opération prend environ cinq minutes. Pour les clusters régionaux, elle prend environ 15 minutes, selon la région du cluster.

  3. Une fois les opérations terminées sans erreur, vous pouvez vérifier que les TPU provisionnés par le cluster ont été supprimés :

    $ gcloud compute tpus list
    

    Les noms des TPU créés par Cloud TPU sont au format suivant :

    $ gke-cluster-name-cluster-id-tpu-tpu-id
    

    Remplacez les éléments suivants :

    • cluster-name : nom de votre cluster existant.
    • cluster-id : ID de votre cluster existant.
    • tpu-id : ID de la ressource Cloud TPU.

    Si des TPU apparaissent, vous pouvez les supprimer manuellement en exécutant la commande suivante :

    $ gcloud compute tpus delete gke-cluster-name-cluster-id-tpu-tpu-id
    

Effectuer un nettoyage

Une fois que vous n'avez plus besoin d'utiliser Cloud TPU sur GKE, nettoyez les ressources pour éviter que des frais supplémentaires ne soient imputés à votre compte de facturation Cloud.

Si vous n'avez pas défini le projet et la zone pour cette session, faites-le maintenant. Consultez les instructions plus haut dans ce guide. Suivez ensuite cette procédure de nettoyage :

  1. Exécutez la commande suivante pour supprimer votre cluster GKE, en remplaçant cluster-name par le nom de votre cluster et "project-name" par le nom de votre projet Google Cloud :

    $ gcloud container clusters delete cluster-name \
    --project=project-name --zone=us-central1-b
    
  2. Lorsque vous avez terminé l'examen des données, exécutez la commande gsutil pour supprimer le bucket Cloud Storage que vous avez créé. Remplacez bucket-name par le nom de votre bucket Cloud Storage.

    $ gsutil rm -r gs://bucket-name