GKE に All Capacity モードで TPU をデプロイする

Standard

プレビュー版 - TPU All Capacity mode

この機能は、 サービス固有の規約の対象となります。 pre-GA の機能は「現状のまま」で提供されるものであり、サポートが制限されることがあります。 詳細については、リリース ステージの説明をご覧ください。

このドキュメントでは、GKE の TPU 全容量モード機能を使用してワークロードをデプロイして管理する方法について説明します。All Capacity モードの予約では、TPU リソースの制御が強化され、予約済みの容量内でワークロードをきめ細かい制御で配置できます。

このドキュメントは、TPU デプロイをきめ細かく制御して Kubernetes コンテナ オーケストレーションを使用する ML エンジニア、プラットフォーム管理者、オペレーターを対象としています。

このドキュメントを読む前に、次の内容をよく理解しておいてください。

• TPU
• TPU システム アーキテクチャ
• GKE の TPU

TPU All Capacity モードとは

TPU Cluster Director によって有効になる TPU All Capacity モードでは、予約済みの TPU 容量を完全に制御できます。TPU Cluster Director は、TPU を予約ベースで制御できる管理サービスです。

以前のマネージド モードでは、ハードウェア障害に対処するために容量の一部が予約されていましたが、All Capacity モードでは、予約済みの TPU リソース全体にアクセスできます。 Google Cloud このモードでは、ハードウェアのステータスを完全に把握できますが、ノード障害と計画メンテナンスの管理はお客様の責任となります。

All Capacity モードの主な機能の詳細については、TPU Cluster Director の概要の All Capacity モードをご覧ください。

GKE の All Capacity モードに関連する用語

次の表に、Ironwood（TPU7x）バージョンでのブロック、サブブロック、キューブのサイズで定義される用語と等価性を示します。キューブは、相互接続された TPU チップの 4x4x4 トポロジです。これは、3 タプル（{A}x{B}x{C}）のトポロジにのみ適用されます。

TPU リソース	コア数	チップ	ホスト	キューブ
1 チップ	2	1	-	-
1 ホスト	8	4	1	-
1 個のサブブロック	128	64	16	1
1 つのブロックには 144 個のサブブロックが含まれています	18432	9,216	2304	144

ブロックで許可されるトポロジの詳細については、トポロジを選択するをご覧ください。

始める前に

作業を始める前に、次のタスクが完了していることを確認してください。

• Google Kubernetes Engine API を有効にする。
Google Kubernetes Engine API を有効化
• このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。gcloud CLI をインストール済みの場合は、gcloud components update コマンドを実行して最新のバージョンを取得します。以前のバージョンの gcloud CLI では、このドキュメントのコマンドを実行できない場合があります。
  注: gcloud CLI がすでにインストールされている場合には、必ず compute/region プロパティを設定してください。主にゾーンクラスタを使用する場合は、代わりに compute/zone を設定します。デフォルトのロケーションを設定することで、gcloud CLI のエラー（One of [--zone, --region] must be supplied: Please specify location など）を防止できます。クラスタのロケーションが設定したデフォルトと異なる場合は、特定のコマンドでロケーションの指定が必要になることがあります。

• バージョン 1.34.0-gke.2201000 以降の既存の Standard クラスタがあることを確認します。新しいクラスタを作成するには、リージョン クラスタの作成をご覧ください。
• 使用するリージョンに十分な TPU の割り当てがあることを確認します。
• v0.2.3 以降の JobSet をインストールします。

制限事項

GKE の TPU All Capacity モードは、Ironwood（TPU7x）バージョンのみをサポートしています。

GKE で TPU All Capacity モードを使用する

このセクションでは、GKE で TPU All Capacity モードを使用するワークフローについて説明します。

1. TPU Cluster Director について理解しておきます。
2. All Capacity モードで TPU 容量をリクエストする。
3. All Capacity モードの予約のトポロジと健全性ステータスを表示する。
4. このドキュメントの手順を完了します。
   1. GKE ノードプールを作成する
   2. ワークロードのスケジュールを設定する
   3. All Capacity モードの予約でノード障害を管理する
5. All Capacity モードの TPU を使用してメンテナンス イベントを管理する。
6. All Capacity モードの TPU ホストの障害を報告して修復する。

