Gerenciar eventos de manutenção com pods do Cloud TPU

Visão geral

As VMs de TPU são instâncias de VMs do Compute Engine com hardware de TPU anexado. As VMs do Compute Engine estão sujeitas a eventos de manutenção delas. Cada TPU é conectada a uma VM do Compute Engine. Portanto, o uso de mais TPUs (por exemplo, em um pod de TPU) aumenta a probabilidade de uma das VMs encontrar um evento de manutenção.

Observação: as diferenças entre VMs e nós de TPU estão descritas em Arquitetura do sistema.

Neste documento, discutimos várias abordagens para processar eventos de manutenção de jobs de treinamento de longa duração nos Cloud TPUs.

Como usar checkpoints para recuperação rápida de eventos de manutenção

Checkpoints são fundamentais para recuperações curtas de eventos de manutenção e precisam ser salvos com frequência: uma boa regra é salvar checkpoints aproximadamente a cada hora. Não criar checkpoints com frequência suficiente corre o risco de perder muito progresso do treinamento devido a eventos de manutenção ou outras interrupções no treinamento.

Importante: quando um nó de TPU falha ou é encerrado para um evento de manutenção, os dados armazenados no disco de inicialização da VM poderão ser perdidos. Portanto, recomendamos que você salvar pontos de verificação em alguma forma de armazenamento externo, como um disco permanente ou um bucket do Cloud Storage. Para mais informações, consulte Opções de armazenamento da TPU

Os checkpoints geralmente se referem a todos os parâmetros salvos usados no treinamento, como pesos do modelo. O tempo necessário para salvar um checkpoint pode variar da ordem de segundos à ordem de minutos.

Embora a maioria dos eventos de manutenção seja recuperada automaticamente e os jobs de treinamento continuem sem intervenção manual, pode haver casos extremos em que o job não é reiniciado e continua automaticamente. Quando isso acontece, é necessário excluir e recriar os recursos da TPU e reiniciar o job de treinamento a partir de um checkpoint salvo. Para mais informações sobre como detectar e recuperar de falhas de recuperação automática, consulte Detectar e recuperar de falhas de TPU.

Os mecanismos usados para salvar e carregar checkpoints são diferentes para cada ML de análise de dados em nuvem. Os modelos do Cloud TPU com suporte geralmente têm checkpoint integrado. Para mais informações sobre checkpoints, consulte: TensorFlow 2.x, PyTorch ou JAX/flax.

Como detectar eventos de manutenção

É possível detectar quando e se um evento de manutenção ocorre na sua TPU usando o seguinte comando gcloud describe:

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

A saída desse comando mostra o estado atual da TPU e uma descrição do evento de manutenção mais recente. O resultado será semelhante a:

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

Logs de eventos de manutenção

É possível conferir os registros históricos de eventos de manutenção no TPU nos registros de auditoria de eventos do sistema.

No menu de navegação do console do Google Cloud, clique em Compute Engine > Instâncias de VM e pesquise, por exemplo:

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

Durante o período de pesquisa, quaisquer interrupções e reparos dos seus workers da TPU são exibidos. Os registros mostram a data e a hora do evento, o tipo evento e para "terminate" o motivo da rescisão protoPayload.metadata.terminateReason:

Gerenciar eventos de manutenção

Há várias maneiras de reduzir interrupções em eventos de manutenção.

1. Salvar os checkpoints periodicamente

   No cenário ideal, quando um "evento de interrupção" acontece, currículos de treinamento a partir do último checkpoint.

2. Novas tentativas de script de treinamento

   O script de treinamento pode parar como resultado de um "evento de interrupção". É possível usar um script bash para repetir continuamente o script de treinamento até que ele seja concluído. Cada nova tentativa precisa continuar a partir do checkpoint mais recente. Portanto, os scripts de nova tentativa precisam ser sempre usados em conjunto com checkpoints.

   Os pipelines de treinamento prontos para produção precisam usar um sistema de gerenciamento de recursos, como o Google Kubernetes Engine (GKE). Para mais informações sobre como usar Google Kubernetes Engine com VMs TPU, consulte Implantar cargas de trabalho de TPU. Implemente um script bash para repetir continuamente script de treinamento até a conclusão. Exemplo:

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

   Observação: execute o script em um Cloud Shell ou em um terminal, não na VM da TPU.

3. Detectar e recuperar de falhas da TPU

   Quando uma TPU não se recupera de um evento de manutenção, é possível usar um script de recuperação para detectar o estado da TPU, além de excluir e recriar a TPU. Confira um exemplo deste script aqui. Consulte Como gerenciar TPUs para ver detalhes sobre excluir e recriar manualmente as TPUs.

   Ao criar ou recriar uma VM de TPU, é possível especificar um script de inicialização com o parâmetro --metadata startup-script. Um script de inicialização é executado sempre que uma VM de TPU é criada. Consulte Executar scripts de instalação padrão para mais informações.