Déployer des charges de travail TPU sur GKE Autopilot


Cette page explique comment accélérer les charges de travail de machine learning (ML) en utilisant des accélérateurs Cloud TPU (TPU) dans des clusters Google Kubernetes Engine (GKE) en mode Autopilot. Ces conseils peuvent vous aider à sélectionner les bibliothèques appropriées pour vos frameworks d'application de ML, à configurer vos charges de travail TPU pour qu'elles s'exécutent de manière optimale sur GKE et à surveiller vos charges de travail après le déploiement.

Cette page s'adresse aux administrateurs et opérateurs de plate-forme, aux spécialistes des données et de l'IA, ainsi qu'aux développeurs d'applications qui souhaitent préparer et exécuter des charges de travail de ML sur des TPU. Pour en savoir plus sur les rôles, les responsabilités et les exemples de tâches courants que nous citons dans le contenu Google Cloud , consultez Rôles utilisateur et tâches courantes de GKE Enterprise.

Avant de lire cette page, assurez-vous de connaître les ressources suivantes :

Fonctionnement des TPU dans Autopilot

Pour utiliser des TPU dans des charges de travail Autopilot, vous devez spécifier les éléments suivants dans le fichier manifeste de la charge de travail :

  • Version du TPU dans le champ spec.nodeSelector.
  • La topologie de TPU dans le champ spec.nodeSelector. La topologie doit être compatible avec la version de TPU spécifiée.
  • Le nombre de puces TPU dans les champs spec.containers.resources.requests et spec.containers.resources.limits.

Lorsque vous déployez la charge de travail, GKE provisionne les nœuds disposant de la configuration TPU demandée et planifie vos pods sur ces nœuds. GKE place chaque charge de travail sur son propre nœud afin que chaque pod puisse accéder à toutes les ressources du nœud avec un risque d'interruption réduit au minimum.

Les TPU dans Autopilot sont compatibles avec les fonctionnalités suivantes :

  1. Pods Spot
  2. Réservations de capacité spécifiques
  3. Pods d'exécution étendus
  4. Démarrage Flex

Planifier la configuration TPU

Avant d'utiliser ce guide pour déployer des charges de travail TPU, planifiez la configuration de votre TPU en fonction de votre modèle et de la quantité de mémoire dont il a besoin. Pour en savoir plus, consultez Planifier la configuration de TPU.

Tarifs

Pour en savoir plus sur la tarification, consultez la page Tarifs du mode Autopilot.

Avant de commencer

Avant de commencer, effectuez les tâches suivantes :

  • Activez l'API Google Kubernetes Engine.
  • Activer l'API Google Kubernetes Engine
  • Si vous souhaitez utiliser Google Cloud CLI pour cette tâche, installez puis initialisez gcloud CLI. Si vous avez déjà installé gcloud CLI, assurez-vous de disposer de la dernière version en exécutant la commande gcloud components update.
  • Assurez-vous de disposer d'un cluster Autopilot exécutant la version 1.32.3-gke.1927000 ou ultérieure de GKE. Pour obtenir des instructions, consultez Créer un cluster Autopilot.
  • Pour utiliser des TPU réservés, assurez-vous de disposer d'une réservation de capacité spécifique existante. Pour obtenir des instructions, consultez la section Consommer des ressources zonales réservées.

Vérifier le quota pour les TPU et les autres ressources GKE

Les sections suivantes vous aident à vous assurer que vous disposez d'un quota suffisant lorsque vous utilisez des TPU dans GKE.

Pour créer des nœuds de tranche TPU, vous devez disposer d'un quota de TPU, sauf si vous utilisez une réservation de capacité existante. Si vous utilisez des TPU réservés, ignorez cette section.

La création de nœuds de tranche TPU dans GKE nécessite un quota d'API Compute Engine (compute.googleapis.com) et non le quota de l'API Cloud TPU (tpu.googleapis.com). Le nom du quota est différent dans les pods Autopilot standards et dans les pods Spot.

