メンテナンス イベントに備える

TPU VM は、TPU ハードウェアが割り当てられた Compute Engine VM のインスタンスです。Compute Engine VM は Compute Engine VM のメンテナンス イベントの対象となります。各 TPU は Compute Engine VM に割り当てられているため、使用する TPU（TPU スライスなど）を増やすと、いずれかの VM でメンテナンス イベントが発生する可能性があります。

このドキュメントでは、Cloud TPU で長時間実行トレーニング ジョブのメンテナンス イベントを処理する方法について説明します。Google Kubernetes Engine（GKE）で TPU のメンテナンス イベントを処理する方法については、GPU と TPU の GKE ノードの中断を管理するをご覧ください。

今後のメンテナンスに関する通知を表示する

今後のホスト メンテナンス イベントの通知を表示できます。インスタンスの今後のメンテナンスの時間枠をモニタリングすると、中断を最小限に抑えながら、今後のメンテナンスを処理するワークロードをプロアクティブに準備できます。詳細については、メンテナンス通知を表示するをご覧ください。

チェックポイントを使用してメンテナンス イベントから迅速に復元する

チェックポイントは、メンテナンス イベントから短時間で復元するうえで重要であるため、頻繁に保存する必要があります。チェックポイントは約 1 時間ごとに保存することをおすすめします。チェックポイントが不十分であると、メンテナンス イベントやその他のトレーニングの中断が原因となり、多くのトレーニングの進行状況が失われるリスクが発生します。

一般に、チェックポイントはトレーニングで使用されるすべての保存済みパラメータ（モデルの重みなど）を参照します。チェックポイントの保存に要する時間は、数秒から数分程度になることがあります。

TPU はほとんどのメンテナンス イベントから自動的に復元でき、トレーニング ジョブは手動で操作しなくても継続されます。ただし、ジョブが再起動せず自動的に続行するケースがまれにあります。この場合、TPU リソースを削除して再作成し、保存したチェックポイントからトレーニング ジョブを再起動する必要があります。自動復元の失敗の検出と復旧方法については、TPU の障害を検出して回復するをご覧ください。

チェックポイントの保存と読み込みのメカニズムは、ML フレームワークごとに異なります。通常、サポートされている Cloud TPU モデルにはチェックポインティングが組み込まれています。チェックポインティングの詳細については、TensorFlow 2.x、PyTorch、または JAX/flax をご覧ください。

Autocheckpoint を使用する

Autocheckpoint 機能を使用すると、メンテナンス イベントが発生したときにスケジュールが設定されていないチェックポイントを保存するようにコードを構成することで、トレーニングの進行状況を保持できます。Autocheckpoint の詳細については、Cloud TPU Autocheckpoint をご覧ください。

トレーニング スクリプトを再試行する

トレーニング スクリプトは、「中断イベント」の結果として停止する場合があります。bash スクリプトを使用すると、トレーニングが完了するまでトレーニング スクリプトを継続的に再試行できます。例:

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

再試行は最新のチェックポイントから続行するため、再試行スクリプトは常にチェックポイントと組み合わせて使用する必要があります。

本番環境に対応したトレーニング パイプラインでは、Google Kubernetes Engine（GKE）などのリソース管理システムを使用する必要があります。TPU VM で Google Kubernetes Engine を使用する方法については、TPU ワークロードをデプロイするをご覧ください。

TPU の障害を検出して回復する

TPU がメンテナンス イベントから回復しない場合、復元スクリプトを使用して TPU の状態を検出し、TPU を削除して再作成できます。復元スクリプトの例については、retry.sh をご覧ください。トレーニング スクリプトを実行しているプロセスがクラッシュした場合は、復元スクリプトを変更してトレーニング スクリプトの実行を再試行できます。

TPU を手動で削除して再作成する方法については、TPU リソースの管理をご覧ください。

コレクション スケジューリングを使用する

Cloud TPU ではコレクション スケジューリングのコンセプトが採用され、トレーニングまたはサービングと推論のワークロードのいずれかをサポートするために使用できる 2 種類のコレクションが用意されています。この機能を使用して Cloud TPU インスタンスをデプロイすると、 Google Cloud はアプリケーションに最適な一意のメンテナンス スケジュールを適用します。各コレクション タイプでは、次の動作が想定されます。

• トレーニング（デフォルト）: このコレクション タイプは、メンテナンス イベント中にサービスを迅速に復元するために、すべてのインスタンスでダウンタイムを最小限に抑え、予期しない中断を制限する必要のある一般的なトレーニング ワークロードに適しています。トレーニング コレクション タイプは、インスタンスのグループのメンテナンス イベントの並列スケジューリングと実行を提供します。

• サービング（--workload-type=AVAILABILITY_OPTIMIZED を使用して利用可能）: このコレクション タイプは、メンテナンス イベント中もサービスの継続性を確保するために、インスタンス（レプリカ）のサブセット全体でダウンタイムを最小限に抑える必要のあるサービング ワークロードまたは推論ワークロードに適しています。サービング コレクション タイプでは、インスタンスのグループに対してメンテナンス イベントのスケジュールと実行が分散されます。サービング コレクションの指定は、TPU v6e でのみサポートされています。

コレクション スケジューリングの詳細については、推論ワークロードの TPU コレクション スケジューリングをご覧ください。

メンテナンス イベントを検出する

TPU でメンテナンス イベントが発生するかどうかと、そのタイミングは、次の gcloud compute tpus tpu-vm describe コマンドを使用して検出できます。

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

このコマンドの出力には、TPU の現在の状態と最新のメンテナンス イベントの説明が表示されます。出力は次のようになります。

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

メンテナンス イベントログを表示する

TPU のメンテナンス イベントの過去のログは、システム イベント監査ログで確認できます。

1. Google Cloud コンソールのナビゲーション メニューで、[ログ エクスプローラ] ページに移動します。

   [ログ エクスプローラ] に移動

2. 次の検索クエリを使用して、終了または再起動された TPU を表示します。

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

   検索期間内の TPU ワーカーの中断と修復のログが結果に表示されます。ログには次のものが含まれます。

   • イベントの発生日時
   • イベントのタイプ
   • 「終了」イベントの場合は、protoPayload.metadata.terminateReason フィールドに終了理由

次のステップ

• メンテナンス通知を表示する
• ホスト メンテナンス イベントを手動で開始する