Gestire gli eventi di manutenzione con i pod Cloud TPU

Panoramica

Le VM TPU sono istanze di VM Compute Engine con hardware TPU collegato. Le VM di Compute Engine sono soggette a eventi di manutenzione delle VM di Compute Engine. Ogni TPU è collegata a una VM Compute Engine, pertanto l'utilizzo di più TPU (ad esempio in un pod di TPU) aumenta la probabilità che una delle tue VM riscontri un evento di manutenzione.

Questo documento illustra vari approcci per gestire gli eventi di manutenzione per i job di addestramento di lunga durata sulle Cloud TPU.

Utilizzo dei checkpoint per un rapido recupero dagli eventi di manutenzione

I checkpoint sono fondamentali per i rapidi ripristini dagli eventi di manutenzione e devono essere salvati di frequente: una buona regola empirica è salvarli ogni ora circa. Se non esegui il checkpoint abbastanza spesso, rischi di perdere molto del progresso della formazione a causa di eventi di manutenzione o altre interruzioni della formazione.

.

In genere, i controlli regolari si riferiscono a tutti i parametri salvati utilizzati nell'addestramento (ad esempio i pesi del modello). Il tempo necessario per salvare un checkpoint può variare da pochi secondi a pochi minuti.

Sebbene la maggior parte degli eventi di manutenzione venga recuperata automaticamente e i job di addestramento continuino senza intervento manuale, potrebbero verificarsi casi limite in cui il job non si riavvia e non continua automaticamente. In questo caso, devi eliminare e ricreare le risorse TPU, e riavviare il job di addestramento da un checkpoint salvato. Per informazioni su come rilevare e recuperare da errori di recupero automatico, consulta Rilevare e recuperare da errori TPU.

I meccanismi utilizzati per salvare e caricare i checkpoint sono diversi per ogni framework di ML. I modelli Cloud TPU supportati hanno generalmente il checkpointing integrato. Per ulteriori informazioni sul checkpointing, consulta TensorFlow 2.x, PyTorch o JAX/flax.

Rilevamento di eventi di manutenzione

Puoi rilevare se e quando si verifica un evento di manutenzione sulla tua TPU utilizzando il seguente comando gcloud describe:

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

L'output di questo comando mostra lo stato attuale della TPU e una descrizione dell'evento di manutenzione più recente. L'output dovrebbe essere simile al seguente:

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

Log degli eventi di manutenzione

Puoi visualizzare i log storici degli eventi di manutenzione sulla TPU negli audit log degli eventi di sistema.

Nel menu di navigazione della console Google Cloud, fai clic su Compute Engine > Istanze VM e cerca, ad esempio:

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

Nel periodo di tempo selezionato per la ricerca, vengono visualizzate eventuali interruzioni e riparazioni dei tuoi worker TPU. I log mostreranno la data e l'ora dell'evento, il tipo di evento e, per gli eventi "terminate", il motivo dell'interruzione in protoPayload.metadata.terminateReason.

Gestire gli eventi di manutenzione

Esistono diversi modi per ridurre le interruzioni dovute agli eventi di manutenzione.

  1. Salvare periodicamente i checkpoint

    Nello scenario ideale, quando si verifica un "evento di interruzione", la formazione riprende dall'ultimo checkpoint.

  2. Nuovi tentativi dello script di addestramento

    Lo script di addestramento potrebbe arrestarsi a causa di un "evento di interruzione". Puoi utilizzare uno script bash per riprovare continuamente allo script di addestramento fino al completamento dell'addestramento. Ogni nuovo tentativo deve continuare dall'ultimo checkpoint, pertanto gli script di nuovo tentativo devono essere sempre utilizzati in combinazione con i checkpoint.

    Le pipeline di addestramento pronte per la produzione devono utilizzare un sistema di gestione delle risorse come Google Kubernetes Engine (GKE). Per ulteriori informazioni sull'utilizzo di Google Kubernetes Engine con le VM TPU, consulta Eseguire il deployment dei workload TPU. Puoi implementare uno script bash per ripetere continuamente lo script di addestramento fino al completamento. Ad esempio:

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

  3. Rileva e recupera da guasti delle TPU

    Quando una TPU non si riprende da un evento di manutenzione, puoi utilizzare uno script di recupero per rilevare lo stato della TPU ed eliminarla e ricrearla. Un esempio di questo script è disponibile qui. Per maggiori dettagli sull'eliminazione e sulla ricostituzione manuale delle TPU, consulta la sezione Gestire le TPU.

    Quando crei o crei di nuovo una VM TPU, puoi specificare uno script di avvio con il parametro --metadata startup-script. Uno script di avvio viene eseguito ogni volta che viene creata una VM TPU. Per ulteriori informazioni, consulta Eseguire script di installazione standard.