GPU と TPU の GKE ノードの中断を管理する

このページでは、人工知能（AI）または ML のワークロードを実行している Google Kubernetes Engine（GKE）ノードで発生する可能性のある中断イベントについて説明します。このページの内容は次のとおりです。

長時間実行される GKE クラスタのライフサイクル中に、GKE インフラストラクチャの基盤となる Compute Engine リソースに対して Google Cloud から発行された自動メンテナンス イベントにより、ワークロードが定期的に中断されます。このような中断が AI / ML ワークロードを実行している GKE ノードに影響する場合、GKE は実行中のワークロードとその基盤となるノードの両方を再起動する必要があります。

GPU と TPU で中断管理が必要となる理由

GKE クラスタは、GKE ノードのライフサイクルを管理します。これらのノードは Compute Engine VM にプロビジョニングされます。Compute Engine VM では、ハードウェアやソフトウェアの更新、メンテナンス、ハードウェア障害など、さまざまな理由でホストイベントが定期的に発生します。ホストイベントは、基盤となる Google Cloud インフラストラクチャに対して発行され、GKE のメンテナンス ポリシーと除外をバイパスします。

一部の例外を除き、ほとんどの Compute Engine VM のホスト メンテナンス ポリシーはライブ マイグレーションに設定されています。つまり、実行中のワークロードが中断されることはありません。ただし、特定のクラスの VM はライブ マイグレーションに対応していません。たとえば、AI / ML ワークロードが実行される GPU や TPU が割り当てられた VM はサポートしていません。また、GKE は、プリエンプションを使用して予約済み TPU とオンデマンド TPU を再起動することもできます。これにより、デフラグのために、より大きな TPU のプロビジョニングが可能になります。

ホストイベントが発生すると、GKE はノードとその Pod を終了します。Pod が Job や Deployment などのより大きなワークロードの一部としてデプロイされている場合、GKE は、影響を受けるノードの Pod を再起動します。こうした中断に対しては、ユーザー自身、あるいはワークロードとジョブを処理するフレームワークによって適切に対応する必要があります。たとえば、AI トレーニング ジョブの状態を保存することで、データの損失を減らします。

正常終了のプロセス

次のワークフローは、Compute Engine によって中断が発行された後に、GKE がノードの停止を実行する方法を示しています。

  1. Compute Engine が、VM メタデータキー maintenance-event の値を TERMINATE_ON_HOST_MAINTENANCE に更新します。

  2. 60 秒以内に、次の処理が行われます。

    1. システム コンポーネントが、メンテナンスの進行中であることを示すため、true に設定された新しいノードラベル（cloud.google.com/active-node-maintenance）を適用します。

    2. GKE は、ノード taint（cloud.google.com/impending-node-termination:NoSchedule）を適用して、新しい Pod がノードでスケジュールされないようにします。突然の終了ではないため、この taint を許容するようにワークロードを変更することをおすすめします。

  3. maintenance-handler コンポーネントが、ワークロード Pod、システム Pod（kube-system など）の順に Pod の強制排除を開始します。

  4. GKE がシャットダウン シグナル SIGTERM を送信し、ノードで実行されているワークロード Pod にまもなくシャットダウンすることを通知します。このアラートにより、Pod は進行中のタスクを完了できます。GKE は、ベスト エフォートでこれらの Pod を正常に停止します。

maintenance-event 通知は、GKE ノードの基盤となる Compute Engine VM で、ノードの終了につながるホストイベントが発生したときに発生します。この場合、Compute Engine は maintenance-event メタデータキーを更新します。ノードが終了する前のメンテナンス通知期間は次のとおりです。

  • GPU: 60 分前
  • TPU: 5 分前

その後、ノードが終了し、交換ノードが割り当てられます。プロセスが完了すると、GKE はラベルと taint をクリアします。GPU または TPU を使用するワークロードの終了ウィンドウを長くするには、ワークロードを正常に終了するように GKE を構成するの説明に従って操作します。

GKE でワークロードの中断を処理する

GKE では、GKE の終了イベントを管理し、クラスタ内のワークロードの中断を減らすため、これらの通知をモニタリングして次の処理を行います。

  • シャットダウンが迫っていることをワークロードに事前に通知する: ホストのメンテナンス イベントのために GKE ノードを停止する必要がある場合、GKE は事前通知期間の開始時にノードで実行中の Pod に SIGTERM シグナルを送信します。SIGTERM などの OS シグナルは、Python や Go などのほとんどの標準ライブラリでネイティブに処理できます。SIGTERM をキャプチャできるフレームワークには、MaxText、Pax、Orbax を使用した JAX などがあります。
  • ワークロードを正常に終了する: Pod の終了猶予期間を使用してワークロードを正常に終了するように GKE を構成できます。Pod は SIGTERM シグナルに応答して、進行中のタスクを終了し、トレーニング状態の保存など、定義した終了アクションを実行します。正常終了中に、GKE は Pod を正常に終了し、クリーンアップ プロセスを実行します。また、アプリケーションで定義されている終了アクションを実行します。たとえば、データ損失を抑えるためにワークロードのデータを保存したり、トレーニング状態を保存します。

