Gestire gli eventi di manutenzione con i pod Cloud TPU

Panoramica

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

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

Utilizzo dei checkpoint per un recupero rapido da eventi di manutenzione

I checkpoint sono fondamentali per brevi recuperi da eventi di manutenzione e devono essere salvati di frequente: una buona regola generale è salvare i checkpoint approssimativamente ogni ora. Non controllare abbastanza spesso rischi di perdere molti progressi durante l'addestramento a causa di eventi di manutenzione o di altre interruzioni.

In genere, i checkpoint si riferiscono a tutti i parametri salvati utilizzati durante l'addestramento (ad esempio le ponderazioni del modello). Il tempo necessario per salvare un checkpoint può variare da secondi a 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 viene riavviato e 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 ripristinare gli errori di ripristino automatico, vedi Rilevamento e ripristino dagli errori TPU di seguito.

I meccanismi utilizzati per salvare e caricare i checkpoint sono diversi per ogni framework ML. In genere i modelli Cloud TPU supportati hanno il checkpoint integrato. Per ulteriori informazioni sul checkpoint, 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:

VM TPU

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

Nodi TPU

$ gcloud compute tpus 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

Gestire gli eventi di manutenzione

Esistono diversi modi per mitigare le interruzioni di eventi di manutenzione.

  1. Salva periodicamente i checkpoint

    Nello scenario ideale, quando si verifica un "evento di interruzione", l'addestramento riprende semplicemente dall'ultimo checkpoint.

  2. Nuovi tentativi di script di addestramento

    Lo script di addestramento potrebbe interrompersi a causa di un "evento di interruzione". Puoi utilizzare uno script bash per riprovare continuamente. Ogni nuovo tentativo dovrebbe continuare dall'ultimo checkpoint, quindi gli script dei nuovi tentativi devono essere sempre usati insieme ai checkpoint.

    Le pipeline di addestramento pronte per la produzione devono utilizzare un sistema di gestione delle risorse come Google Kubernetes Engine (GKE). GKE è attualmente supportato con i nodi TPU, ma non con le VM TPU. In caso contrario, puoi implementare uno script bash per riprovare a seguire lo script di addestramento fino al completamento. Ad esempio:

    Con nodo TPU:

    (Dalla tua VM) while ! python3 [training command]; do sleep 1; done

    Con VM TPU:

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

    Tieni presente che devi eseguire il comando VM TPU da una shell Google Cloud o da un terminale, non dalla VM TPU.

  3. Rilevamento e ripristino dagli errori TPU

    Se una TPU non si recupera da un evento di manutenzione, puoi utilizzare uno script di recupero per rilevare lo stato della TPU, nonché eliminare e ricreare la TPU. Un esempio di questo script è disponibile qui. Vedi Gestione delle TPU per maggiori dettagli sull'eliminazione e sulla creazione manuale di TPU.

    Durante la creazione o la ricreazione di 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 la sezione Eseguire script di installazione standard.