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

Standard

このドキュメントでは、GKE の TPU All Capacity モード機能を使用してワークロードをデプロイして管理する方法について説明します。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 の割り当てがあることを確認します。
• JobSet v0.2.3 以降をインストールします。

制限事項

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 の All Capacity モードでは、次の方法でノードプールを作成できます。

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

GKE が TPU All Capacity 予約のブロックまたはサブブロックを選択する

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

ノードプールを作成するには、gcloud container node-pools create コマンドを使用し、--reservation フラグを指定します。 --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 を使用したメンテナンス イベントの管理をご覧ください。

クリーンアップ

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

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

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

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

   • NODE_POOL_NAME: ノードプールの名前。
   • CLUSTER_NAME: クラスタの名前。
   • LOCATION: クラスタの Compute Engine のロケーション。

次のステップ

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