Pour vérifier la limite et l'utilisation actuelle de votre quota d'API Compute Engine pour les TPU, procédez comme suit :

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

    Accéder à la section "Quotas"

  2. Dans la zone  Filtre, 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 de la valeur dans le sélecteur de nœud cloud.google.com/gke-tpu-accelerator. Par exemple, si vous envisagez de créer des nœuds TPU v5e à la demande dont la valeur dans le sélecteur de nœud cloud.google.com/gke-tpu-accelerator est tpu-v5-lite-podslice, saisissez Name: TPU v5 Lite PodSlice chips.

      Version du TPU, cloud.google.com/gke-tpu-accelerator Propriété et nom du quota pour les instances à la demande Propriété et nom du quota pour les instances Spot2
      TPU v3,
      tpu-v3-device
      Dimensions (e.g. location):
      tpu_family:CT3
      Non applicable
      TPU v3,
      tpu-v3-slice
      Dimensions (e.g. location):
      tpu_family:CT3P
      Non applicable
      TPU v4,
      tpu-v4-podslice
      Name:
      TPU v4 PodSlice chips
      Name:
      Preemptible TPU v4 PodSlice chips
      TPU v5e,
      tpu-v5-lite-podslice
      Name:
      TPU v5 Lite PodSlice chips
      Name:
      Preemptible TPU v5 Lite Podslice
      chips
      TPU v5p,
      tpu-v5p-slice
      Name:
      TPU v5p chips
      Name:
      Preemptible TPU v5p chips
      TPU Trillium,
      tpu-v6e-slice
      Dimensions (e.g. location):
      tpu_family:CT6E
      Name:
      Preemptible TPU slices v6e
    2. Sélectionnez la propriété Dimensions (par exemple, emplacements) et saisissez region: suivi du nom de la région dans laquelle vous prévoyez de créer des TPU dans GKE. Par exemple, saisissez region:us-west4 si vous envisagez de créer des nœuds de tranche TPU dans 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.

Si aucun quota ne correspond au filtre que vous avez saisi, le projet ne dispose d'aucun quota pour la région dont vous avez besoin et vous devez demander un ajustement de quota TPU.

Lors de la création d'une réservation TPU, les valeurs de limite et d'utilisation actuelle du quota correspondant augmentent du nombre de puces dans la réservation TPU. Par exemple, lorsqu'une réservation est créée pour 16 puces TPU v5e dont la valeur dans le sélecteur de nœuds cloud.google.com/gke-tpu-accelerator est tpu-v5-lite-podslice, les valeurs Limite et Utilisation actuelle pour le quota TPU v5 Lite PodSlice chips dans la région concernée augmentent de 16.

Quotas pour les ressources GKE supplémentaires

Vous devrez peut-être augmenter les quotas suivants liés à GKE dans les régions où GKE crée vos ressources.

  • Quota SSD Persistent Disk (Go) : 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 aussi haut que le produit du nombre maximal de nœuds GKE que vous prévoyez de créer et 100 Go (nœuds * 100 Go).
  • Quota d'adresses IP en cours d'utilisation : chaque nœud Kubernetes consomme une adresse IP. Par conséquent, ce quota doit être défini au moins aussi haut que le nombre maximal de nœuds GKE que vous prévoyez de créer.
  • Assurez-vous que max-pods-per-node correspond à la plage de sous-réseau : chaque nœud Kubernetes utilise des plages d'adresses IP secondaires pour les pods. Par exemple, max-pods-per-node sur 32 nécessite 64 adresses IP, ce qui se traduit par un sous-réseau /26 par nœud. Notez que cette plage ne doit pas être partagée avec un autre cluster. Pour éviter d'épuiser la plage d'adresses IP, utilisez l'option --max-pods-per-node pour limiter le nombre de pods pouvant être planifiés sur un nœud. Le quota pour max-pods-per-node doit être défini au moins aussi haut que le nombre maximal de nœuds GKE que vous prévoyez de créer.

Pour demander une augmentation de quota, consultez Demander un ajustement de quota.

Options de provisionnement des TPU dans GKE

GKE Autopilot vous permet d'utiliser des TPU directement dans des charges de travail individuelles à l'aide de Kubernetes nodeSelectors.

Vous pouvez également demander des TPU à l'aide de classes de calcul personnalisées. Les classes de calcul personnalisées permettent aux administrateurs de plate-forme de définir une hiérarchie de configurations de nœuds pour que GKE puisse les hiérarchiser lors des décisions de scaling des nœuds. Les charges de travail s'exécutent ainsi sur le matériel de votre choix.

Pour obtenir des instructions, consultez la section Provisionner des TPU de manière centralisée avec des classes de calcul personnalisées.

Préparer votre application TPU