All Capacity モードの予約内にノードプールを作成する

GKE の容量モードでは、次の方法でノードプールを作成できます。

• GKE が TPU All Capacity 予約のブロックまたはサブブロックを選択するノードプール。
• TPU All Capacity モードの予約内の特定のブロックまたはサブブロックをターゲットとするノードプール。

GKE は、TPU All Capacity 予約のブロックまたはサブブロックを選択します。

このモードでは、GKE は TPU All Capacity 予約内のノードプールの配置を選択します。このプロセスは、オンデマンド VM や Spot VM など、他の TPU プロビジョニング オプションを使用してノードプールを作成する場合と同様です。

ノードプールを作成するには、--reservation フラグを指定して gcloud container node-pools create コマンドを使用します。--reservation フラグの値として、TPU 予約の完全なリソース名を指定します。

ノードプール作成コマンドの例については、ノードプールを手動で作成するをご覧ください。

予約内のブロックまたはサブブロックをターゲットにする

TPU All Capacity モードでは、TPU 予約内の特定のブロックまたはサブブロックを並列ワークロードのターゲットにできます。この機能は、レイテンシを最小限に抑えるために TPU チップ間の近接性を必要とするワークロードに役立ちます。

1. 使用可能なブロック、サブブロック、ホスト構成を確認します。All Capacity モードの予約のトポロジと健全性ステータスを表示するの手順を完了します。

2. ワークロード ポリシーを作成します。

   注: ノードプールごとに新しいワークロード ポリシーを作成する必要はありません。ワークロード ポリシーは、プロジェクト、リージョン、トポロジごとに一意です。要件に一致するワークロード ポリシーがすでにある場合は、この手順をスキップして、次の手順で既存のポリシーを再利用できます。ワークロード ポリシーのリストを表示するには、gcloud compute resource-policies list --filter="region:REGION" コマンドを使用します。

   gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
       --type=HIGH_THROUGHPUT \
       --accelerator-topology=TPU_TOPOLOGY \
       --project=PROJECT_ID \
       --region=REGION
   

   次のように置き換えます。

   • WORKLOAD_POLICY_NAME: ワークロード ポリシーの名前。
   • TPU_TOPOLOGY: TPU Ironwood（TPU7x）トポロジ。 例: 2x2x2サポートされているすべての Ironwood（TPU7x）トポロジについては、GKE で TPU を計画するをご覧ください。
   • PROJECT_ID: 実際の Google Cloud プロジェクト ID。
   • REGION: ワークロード ポリシーのリージョン。 ワークロード ポリシーはリージョン リソースであり、同じトポロジを共有するノードプール間で再利用できます。

3. ノードプールを作成し、予約の特定のブロックまたはサブブロックをターゲットにするには、--reservation フラグを使用して、予約内のターゲット ブロックまたはサブブロックの完全なリソース名を指定します。

   予約内の特定のブロックをターゲットにするには、次のコマンドを使用します。

   gcloud container node-pools create NODE_POOL_NAME \
       --cluster=CLUSTER_NAME \
       --machine-type=tpu7x-standard-4t \
       --placement-policy=WORKLOAD_POLICY_NAME \
       --zone=ZONE \
       --reservation=project/PROJECT/reservation/RESERVATION_NAME/reservationBlocks/BLOCK_NAME
   

   ブロック内の特定のサブブロックをターゲットにするには、次のコマンドを使用します。

   gcloud container node-pools create NODE_POOL_NAME \
       --cluster=CLUSTER_NAME \
       --machine-type=tpu7x-standard-4t \
       --placement-policy=WORKLOAD_POLICY_NAME \
       --zone=ZONE \
       --reservation=project/PROJECT/reservation/RESERVATION_NAME/reservationBlocks/BLOCK_NAME/reservationSubBlocks/SUB_BLOCK_NAME
   

   次のように置き換えます。

   • NODE_POOL_NAME: 新しいノードプールの名前。
   • CLUSTER_NAME: GKE クラスタの名前。
   • WORKLOAD_POLICY_NAME: 作成したワークロード ポリシーの名前。
   • ZONE: ノードプールのゾーン（us-central1-a など）。
   • PROJECT: 実際の Google Cloud プロジェクト ID。
   • RESERVATION_NAME: TPU 予約の名前。
   • BLOCK_NAME: 予約内の特定のブロック。
   • SUB_BLOCK_NAME: 予約内の特定のサブブロック。

   上記のコマンドでは、

