Administra eventos de mantenimiento con pods de Cloud TPU

Descripción general

Los nodos TPU y las VM de TPU son instancias de VM de Compute Engine con hardware de TPU adjunto. Las VM de Compute Engine están sujetas a los eventos de mantenimiento de VM de Compute Engine. Cada TPU está conectada a una VM de Compute Engine, por lo que el uso de más TPU (por ejemplo, en un pod de TPU) aumenta la probabilidad de que una de tus VM encuentre un evento de mantenimiento.

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.

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 las fallas de recuperación automáticas y recuperarse de ellas, consulta Detecta fallas de TPU y recupérate de ellas.

Los mecanismos que se usan para guardar y cargar puntos de control son diferentes en cada framework de AA. Los modelos de Cloud TPU compatibles suelen tener un punto de control integrado. 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 ocurre un evento de mantenimiento en la TPU y cuándo ocurre con el siguiente comando gcloud describe:

VMs de TPU

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

Nodos TPU

$ gcloud compute tpus 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 ser similar al 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 los eventos de mantenimiento en la 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 busca, por ejemplo:

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

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

Controla eventos de mantenimiento

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

Guarda puntos de control de forma periódica

En la situación ideal, cuando ocurre un “evento de interrupción”, el entrenamiento se reanuda desde el último punto de control.
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 de Google Kubernetes Engine con la arquitectura de VM de TPU, consulta Implementa cargas de trabajo de TPU. Para obtener más información sobre el uso de Google Kubernetes Engine con la arquitectura de nodo TPU, consulta Ejecuta aplicaciones de TPU en Google Kubernetes Engine. De lo contrario, puedes implementar una secuencia de comandos bash para reintentar de forma continua la secuencia de comandos de entrenamiento hasta que se complete. Por ejemplo:

Con un nodo TPU:

Desde tu VM bash while ! python3 [training command]; do sleep 1; done

Con VM de TPU:
```
while ! gcloud compute tpus tpu-vm ssh ${TPU_NAME} --command "python3 [training command]"; do sleep 1; done
```
(ten en cuenta que debes ejecutar el comando de VM de TPU desde Cloud Shell o desde una terminal, no desde la VM de TPU).
Detecta y recupera fallas de la TPU

Cuando una TPU no se recupera de un evento de mantenimiento, puedes usar una secuencia de comandos de recuperación para detectar el estado de la TPU y borrarla y volver a crearla. Puedes encontrar un ejemplo de esta secuencia de comandos aquí. Consulta Administra las TPU para obtener detalles sobre cómo borrarlas y volver a crearlas de forma manual.

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