Les charges de travail TPU présentent les exigences de préparation suivantes.

  1. Les frameworks tels que JAX, PyTorch et TensorFlow accèdent aux VM des TPU à l'aide de la bibliothèque partagée libtpu. libtpu inclut le compilateur XLA, le logiciel d'environnement d'exécution TPU et le pilote TPU. Chaque version de PyTorch et de JAX nécessite une certaine version de libtpu.so. Pour utiliser des TPU dans GKE, veillez à utiliser les versions suivantes :
    Type de TPU Version de libtpu.so
    TPU Trillium (v6e)
    tpu-v6e-slice
    TPU v5e
    tpu-v5-lite-podslice
    TPU v5p
    tpu-v5p-slice
    • Version jax[tpu] recommandée : 0.4.19 ou version ultérieure.
    • Version de torchxla[tpuvm] recommandée : il est suggéré d'utiliser une version nocturne du 23 octobre 2023.
    TPU v4
    tpu-v4-podslice
    TPU v3
    tpu-v3-slice
    tpu-v3-device
  2. Définissez les variables d'environnement suivantes pour le conteneur qui demande les ressources TPU suivantes :
    • TPU_WORKER_ID : entier unique pour chaque pod. Cet ID indique un ID de nœud de calcul unique dans la tranche de TPU. Les valeurs acceptées pour ce champ sont comprises entre zéro et le nombre de pods moins un.
    • TPU_WORKER_HOSTNAMES : liste de noms d'hôtes ou d'adresses IP de VM TPU séparés par une virgule qui doivent communiquer entre eux au sein de la tranche. Il doit y avoir un nom d'hôte ou une adresse IP pour chaque VM TPU de la tranche. La liste d'adresses IP ou de noms d'hôte est classée et indexée par zéro par le paramètre TPU_WORKER_ID.
    • GKE injecte automatiquement ces variables d'environnement en utilisant un webhook en mutation lorsqu'un job est créé avec completionMode: Indexed, subdomain, parallelism > 1 et en demandant des propriétés google.com/tpu. GKE ajoute un service sans tête afin que les enregistrements DNS soient ajoutés pour les pods qui sauvegardent le service.

Une fois la préparation de la charge de travail terminée, vous pouvez exécuter un job qui utilise des TPU.

Demander des TPU dans une charge de travail

Cette section explique comment créer un job qui demande des TPU dans Autopilot. Pour toute charge de travail nécessitant des TPU, vous devez spécifier les éléments suivants :

  • Sélecteurs de nœuds pour la version et la topologie du TPU
  • Nombre de puces TPU pour un conteneur dans votre charge de travail

Pour obtenir la liste des versions et des topologies de TPU compatibles, ainsi que le nombre correspondant de puces et de nœuds TPU dans une tranche, consultez la page Choisir la version de TPU.

Remarques concernant les requêtes TPU dans les charges de travail

Un seul conteneur d'un pod peut utiliser des TPU. Le nombre de puces demandé par un conteneur doit être égal au nombre de puces TPU associées à un nœud de la tranche. Par exemple, si vous demandez un TPU v5e (tpu-v5-lite-podslice) avec une topologie 2x4, vous pouvez demander l'un des éléments suivants :

  • 4 puces, qui créent deux nœuds multi-hôtes avec quatre puces TPU chacun
  • 8 puces, qui créent un nœud à hôte unique avec huit puces TPU

Pour optimiser votre rentabilité, il est recommandé de toujours utiliser toutes les puces TPU de la tranche que vous demandez. Si vous demandez une tranche multi-hôte de deux nœuds avec quatre puces TPU chacun, vous devez déployer une charge de travail qui s'exécute sur les deux nœuds et consomme les huit puces de la tranche.

Créer une charge de travail qui demande des TPU

Les étapes suivantes permettent de créer un job qui demande des TPU. Si vous avez des charges de travail qui s'exécutent sur des tranches de TPU multi-hôte, vous devez également créer un service sans adresse IP de cluster qui sélectionne votre charge de travail par son nom. Ce service sans tête IP de cluster permet aux pods de différents nœuds de la tranche multi-hôte de communiquer entre eux en mettant à jour la configuration DNS de Kubernetes de sorte qu'elle pointe vers les pods de la charge de travail.

  1. Enregistrez le manifeste suivant sous le nom tpu-autopilot.yaml :

    apiVersion: v1
    kind: Service
    metadata:
      name: headless-svc
    spec:
      clusterIP: None
      selector:
        job-name: tpu-job
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: tpu-job
    spec:
      backoffLimit: 0
      completions: 4
      parallelism: 4
      completionMode: Indexed
      template:
        spec:
          # Optional: Run in GKE Sandbox
          # runtimeClassName: gvisor
          subdomain: headless-svc
          restartPolicy: Never
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: TPU_TYPE
            cloud.google.com/gke-tpu-topology: TOPOLOGY
          containers:
          - name: tpu-job
            image: python:3.10
            ports:
            - containerPort: 8471 # Default port using which TPU VMs communicate
            - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
            command:
            - bash
            - -c
            - |
              pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
              python -c 'import jax; print("TPU cores:", jax.device_count())'
            resources:
              requests:
                cpu: 10
                memory: MEMORY_SIZE
                google.com/tpu: NUMBER_OF_CHIPS
              limits:
                cpu: 10
                memory: MEMORY_SIZE
                google.com/tpu: NUMBER_OF_CHIPS
    

    Remplacez les éléments suivants :

    • TPU_TYPE : type de TPU à utiliser, tel que tpu-v4-podslice. Il doit s'agir d'une valeur compatible avec GKE.
    • TOPOLOGY : disposition des puces TPU dans la tranche, par exemple 2x2x4. Il doit s'agir d'une topologie compatible avec le type de TPU sélectionné.
    • NUMBER_OF_CHIPS : nombre de puces TPU que le conteneur doit utiliser. La valeur doit être identique pour limits et requests.
    • MEMORY_SIZE : quantité maximale de mémoire utilisée par le TPU. Les limites de mémoire dépendent de la version et de la topologie de TPU que vous utilisez. Pour en savoir plus, consultez Valeurs minimales et maximales pour les accélérateurs.
    • ** Facultatif runtimeClassname: gvisor : paramètre qui vous permet d'exécuter ce pod dans GKE Sandbox. Pour l'utiliser, annulez la mise en commentaire de cette ligne. GKE Sandbox est compatible avec les TPU version v4 et ultérieures. Pour en savoir plus, consultez GKE Sandbox.
  2. Déployez le job :

    kubectl create -f tpu-autopilot.yaml
    

    Lorsque vous créez ce job, GKE effectue automatiquement les opérations suivantes :

    1. Il provisionne des nœuds pour exécuter les pods. Selon le type de TPU, la topologie et les demandes de ressources que vous avez spécifiés, ces nœuds sont des tranches à hôte unique ou à plusieurs hôtes.
    2. Il ajoute des rejets aux pods et des tolérances aux nœuds pour éviter que vos autres charges de travail s'exécutent sur les mêmes nœuds que les charges de travail TPU.
  3. Une fois que vous avez terminé cette section, vous pouvez éviter de continuer à payer des frais en supprimant la charge de travail que vous avez créée :

    kubectl delete -f tpu-autopilot.yaml
    

