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

Informações gerais

Os nós e 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 de VMs do Compute Engine. Cada TPU está conectada a uma VM do Compute Engine. Portanto, usar mais TPUs (por exemplo, em um pod de TPU) aumenta a probabilidade de uma das suas VMs encontrar um evento de manutenção.

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.

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 informações sobre como detectar e recuperar falhas de recuperação automática, consulte Detectar e recuperar falhas de TPU.

Os mecanismos usados para salvar e carregar checkpoints são diferentes para cada framework de ML. Os modelos de Cloud TPU compatíveis geralmente têm checkpoint integrado. Para mais informações sobre checkpoint, consulte TensorFlow 2.x, PyTorch ou JAX/flax.

Detectar eventos de manutenção

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

VMs de TPU

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

Nós da TPU

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

A saída desse comando exibe o estado atual da TPU e uma descrição do evento de manutenção mais recente. A saída vai ser semelhante a esta:

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

Registros de eventos de manutenção

É possível visualizar os registros históricos de eventos de manutenção na sua 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"

No período de pesquisa, são exibidas todas as interrupções e reparos dos workers da TPU. Os registros mostrarão a data e a hora do evento, o tipo de evento e, para eventos "encerrado", o motivo do encerramento em 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 ocorre um "evento de interrupção", o treinamento é retomado do checkpoint mais recente.

  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 Production precisam usar um sistema de gerenciamento de recursos, como o Google Kubernetes Engine (GKE). Para mais informações sobre como usar o Google Kubernetes Engine com a arquitetura de VM de TPU, consulte Implantar cargas de trabalho de TPU. Para mais informações sobre como usar o Google Kubernetes Engine com a arquitetura de nós de TPU, consulte Executar aplicativos de TPU no Google Kubernetes Engine. Caso contrário, implemente um script bash para repetir continuamente o script de treinamento até a conclusão. Exemplo:

    Com o nó da TPU:

    (a partir da VM) bash while ! python3 [training command]; do sleep 1; done

    Com a VM da TPU:

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

    Observe que é necessário executar o comando da VM da TPU 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 desse script aqui. Consulte Como gerenciar TPUs para ver detalhes sobre como excluir e recriar manualmente as TPUs.

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