Administra eventos de mantenimiento con pods de Cloud TPU

Descripción general

Las VMs de TPU son instancias de VMs de Compute Engine con hardware de TPU conectado. Las VMs de Compute Engine están sujetas a eventos de mantenimiento de VMs de Compute Engine. Cada TPU está conectada a una VM de Compute Engine, por lo que usar más TPU (por ejemplo, en un Pod de TPU) aumenta la probabilidad de que una de tus VMs experimente un evento de mantenimiento.

Nota: Las diferencias entre las VM de TPU y los nodos TPU se describen en Arquitectura del sistema.

En este documento se analizan varios enfoques para controlar eventos de mantenimiento de trabajos de entrenamiento de larga duración en Cloud TPU.

Usa puntos de control para una recuperación rápida de los eventos de mantenimiento

Los puntos de control son clave para las recuperaciones cortas de los eventos de mantenimiento y deben guardarse con frecuencia: una buena regla general es guardar puntos de control aproximadamente a cada hora. Si no guardas puntos de control a menudo corres el riesgo de perder mucho progreso del entrenamiento debido a eventos de mantenimiento o a otras interrupciones del entrenamiento.

Importante: Cuando un nodo TPU falla o se apaga por un evento de mantenimiento, cualquier es posible que se pierdan los datos almacenados en el disco de arranque de la VM. Por lo tanto, te recomendamos guardar puntos de control en algún tipo de almacenamiento externo, como un disco persistente o un bucket de Cloud Storage. Para obtener más información, consulta Opciones de almacenamiento de TPU

En general, los puntos de control se refieren a todos los parámetros guardados que se usan en el entrenamiento (como los pesos de los modelos). El tiempo que lleva guardar un punto de control puede variar desde segundos hasta minutos.

Aunque la mayoría de los eventos de mantenimiento se recuperan de forma automática y los trabajos de entrenamiento continúan sin intervención manual, puede haber casos extremos en los que el trabajo no se reinicie y continúe de forma automática. Cuando esto sucede, debes borrar y volver a crear los recursos de TPU y reiniciar el trabajo de entrenamiento desde un punto de control guardado. Para obtener información sobre cómo detectar y recuperar de errores de recuperación, consulta Detecta y recupérate de las fallas de TPU.

Los mecanismos que se usan para guardar y cargar los puntos de control son diferentes para cada AA en un framework de aplicaciones. Por lo general, los modelos de Cloud TPU compatibles tienen puntos de control integrados. Para obtener más información sobre los puntos de control, consulta TensorFlow 2.x, PyTorch o JAX/flax.

Detecta eventos de mantenimiento

Puedes detectar si se produce un evento de mantenimiento en tu TPU y cuándo lo haces con el siguiente comando gcloud describe:

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

El resultado de este comando muestra el estado actual de la TPU y una descripción del evento de mantenimiento más reciente. El resultado debería verse similar a lo siguiente:

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

Registros de eventos de mantenimiento

Puedes ver los registros históricos de eventos de mantenimiento en tu TPU en los registros de auditoría de eventos del sistema.

En el menú de navegación de la consola de Google Cloud, haz clic en Compute Engine > Instancias de VM y búsquedas, por ejemplo:

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

Dentro del período de búsqueda, las interrupciones y reparaciones de tus trabajadores de TPU que se muestran. Los registros mostrarán la fecha y hora del evento, el tipo de del evento y para "terminate" eventos, el motivo del cierre en protoPayload.metadata.terminateReason

Controla eventos de mantenimiento

Existen varias formas de mitigar las interrupciones del evento de mantenimiento.

1. Guarda puntos de control de forma periódica

   En la situación ideal, cuando un "evento de interrupción" sucede, el entrenamiento se reanuda desde el último punto de control.

2. Reintentos de la secuencia de comandos de entrenamiento

   La secuencia de comandos de entrenamiento puede detenerse como resultado de un "evento de interrupción". Puedes usar una secuencia de comandos bash para reintentar la secuencia de comandos de entrenamiento de forma continua hasta que se complete. Cada reintento debe continuar desde el último punto de control, por lo que las secuencias de comandos de reintento siempre deben usarse junto con los puntos de control.

   Las canalizaciones de entrenamiento listas para la producción deben usar un sistema de administración de recursos como Google Kubernetes Engine (GKE). Para obtener más información sobre el uso Google Kubernetes Engine con VM de TPU, consulta Implementa cargas de trabajo de TPU. Puedes implementar una secuencia de comandos bash para reintentar continuamente el de la secuencia de comandos de entrenamiento hasta que se complete. Por ejemplo:

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

   Nota: Debes ejecutar la secuencia de comandos desde Cloud Shell o desde una terminal, no desde la VM de TPU).

3. Detecta y recupera fallas de la TPU

   Cuando una TPU no se recupera de un evento de mantenimiento, puedes usar una para detectar el estado de la TPU y borrarla y volver a crearla. Ejemplo Consulte este artículo de la secuencia de comandos. Consulta Administra las TPU para obtener más información sobre borrar y volver a crear manualmente las TPU.

   Cuando crees o vuelvas a crear una VM de TPU, puedes especificar una secuencia de comandos de inicio con el parámetro --metadata startup-script Se ejecuta una secuencia de comandos de inicio se crea la VM de TPU. Consulta Ejecuta secuencias de comandos de instalación estándar para obtener más información.