Créer une charge de travail qui demande des TPU et la planification de la collecte

Dans TPU Trillium, vous pouvez utiliser la planification de la collecte pour regrouper les nœuds de tranche TPU. Le regroupement de ces nœuds de tranche TPU permet d'ajuster plus facilement le nombre de répliques pour répondre à la demande de charge de travail. Google Cloud contrôle les mises à jour logicielles pour s'assurer qu'un nombre suffisant de tranches de la collection est toujours disponible pour traiter le trafic.

TPU Trillium est compatible avec la planification de la collecte pour les pools de nœuds à hôte unique et multi-hôtes qui exécutent des charges de travail d'inférence. La planification de la collecte dépend du type de tranche TPU que vous utilisez :

  • Tranche de TPU multi-hôte : GKE regroupe les tranches de TPU multi-hôtes pour former une collection. Chaque pool de nœuds GKE est une réplique de cette collection. Pour définir une collection, créez une tranche TPU multi-hôte et attribuez-lui un nom unique. Pour ajouter d'autres tranches de TPU à la collection, créez un autre pool de nœuds de tranche TPU multi-hôte avec le même nom de collection et le même type de charge de travail.
  • Tranche de TPU à hôte unique : GKE considère l'ensemble du pool de nœuds de tranche de TPU à hôte unique comme une collection. Pour ajouter des tranches de TPU à la collection, vous pouvez redimensionner le pool de nœuds de tranche de TPU à hôte unique.

Pour en savoir plus sur les limites de la planification des collectes, consultez Fonctionnement de la planification des collectes.

Utiliser une tranche de TPU multi-hôte

La planification de la collecte dans les nœuds de tranche TPU multi-hôtes est disponible pour les clusters Autopilot dans la version 1.31.2-gke.1537000 et ultérieures. Les nœuds de tranche TPU multi-hôtes avec une topologie 2x4 ne sont compatibles qu'avec la version 1.31.2-gke.1115000 ou ultérieure. Pour créer des nœuds de tranche TPU multi-hôte et les regrouper dans une collection, ajoutez les libellés Kubernetes suivants à la spécification de votre charge de travail :

  • cloud.google.com/gke-nodepool-group-name : chaque collection doit avoir un nom unique au niveau du cluster. La valeur de l'étiquette cloud.google.com/gke-nodepool-group-name doit respecter les exigences concernant les étiquettes de cluster.
  • cloud.google.com/gke-workload-type: HIGH_AVAILABILITY

    Par exemple, le bloc de code suivant définit une collection avec une tranche de TPU multihôte :

      nodeSelector:
        cloud.google.com/gke-nodepool-group-name: ${COLLECTION_NAME}
        cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
        cloud.google.com/gke-tpu-topology: 4x4
    ...
    

Utiliser une tranche de TPU à hôte unique

La planification de la collecte dans les nœuds de tranche TPU à hôte unique est disponible pour les clusters Autopilot dans la version 1.31.2-gke.1088000 et ultérieure. Pour créer des nœuds de tranche TPU à hôte unique et les regrouper dans une collection, ajoutez le libellé cloud.google.com/gke-workload-type:HIGH_AVAILABILITY dans la spécification de votre charge de travail.