ワークロードのスケジュール

All Capacity モードの TPU VM を使用してノードプールを作成したら、他の TPU ノードプールと同様にワークロードをデプロイできます。TPU All Capacity モードの場合、標準の SLO ベースの予約を使用するノードプールと比較して、ワークロードのスケジューリングに違いはありません。

TPU を使用するワークロードの詳細と例については、TPU スライスノードでワークロードを実行するとマルチスライス ワークロードを実行するをご覧ください。

ノード障害を管理する

TPU All Capacity モードの予約は、ホールドバックなしの予約です。ホールドバックなしとは、TPU 容量全体を受け取ることを意味します。これには、通常はマネージド容量モードでフェイルオーバー用にGoogle Cloud が保持する部分も含まれます。

TPU All Capacity モードでは、ハードウェア障害などの問題により VM に障害が発生した場合、 Google Cloud は同じホストで VM の復元を試みます（インプレース修復）。そのため、インフラストラクチャの障害時にワークロードの再スケジュールに対応できるように、予備容量を維持することをおすすめします。

ノードの障害と復旧

TPU All Capacity モードの予約でノードに障害が発生すると、次のイベントが発生します。

• Google Cloud は、失敗した Compute Engine VM インスタンスの修復イベントを開始します。このプロセスでは、VM を RUNNING ステータスに復元し、GKE ノードのステータスを READY に戻そうとします。
• TPU VM が修復状態になり、そのノードで実行されているワークロードは、フェイルオーバー ポリシーに応じて失敗する可能性があります。1 つ以上の VM で障害が発生していても、ノードプールのステータスは ERROR に変わりません。

ノードのヘルス ステータスをモニタリングする手順は次のとおりです。

1. ノードプール内のノードを一覧表示します。

   kubectl get nodes
   

   失敗したノードのステータスは NotReady です。

2. Compute Engine ノードのステータスをモニタリングします。

   All Capacity モードの TPU VM の場合は、gcloud compute instances describe を使用します。このコマンドは、トポロジの物理ステータスも提供し、ホスト、サブブロック、ブロックの詳細を確認できます。

   gcloud compute instances describe VM_NAME \
           --format="table[box,title=VM-Position](resourceStatus.physical_host_topology:label=location)" \
           --zone=ZONE
   

   次のように置き換えます。

   • VM_NAME: TPU VM インスタンスの名前。
   • ZONE: VM のゾーン（例: us-central1-a）。

   All Capacity モードの容量に関するトポロジと健全性情報を取得する手順については、All Capacity モードの予約のトポロジと健全性ステータスを表示するをご覧ください。

メンテナンスを管理する

中断の可能性に対処し、ワークロードの復元力を維持するには、GKE メンテナンス ポリシーを使用して個々のノードのメンテナンスを管理します。詳細については、メンテナンスの時間枠とメンテナンスの除外をご覧ください。

GKE は、All Capacity モードの予約で TPU VM のグループ メンテナンスをサポートしていません。サブブロック、ブロック、予約レベルでグループ メンテナンスを実行するには、Compute Engine API を使用します。詳細については、All Capacity モードの TPU を使用したメンテナンス イベントの管理をご覧ください。

クリーンアップ

Google Cloud アカウントで不要な料金が発生しないようにするには、ワークロードのスケジュールがなくなった TPU ノードプールを削除します。実行中のワークロードを正常に終了する必要がある場合は、kubectl drain コマンドを使用してワークロードをクリーンアップしてから、ノードプールを削除します。

1. TPU ノードプールを削除します。

   gcloud container node-pools delete NODE_POOL_NAME \
       --location=LOCATION \
       --cluster=CLUSTER_NAME
   

   次のように置き換えます。

   • NODE_POOL_NAME: ノードプールの名前。
   • CLUSTER_NAME: クラスタの名前。
   • LOCATION: クラスタのコンピューティング ロケーション。

次のステップ

• GKE の TPU の詳細を確認する。
• GKE Standard に TPU ワークロードをデプロイする方法を確認する。
• TPU Cluster Director の詳細を確認する。