ワークロードを正常に終了するように GKE を構成する

このセクションでは、アプリケーションのライフサイクルを管理し、ワークロードの中断を最小限に抑えるように GKE を構成します。猶予期間を構成しない場合、猶予期間はデフォルトで 30 秒になります。

GKE は、これらの Pod を正常に終了し、トレーニング状態の保存など、定義した終了アクションを実行するために最善を尽くします。GKE は、猶予期間の開始時に Pod に SIGTERM シグナルを送信します。猶予期間の終了までに Pod が終了しなかった場合、GKE は Pod 内のコンテナでまだ実行中のプロセスにフォローアップの SIGKILL シグナルを送信します。

ワークロードの正常終了時間を構成するには、GPU または TPU の手順に沿って操作します。

GPU

Pod マニフェストで、spec.terminationGracePeriodSeconds フィールドに最大 3,600 秒（60 分）までの値を設定します。たとえば、通知時間を 10 分にするには、次のように Pod マニフェストの spec.terminationGracePeriodSeconds フィールドに 600 秒を設定します。

    spec:
      terminationGracePeriodSeconds: 600

進行中のタスクが通知期間内に完了できるように、終了猶予期間には十分な時間を設定することをおすすめします。

TPU

クリーンアップ プロセスの実行に最大時間を割り当てるには、Pod マニフェストの spec.terminationGracePeriodSeconds フィールドを 300 秒（5 分）に設定します。例:

    spec:
      terminationGracePeriodSeconds: 300

進行中のタスクが通知期間内に完了できるように、終了猶予期間には十分な時間を設定することをおすすめします。

ワークロードで Orbax を使用して MaxText、Pax、JAX などの ML フレームワークを利用している場合、ワークロードは SIGTERM シグナルをキャプチャして、チェックポイント プロセスを開始できます。詳細については、TPU の Autocheckpoint をご覧ください。

アクティブな正常終了の進行状況をモニタリングする

コントロール プレーンで 1.29.1-gke.1425000 以降が実行されている GKE クラスタの場合、GKE は gpu-maintenance-handler というノードレベルのコンポーネントをデプロイします。このコンポーネントは、対応するコントロール プレーン コンポーネントと一緒にすべての GPU ノードと TPU ノードで実行されます。これらのコンポーネントは次の処理を行います。

  • 正常終了イベントを処理します。
  • SIGTERM シグナルをノードで実行中のワークロードに転送し、GKE VM で差し迫った中断イベントに対応します。これらの中断は、Pod の強制排除リクエストと削除リクエストとしてログに記録されます。

GKE は、シャットダウン直前のノードにラベルと taint を追加します。GKE は、各 GPU ノードと TPU ノードで実行されているシステム コンポーネント maintenance-handler を使用して、メンテナンスなどのホストイベント通知をモニタリングします。

GKE は、次の正常終了イベントをログに記録します。

  • GCE ホストのメンテナンス イベントなど、ノードの終了を伴う差し迫った中断を検出した場合: GKE は、次のノードラベルを追加します。cloud.google.com/active-node-maintenance は true に設定されます。
  • 新しいワークロードのスケジューリングを制限する場合: GKE は cloud.google.com/impending-node-termination:NoSchedule taint を適用します。

GKE が正常終了を完了すると、ラベルと taint がクリアされます。

中断によって発生したアクティブな正常終了のステータスをモニタリングするには、コンソールまたは Google Cloud CLI を使用して gpu-maintenance-handler ログを表示します。

gcloud

  1. 次のコマンドを実行して、gpu-maintenance-handler のインスタンスを実行しているノードと Pod の名前を探します。

    kubectl get pods -l name=maintenance-handler -A -o wide

    出力の各行には、ノード名、Pod 名、ステータスが含まれます。

  2. ログを調べる

    kubectl logs -n=kube-system MAINTENANCE_HANDLER_POD_NAME

    MAINTENANCE_HANDLER_POD_NAME は、ハンドラ インスタンスの名前に置き換えます。

    メンテナンス イベントが検出されると、Pod はメッセージを記録し、ラベルを適用して、強制排除を開始します。

  3. ノードのラベルと taint を確認します。

    kubectl describe node NODE_NAME

    NODE_NAME は、表示するノードの名前に置き換えます。

    出力には、監視するノードラベルとタグのリストが表示されます。

コンソール

  1. Google Cloud コンソールの [ログ エクスプローラ] ページに移動します。

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

  2. [クエリ] フィールドに次のクエリを規定します。

    resource.type="k8s_container"
resource.labels.namespace_name="kube-system"
resource.labels.container_name="maintenance-handler"
resource.labels.cluster_name="CLUSTER_NAME"

    CLUSTER_NAME は、クラスタの名前に置き換えます。

次のステップ