Par exemple, le bloc de code suivant définit une collection avec une tranche de TPU à hôte unique :

  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 2x2
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
  ...

Provisionner des TPU de manière centralisée avec des classes de calcul personnalisées

Pour provisionner des TPU avec une classe de calcul personnalisée qui suit les règles des TPU et déployer la charge de travail, procédez comme suit :

  1. Enregistrez le manifeste suivant sous le nom tpu-compute-class.yaml :

    apiVersion: cloud.google.com/v1
    kind: ComputeClass
    metadata:
      name: tpu-class
    spec:
      priorities:
      - tpu:
          type: tpu-v5-lite-podslice
          count: 4
          topology: 2x4
      - spot: true
        tpu:
          type: tpu-v5-lite-podslice
          count: 4
          topology: 2x4
      - flexStart:
          enabled: true
        tpu:
          type: tpu-v6e-slice
          count: 4
          topology: 2x4
      nodePoolAutoCreation:
        enabled: true
    
  2. Déployez la classe de calcul :

    kubectl apply -f tpu-compute-class.yaml
    

    Pour en savoir plus sur les classes de calcul personnalisées et les TPU, consultez Configuration des TPU.

  3. Enregistrez le manifeste suivant sous le nom tpu-job.yaml :

    apiVersion: v1
    kind: Service
    metadata:
      name: headless-svc
    spec:
      clusterIP: None
      selector:
        job-name: tpu-job
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: tpu-job
    spec:
      backoffLimit: 0
      completions: 4
      parallelism: 4
      completionMode: Indexed
      template:
        spec:
          subdomain: headless-svc
          restartPolicy: Never
          nodeSelector:
            cloud.google.com/compute-class: tpu-class
          containers:
          - name: tpu-job
            image: python:3.10
            ports:
            - containerPort: 8471 # Default port using which TPU VMs communicate
            - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
            command:
            - bash
            - -c
            - |
              pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
              python -c 'import jax; print("TPU cores:", jax.device_count())'
            resources:
              requests:
                cpu: 10
                memory: MEMORY_SIZE
                google.com/tpu: NUMBER_OF_CHIPS
              limits:
                cpu: 10
                memory: MEMORY_SIZE
                google.com/tpu: NUMBER_OF_CHIPS
    

    Remplacez les éléments suivants :

    • NUMBER_OF_CHIPS : nombre de puces TPU que le conteneur doit utiliser. La valeur doit être identique pour limits et requests, et égale à la valeur du champ tpu.count dans la classe de calcul personnalisée sélectionnée.
    • MEMORY_SIZE : quantité maximale de mémoire utilisée par le TPU. Les limites de mémoire dépendent de la version et de la topologie de TPU que vous utilisez. Pour en savoir plus, consultez Valeurs minimales et maximales pour les accélérateurs.
    • NUMBER_OF_CHIPS : nombre de puces TPU que le conteneur doit utiliser. La valeur doit être identique pour limits et requests.
  4. Déployez le job :

    kubectl create -f tpu-job.yaml
    

    Lorsque vous créez ce job, GKE effectue automatiquement les opérations suivantes :

    • Il provisionne des nœuds pour exécuter les pods. Selon le type de TPU, la topologie et les demandes de ressources que vous avez spécifiés, ces nœuds sont des tranches à hôte unique ou à plusieurs hôtes. En fonction de la disponibilité des ressources TPU dans la priorité la plus élevée, GKE peut revenir à des priorités inférieures pour maximiser la disponibilité.
    • Il ajoute des rejets aux pods et des tolérances aux nœuds pour éviter que vos autres charges de travail s'exécutent sur les mêmes nœuds que les charges de travail TPU.

    Pour en savoir plus, consultez À propos des classes de calcul personnalisées.

  5. Une fois que vous avez terminé cette section, vous pouvez éviter de continuer à payer des frais en supprimant les ressources que vous avez créées :

    kubectl delete -f tpu-job.yaml
    

Exemple : Afficher le nombre total de puces TPU dans une tranche multi-hôte

La charge de travail suivante renvoie le nombre de puces TPU sur tous les nœuds d'une tranche TPU multi-hôte. Pour créer une tranche multi-hôte, la charge de travail comporte les paramètres suivants :

  • Version du TPU : TPU v4
  • Topologie : 2x2x4

