Gérer les événements de maintenance avec les pods Cloud TPU

Présentation

Les nœuds TPU et les VM TPU sont des instances de VM Compute Engine associées à du matériel TPU. Les VM Compute Engine sont soumises aux événements de maintenance de VM Compute Engine. Chaque TPU est connecté à une VM Compute Engine. Par conséquent, l'utilisation de plusieurs TPU (par exemple, dans un pod TPU) augmente la probabilité que l'une de vos VM rencontre un événement de maintenance.

Ce document présente différentes approches de gestion des événements de maintenance pour les tâches d'entraînement de longue durée sur Cloud TPU.

Utiliser des points de contrôle pour une récupération rapide après des événements de maintenance

Les points de contrôle sont essentiels pour récupérer rapidement les données après un événement de maintenance et doivent être enregistrés fréquemment. Il est recommandé de les enregistrer environ toutes les heures. L'absence de points de contrôle fréquents risque de causer la perte d'une grande partie de la progression de l'entraînement en raison d'événements de maintenance ou d'autres interruptions.

Les points de contrôle font généralement référence à tous les paramètres enregistrés utilisés lors de l'entraînement (tels que les pondérations du modèle). Le temps nécessaire à l'enregistrement d'un point de contrôle est de quelques secondes à quelques minutes.

Bien que la plupart des événements de maintenance soient automatiquement récupérés et que les tâches d'entraînement se poursuivent sans intervention manuelle, il peut arriver que la tâche ne redémarre pas et continue automatiquement. Dans ce cas, vous devez supprimer et recréer les ressources TPU, puis redémarrer la tâche d'entraînement à partir d'un point de contrôle enregistré. Pour en savoir plus sur la détection et la récupération des défaillances de récupération automatique, consultez la section Détecter et récupérer les défaillances de TPU ci-dessous.

Les mécanismes permettant d'enregistrer et de charger des points de contrôle sont différents pour chaque framework de ML. Les modèles Cloud TPU compatibles intègrent généralement des points de contrôle. Pour en savoir plus sur les points de contrôle, consultez les pages TensorFlow 2.x, PyTorch ou JAX/flax.

Détecter des événements de maintenance

Événement de maintenance récent

Vous pouvez détecter si un événement de maintenance se produit sur votre TPU et, le cas échéant, à quel moment à l'aide de la commande gcloud describe suivante:

VM TPU

$ gcloud compute tpus tpu-vm describe tpu-name --zone=zone  | grep 'health'

Nœuds TPU

$ gcloud compute tpus describe tpu-name --zone=zone | grep 'health'

Le résultat de cette commande affiche l'état actuel du TPU et fournit une description du dernier événement de maintenance. Le résultat doit se présenter comme suit:

health: HEALTHY
healthDescription: The TPU had a maintenance event at 2022-01-26T03:44:36.265703305Z

Journaux des événements de maintenance

Vous pouvez consulter l'historique des événements de maintenance de votre TPU dans les journaux d'audit des événements système.

Dans le menu de navigation de la console Google Cloud, cliquez sur Compute Engine > Instances de VM, puis recherchez, par exemple:

"tpu.nodes.terminate" OR "tpu.nodes.restart"

Dans le délai de recherche, toutes les interruptions et réparations de vos nœuds de calcul TPU sont affichées. Les journaux indiquent la date et l'heure de l'événement, le type d'événement et, pour les événements "terminés", le motif de l'arrêt dans protoPayload.metadata.terminateReason.

Gérer les événements de maintenance

Il existe plusieurs façons de limiter les interruptions liées aux événements de maintenance.

  1. Enregistrer régulièrement les points de contrôle

    Dans un scénario idéal, lorsqu'un "événement d'interruption" se produit, l'entraînement reprend simplement à partir du dernier point de contrôle.

  2. Nouvelles tentatives d'exécution de script d'entraînement

    Le script d'entraînement peut s'arrêter en raison d'un "événement d'interruption". Vous pouvez utiliser un script bash pour relancer le script d'entraînement en continu jusqu'à la fin de l'entraînement. Chaque nouvelle tentative doit poursuivre à partir du dernier point de contrôle. Par conséquent, les scripts de nouvelle tentative doivent toujours être utilisés conjointement avec les points de contrôle.

    Les pipelines d'entraînement prêts pour la production doivent utiliser un système de gestion des ressources tel que Google Kubernetes Engine (GKE). GKE est actuellement compatible avec les nœuds TPU, mais pas avec les VM TPU. Sinon, vous pouvez implémenter un script bash pour relancer le script d'entraînement en continu jusqu'à la fin de l'opération. Exemple :

    Avec un nœud TPU :

    (Depuis votre VM) while ! python3 [training command]; do sleep 1; done

    Avec une VM TPU:

    while ! gcloud compute tpus tpu-vm ssh ${TPU_NAME} --command "python3 [training command]"; do sleep 1; done
    

    Notez que vous devez exécuter la commande de VM TPU à partir d'une interface système Google Cloud ou d'un terminal, et non à partir de la VM TPU.

  3. Détecter et résoudre les défaillances de TPU

    Lorsqu'un TPU ne récupère pas d'un événement de maintenance, vous pouvez utiliser un script de récupération pour détecter son état, puis supprimer et recréer le TPU. Vous trouverez un exemple de ce script ici. Consultez la section Gérer les TPU pour plus de détails sur la suppression et la recréation manuelles de TPU.

    Lorsque vous créez ou recréez une VM TPU, vous pouvez spécifier un script de démarrage avec le paramètre --metadata startup-script. Un script de démarrage s'exécute chaque fois qu'une VM TPU est créée. Pour en savoir plus, consultez la section Exécuter des scripts d'installation standards.