Cette version et cette sélection de topologie génèrent une tranche multi-hôte.

  1. Enregistrez le manifeste suivant sous le nom available-chips-multihost.yaml :
    apiVersion: v1
    kind: Service
    metadata:
      name: headless-svc
    spec:
      clusterIP: None
      selector:
        job-name: tpu-available-chips
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: tpu-available-chips
    spec:
      backoffLimit: 0
      completions: 4
      parallelism: 4
      completionMode: Indexed
      template:
        spec:
          subdomain: headless-svc
          restartPolicy: Never
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
            cloud.google.com/gke-tpu-topology: 2x2x4
          containers:
          - name: tpu-job
            image: python:3.10
            ports:
            - containerPort: 8471 # Default port using which TPU VMs communicate
            - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
            command:
            - bash
            - -c
            - |
              pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
              python -c 'import jax; print("TPU cores:", jax.device_count())'
            resources:
              requests:
                cpu: 10
                memory: 407Gi
                google.com/tpu: 4
              limits:
                cpu: 10
                memory: 407Gi
                google.com/tpu: 4
  2. Déployez le fichier manifeste :
    kubectl create -f available-chips-multihost.yaml
    

    GKE exécute une tranche de TPU v4 avec quatre VM (tranche TPU multi-hôte). La tranche comporte 16 puces TPU interconnectées.

  3. Vérifiez que la tâche a créé quatre pods :
    kubectl get pods
    

    Le résultat ressemble à ce qui suit :

    NAME                       READY   STATUS      RESTARTS   AGE
    tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
    tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
    tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
    tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
    
  4. Obtenez les journaux de l'un des pods :
    kubectl logs POD_NAME
    

    Remplacez POD_NAME par le nom de l'un des pods créés. Par exemple, tpu-job-podslice-0-5cd8r.

    Le résultat ressemble à ce qui suit :

    TPU cores: 16
    
  5. Facultatif : Supprimez la charge de travail :
    kubectl delete -f available-chips-multihost.yaml
    

Exemple : Afficher les puces TPU dans un seul nœud

La charge de travail suivante est un pod statique qui affiche le nombre de puces TPU associées à un nœud spécifique. Pour créer un nœud à hôte unique, la charge de travail comporte les paramètres suivants :

  • Version du TPU : TPU v5e
  • Topologie : 2x4

Cette version et cette sélection de topologie génèrent une tranche à hôte unique.

  1. Enregistrez le manifeste suivant sous le nom available-chips-singlehost.yaml :
    apiVersion: v1
    kind: Pod
    metadata:
      name: tpu-job-jax-v5
    spec:
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("Total TPU chips:", jax.device_count())'
        resources:
          requests:
            google.com/tpu: 8
          limits:
            google.com/tpu: 8
  2. Déployez le fichier manifeste :
    kubectl create -f available-chips-singlehost.yaml
    

    GKE provisionne des nœuds avec huit tranches de TPU à hôte unique qui utilisent les TPU v5e. Chaque nœud TPU possède huit puces TPU (tranche de TPU à hôte unique).

  3. Récupérez les journaux du pod :
    kubectl logs tpu-job-jax-v5
    

    Le résultat ressemble à ce qui suit :

    Total TPU chips: 8
    
  4. Facultatif : Supprimez la charge de travail :
      kubectl delete -f available-chips-singlehost.yaml
      

Observer et surveiller les TPU

Tableau de bord

L'observabilité des pools de nœuds dans la consoleGoogle Cloud est en disponibilité générale. Pour afficher l'état de vos pools de nœuds TPU multi-hôtes sur GKE, accédez au tableau de bord État des pools de nœuds TPU GKE fourni par Cloud Monitoring :

Accéder à l'état du pool de nœuds TPU GKE

Ce tableau de bord vous fournit des informations complètes sur l'état de vos pools de nœuds TPU multi-hôtes. Pour en savoir plus, consultez Surveiller les métriques sur l'état des nœuds et des pools de nœuds TPU.

Sur la page Clusters Kubernetes de la consoleGoogle Cloud , l'onglet Observabilité affiche également les métriques d'observabilité des TPU, telles que l'utilisation des TPU, sous l'en-tête Accélérateurs > TPU. Pour en savoir plus, consultez Afficher les métriques d'observabilité.

Le tableau de bord TPU n'est renseigné que si les métriques système sont activées dans votre cluster GKE.

Métriques d'exécution

Dans GKE version 1.27.4-gke.900 ou ultérieure, les charges de travail TPU qui utilisent la version JAX 0.4.14 ou ultérieure et spécifient containerPort: 8431 exportent les métriques d'utilisation des TPU en tant que métriques système GKE. Les métriques suivantes sont disponibles dans Cloud Monitoring pour surveiller les performances d'exécution de votre charge de travail TPU :

  • Cycle d'utilisation : durée exprimée en pourcentage de la dernière période d'échantillonnage (60 secondes) pendant laquelle les TensorCores ont traité activement les données sur une puce TPU. Un pourcentage plus élevé signifie une meilleure utilisation du TPU.
  • Mémoire utilisée : quantité de mémoire d'accélérateur allouée, en octets. Cette valeur est échantillonnée toutes les 60 secondes.
  • Mémoire totale : mémoire totale de l'accélérateur, en octets. Échantillonné toutes les 60 secondes.

Ces métriques se trouvent dans le schéma de nœud Kubernetes (k8s_node) et de conteneur Kubernetes (k8s_container).

Conteneur Kubernetes :

  • kubernetes.io/container/accelerator/duty_cycle
  • kubernetes.io/container/accelerator/memory_used
  • kubernetes.io/container/accelerator/memory_total

Nœud Kubernetes :

  • kubernetes.io/node/accelerator/duty_cycle
  • kubernetes.io/node/accelerator/memory_used
  • kubernetes.io/node/accelerator/memory_total

Surveiller les métriques d'état des nœuds et des pools de nœuds TPU

Lorsqu'une tâche d'entraînement présente une erreur ou se termine par un échec, vous pouvez vérifier les métriques liées à l'infrastructure sous-jacente pour déterminer si l'interruption est due à un problème avec le nœud ou le pool de nœuds sous-jacents.

État du nœud

Dans GKE version 1.32.1-gke.1357001 ou ultérieure, la métrique système GKE suivante expose l'état d'un nœud GKE :

  • kubernetes.io/node/status_condition

Le champ condition indique les conditions du nœud, telles que Ready, DiskPressure et MemoryPressure. Le champ status indique l'état signalé de la condition, qui peut être True, False ou Unknown. Il s'agit d'une métrique avec le type de ressource surveillée k8s_node.

Cette requête PromQL indique si un nœud particulier est Ready :

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    node_name="NODE_NAME",
    condition="Ready",
    status="True"}

Pour résoudre les problèmes dans un cluster, vous pouvez examiner les nœuds qui ont présenté d'autres conditions :

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition!="Ready",
    status="True"}

Vous pouvez examiner plus précisément les nœuds qui ne sont pas Ready :

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition="Ready",
    status="False"}

Si aucune donnée n'est disponible, les nœuds sont prêts. L'état est échantillonné toutes les 60 secondes.

Vous pouvez utiliser la requête suivante pour comprendre l'état des nœuds dans l'ensemble du parc :

avg by (condition,status)(
  avg_over_time(
    kubernetes_io:node_status_condition{monitored_resource="k8s_node"}[${__interval}]))

État du pool de nœuds

La métrique système GKE suivante pour la ressource surveillée k8s_node_pool expose l'état d'un pool de nœuds GKE :

  • kubernetes.io/node_pool/status

Cette métrique n'est signalée que pour les pools de nœuds TPU multi-hôtes.

Le champ status indique l'état du pool de nœuds, par exemple Provisioning, Running, Error, Reconciling ou Stopping. Les mises à jour d'état se produisent une fois les opérations de l'API GKE terminées.

Pour vérifier si un pool de nœuds particulier présente l'état Running, utilisez la requête PromQL suivante :

kubernetes_io:node_pool_status{
    monitored_resource="k8s_node_pool",
    cluster_name="CLUSTER_NAME",
    node_pool_name="NODE_POOL_NAME",
    status="Running"}

Pour surveiller le nombre de pools de nœuds dans votre projet, regroupés par état, utilisez la requête PromQL suivante :

count by (status)(
  count_over_time(
    kubernetes_io:node_pool_status{monitored_resource="k8s_node_pool"}[${__interval}]))

Disponibilité du pool de nœuds

La métrique système GKE suivante indique si un pool de nœuds TPU multi-hôte est disponible :

  • kubernetes.io/node_pool/multi_host/available

La métrique a une valeur de True si tous les nœuds du pool de nœuds sont disponibles et False dans le cas contraire. La métrique est échantillonnée toutes les 60 secondes.

Pour vérifier la disponibilité des pools de nœuds TPU multi-hôtes dans votre projet, utilisez la requête PromQL suivante :

avg by (node_pool_name)(
  avg_over_time(
    kubernetes_io:node_pool_multi_host_available{
      monitored_resource="k8s_node_pool",
      cluster_name="CLUSTER_NAME"}[${__interval}]))

Nombre d'interruptions de nœud

La métrique système GKE suivante indique le nombre d'interruptions pour un nœud GKE depuis le dernier échantillon (la métrique est échantillonnée toutes les 60 secondes) :

  • kubernetes.io/node/interruption_count

Les champs interruption_type (par exemple, TerminationEvent, MaintenanceEvent ou PreemptionEvent) et interruption_reason (par exemple, HostError, Eviction ou AutoRepair) peuvent aider à expliquer pourquoi un nœud a été interrompu.

Pour obtenir une répartition des interruptions et de leurs causes dans les nœuds TPU des clusters de votre projet, utilisez la requête PromQL suivante :

  sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node"}[${__interval}]))

Pour n'afficher que les événements de maintenance de l'hôte, mettez à jour la requête afin de filtrer la valeur HW/SW Maintenance pour interruption_reason. Utilisez la requête PromQL suivante :

```promql
sum by (interruption_type,interruption_reason)(
  sum_over_time(
kubernetes_io:node_interruption_count{monitored_resource="k8s_node", interruption_reason="HW/SW Maintenance"}[${__interval}]))
```

Pour afficher le nombre d'interruptions agrégé par pool de nœuds, utilisez la requête PromQL suivante :

```promql
sum by (node_pool_name,interruption_type,interruption_reason)(
  sum_over_time(
    kubernetes_io:node_pool_interruption_count{monitored_resource="k8s_node_pool", interruption_reason="HW/SW Maintenance", node_pool_name=NODE_POOL_NAME }[${__interval}]))
```

Temps de restauration (TTR) du pool de nœuds

La métrique système GKE suivante indique la distribution des durées de période de récupération pour les pools de nœuds TPU multi-hôtes GKE :

  • kubernetes.io/node_pool/accelerator/times_to_recover

Chaque échantillon enregistré dans cette métrique indique un événement de récupération unique pour le pool de nœuds après une période d'indisponibilité.

Cette métrique est utile pour suivre le temps de récupération et le temps entre les interruptions des pools de nœuds TPU multi-hôte.

Vous pouvez utiliser la requête PromQL suivante pour calculer le délai moyen de rétablissement (MTTR) des sept derniers jours dans votre cluster :

sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_sum{
    monitored_resource="k8s_node_pool", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_count{
    monitored_resource="k8s_node_pool",cluster_name="CLUSTER_NAME"}[7d]))

Temps entre interruptions (TBI) du pool de nœuds

Le temps entre les interruptions du pool de nœuds mesure la durée de fonctionnement de votre infrastructure avant qu'elle ne subisse une interruption. Il est calculé comme la moyenne sur une période donnée, où le numérateur mesure la durée totale pendant laquelle votre infrastructure était opérationnelle et le dénominateur mesure le nombre total d'interruptions de votre infrastructure.

L'exemple PromQL suivant montre le temps moyen entre les interruptions (MTBI) sur sept jours pour le cluster donné :

sum(count_over_time(
  kubernetes_io:node_memory_total_bytes{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_interruption_count{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))

Métriques d'hôte

Dans GKE version 1.28.1-gke.1066000 ou ultérieure, les VM dans une tranche TPU exportent les métriques d'utilisation des TPU en tant que métriques système GKE. Les métriques suivantes sont disponibles dans Cloud Monitoring pour surveiller les performances de votre hôte TPU :

  • Utilisation de TensorCore : pourcentage actuel du TensorCore utilisé. La valeur TensorCore correspond à la somme des unités de multiplication de matrice (MXU) plus l'unité vectorielle. La valeur d'utilisation TensorCore correspond à la division des opérations TensorCore effectuées au cours de la période d'échantillonnage précédente (60 secondes) par le nombre d'opérations TensorCore compatibles sur la même période. Plus la valeur est élevée, plus l'utilisation est optimale.
  • Utilisation de la bande passante mémoire : pourcentage actuel de bande passante mémoire utilisée par l'accélérateur. Calculé en divisant la bande passante mémoire utilisée sur une période d'échantillonnage (60 secondes) par la bande passante maximale acceptée sur la même période d'échantillonnage.

Ces métriques se trouvent dans le schéma de nœud Kubernetes (k8s_node) et de conteneur Kubernetes (k8s_container).

Conteneur Kubernetes :

  • kubernetes.io/container/accelerator/tensorcore_utilization
  • kubernetes.io/container/accelerator/memory_bandwidth_utilization

Nœud Kubernetes :

  • kubernetes.io/node/accelerator/tensorcore_utilization
  • kubernetes.io/node/accelerator/memory_bandwidth_utilization

Pour en savoir plus, consultez Métriques de Kubernetes et Métriques système de GKE.

Journalisation

Les journaux émis par les conteneurs s'exécutant sur des nœuds GKE, y compris les VM TPU, sont collectés par l'agent de journalisation GKE, envoyés à Logging et visibles dans Logging.

Recommandations pour les charges de travail TPU dans Autopilot

Les recommandations suivantes peuvent améliorer l'efficacité de vos charges de travail TPU :

  • Utilisez des pods à durée d'exécution prolongée pendant un délai de grâce allant jusqu'à sept jours avant que GKE n'arrête vos pods pour effectuer des scalings à la baisse ou les mises à niveau de nœuds. Vous pouvez utiliser des intervalles de maintenance et des exclusions avec des pods avec une durée d'exécution étendue pour retarder davantage les mises à niveau automatiques des nœuds.
  • Utilisez des réservations de capacité pour vous assurer que vos charges de travail reçoivent les TPU demandés sans être placées dans une file d'attente pour disponibilité.

Pour apprendre à configurer Cloud TPU dans GKE, consultez les ressources suivantes : Google Cloud