GKE Autopilot に TPU ワークロードをデプロイする

Autopilot

このページでは、Google Kubernetes Engine（GKE）の Autopilot クラスタに Cloud TPU アクセラレータ（TPU）を使用してデプロイすることで、ML ワークロードを高速化する方法について説明します。このガイダンスは、ML アプリケーションフレームワークに適したライブラリを選択したり、GKE で最適に実行されるように TPU ワークロードを設定したりするのに役立ちます。また、デプロイ後にワークロードをモニタリングすることもできます。

このページは、TPU で ML ワークロードを準備して実行するプラットフォーム管理者とオペレーター、データおよび AI スペシャリスト、アプリケーションデベロッパーを対象としています。 Google Cloud のコンテンツで参照する一般的なロール、責任、タスク例の詳細については、GKE ユーザーの一般的なロールとタスクをご覧ください。

このページを読む前に、次のリソースについてよく理解しておいてください。

Autopilot での TPU の仕組み

Autopilot ワークロードで TPU を使用するには、ワークロードマニフェストで次のように指定します。

spec.nodeSelector フィールドの TPU バージョン。
spec.nodeSelector フィールドの TPU トポロジ。トポロジは、指定されたバージョンの TPU でサポートされている必要があります。
spec.containers.resources.requests フィールドと spec.containers.resources.limits フィールドの TPU チップの数。

ワークロードをデプロイすると、GKE はリクエストされた TPU 構成を持つノードをプロビジョニングし、ノードに Pod をスケジュールします。GKE は各ワークロードを独自のノードに配置するため、各 Pod は中断のリスクを最小限に抑えつつ、ノードのすべてのリソースにアクセスできます。

Autopilot の TPU は、次の機能に対応しています。

TPU 構成を計画する

このガイドに沿って TPU ワークロードをデプロイする前に、モデルと必要なメモリ量に基づいて TPU 構成を計画します。詳しくは、TPU 構成の計画をご覧ください。

料金

料金については、Autopilot の料金をご覧ください。

始める前に

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

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 など）を防止できます。クラスタのロケーションが設定したデフォルトと異なる場合は、特定のコマンドでロケーションの指定が必要になることがあります。

GKE バージョン 1.32.3-gke.1927000 以降を実行している Autopilot クラスタがあることを確認します。手順については、Autopilot クラスタを作成するをご覧ください。
予約済みの TPU を使用するには、特定の容量がすでに予約してあることを確認してください。手順については、予約を使用するをご覧ください。

TPU とその他の GKE リソースの割り当てを確保する

以降のセクションでは、GKE で TPU を使用するときに十分な割り当てを確保できるようにします。

TPU スライスノードを作成するには、既存の容量予約を使用する場合を除き、TPU の割り当てが必要です。予約済みの TPU を使用している場合は、このセクションをスキップしてください。

GKE で TPU スライスノードを作成するには、Cloud TPU API の割り当て（tpu.googleapis.com）ではなく、Compute Engine API の割り当て（compute.googleapis.com）が必要です。割り当ての名前は、通常の Autopilot Pod と Spot Pod で異なります。

TPU 用の Compute Engine API の割り当ての上限と現在の使用量を確認するには、次の操作を行います。

Google Cloud コンソールで [割り当て] ページに移動します。

[割り当て] に移動

[ フィルタ] ボックスで次の操作を行います。

次の表を使用して、TPU のバージョンと cloud.google.com/gke-tpu-accelerator ノードセレクタの値に基づいて割り当てのプロパティを選択してコピーします。たとえば、cloud.google.com/gke-tpu-accelerator ノードセレクタの値が tpu-v5-lite-podslice のオンデマンド TPU v5e ノードを作成する場合は、「Name: TPU v5 Lite PodSlice chips」と入力します。

TPU バージョン、`cloud.google.com/gke-tpu-accelerator`	オンデマンドインスタンスの割り当てのプロパティと名前	Spot² インスタンスの割り当てのプロパティと名前
TPU v3、 `tpu-v3-device`	Dimensions (e.g. location): tpu_family:CT3	該当なし
TPU v3、 `tpu-v3-slice`	Dimensions (e.g. location): tpu_family:CT3P	該当なし
TPU v4、 `tpu-v4-podslice`	Name: TPU v4 PodSlice chips	Name: Preemptible TPU v4 PodSlice chips
TPU v5e、 `tpu-v5-lite-podslice`	Name: TPU v5 Lite PodSlice chips	Name: Preemptible TPU v5 Lite Podslice chips
TPU v5p、 `tpu-v5p-slice`	Name: TPU v5p chips	Name: Preemptible TPU v5p chips
TPU Trillium、 `tpu-v6e-slice`	Dimensions (e.g. location): tpu_family:CT6E	Name: Preemptible TPU slices v6e
Ironwood（TPU7x）（プレビュー）、 `tpu7x`	Dimensions (e.g. location): tpu_family:tpu7x	Name: Preemptible TPU slices tpu7x

[項目（ロケーションなど）] プロパティを選択し、「region:」に続けて、GKE で TPU を作成するリージョンの名前を入力します。たとえば、ゾーン us-west4-a で TPU スライスノードを作成する場合は、「region:us-west4」と入力します。TPU の割り当てはリージョン単位であるため、同じリージョン内のすべてのゾーンで同じ TPU の割り当てが消費されます。

入力したフィルタに一致する割り当てがない場合、プロジェクトには目的のリージョンで指定した割り当てのいずれも付与されていないため、TPU 割り当ての調整をリクエストする必要があります。

TPU 予約が作成されると、対応する割り当ての上限と現在の使用量の値は、TPU 予約のチップ数の分だけ増加します。たとえば、cloud.google.com/gke-tpu-accelerator ノードセレクタの値が tpu-v5-lite-podslice である 16 個の TPU v5e チップの予約を作成すると、関連するリージョンの TPU v5 Lite PodSlice chips 割り当ての上限と現在の使用量の両方が 16 増加します。

追加の GKE リソースの割り当て

GKE がリソースを作成するリージョンで、次の GKE 関連の割り当てを増やす必要がある場合があります。

Persistent Disk SSD（GB）割り当て: 各 Kubernetes ノードのブートディスクにはデフォルトで 100 GB 必要です。そのため、この割り当ては、少なくとも、作成する予定の GKE ノードの最大数と 100 GB の積（ノード × 100 GB）以上の値に設定する必要があります。
使用中の IP アドレスの割り当て: 各 Kubernetes ノードは 1 つの IP アドレスを消費します。そのため、この割り当てには、少なくとも作成することが予想される GKE ノードの最大数を設定する必要があります。
max-pods-per-node がサブネット範囲と一致していることを確認する: 各 Kubernetes ノードは、Pod にセカンダリ IP 範囲を使用します。たとえば、max-pods-per-node が 32 の場合、64 個の IP アドレスが必要になります。これはノードあたり /26 サブネットに相当します。この範囲は他のクラスタと共有しないでください。IP アドレス範囲が使い果たされないようにするには、--max-pods-per-node フラグを使用して、ノードでスケジュールできる Pod の数を制限します。max-pods-per-node の割り当ては、作成する予定の GKE ノードの最大数と同じ値に設定する必要があります。

割り当ての増加をリクエストするには、割り当ての調整をリクエストするをご覧ください。

TPU アプリケーションを準備する

TPU ワークロードには、次の準備要件があります。

JAX、PyTorch、TensorFlow などのフレームワークは、libtpu 共有ライブラリを使用して TPU VM にアクセスします。libtpu には、XLA コンパイラ、TPU ランタイムソフトウェア、TPU ドライバが含まれています。PyTorch と JAX の各リリースには、特定の libtpu.so バージョンが必要です。パッケージバージョンの競合を避けるため、JAX AI 画像を使用することをおすすめします。GKE で TPU を使用するには、次のバージョンを使用してください。 tpu7x

TPU タイプ	`libtpu.so` のバージョン
Ironwood（TPU7x）（プレビュー）	推奨される JAX AI 画像: jax0.8.1-rev1 以降推奨される jax[tpu] バージョン: v0.8.1
TPU Trillium（v6e） `tpu-v6e-slice`	推奨される JAX AI 画像: jax0.4.35-rev1 以降推奨される jax[tpu] バージョン: v0.4.9 以降推奨される torchxla[tpuvm] のバージョン: v2.1.0 以降
TPU v5e `tpu-v5-lite-podslice`	推奨される JAX AI 画像: jax0.4.35-rev1 以降推奨される jax[tpu] バージョン: v0.4.9 以降推奨される torchxla[tpuvm] のバージョン: v2.1.0 以降
TPU v5p `tpu-v5p-slice`	推奨される JAX AI 画像: jax0.4.35-rev1 以降推奨される jax[tpu] バージョン: 0.4.19 以降。推奨される torchxla[tpuvm] バージョン: 毎晩更新されるバージョンの 2023 年 10 月 23 日付のビルドを使用することをおすすめします。
TPU v4 `tpu-v4-podslice`	推奨される JAX AI 画像: jax0.4.35-rev1 以降推奨される jax[tpu]: v0.4.4 以降推奨される torchxla[tpuvm]: v2.0.0 以降
TPU v3 `tpu-v3-slice` `tpu-v3-device`	推奨される JAX AI 画像: jax0.4.35-rev1 以降推奨される jax[tpu]: v0.4.4 以降推奨される torchxla[tpuvm]: v2.0.0 以降

ワークロードマニフェストに Kubernetes ノードセレクタを追加して、定義した TPU マシンタイプと TPU トポロジで、GKE が TPU ワークロードをスケジュールできるようにします。
```
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
    cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
    cloud.google.com/placement-policy-name: WORKLOAD_POLICY # Required only for Ironwood (TPU7x)
  
```
次のように置き換えます。
- TPU_ACCELERATOR: TPU アクセラレータの名前。たとえば、tpu7x-standard-4t を使用します。
- TPU_TOPOLOGY: TPU スライスの物理トポロジ。トポロジの形式は TPU のバージョンによって異なります。たとえば、2x2x2 を使用します。詳細については、GKE で TPU を計画するをご覧ください。
- WORKLOAD_POLICY: TPU Pod の配置に使用するワークロードポリシーの名前。このノードセレクタは、Ironwood（TPU7x）でのみ必要です。

ワークロードの準備が完了したら、TPU を使用する Job を実行できます。

GKE で TPU をプロビジョニングする場合のオプション

GKE で TPU をプロビジョニングするには、次の構成オプションがあります。

ワークロードリクエスト: spec.nodeSelector フィールドで TPU のバージョンとトポロジを指定し、spec.containers.resources セクションで TPU チップの数を指定します。ワークロードをデプロイすると、GKE は正しい TPU 構成でノードを自動的にプロビジョニングし、各ワークロードを独自の専用ノードに配置して、ノードのリソースに完全にアクセスできるようにします。手順については、ワークロードで TPU をリクエストするをご覧ください。

カスタムのコンピューティングクラスを使用して TPU を一元的にプロビジョニングする

以下のセクションでは、カスタム ComputeClass を作成し、ComputeClass で定義された TPU を消費する Job を作成する方法について説明します。

カスタム ComputeClass を作成する

TPU ルールに従うカスタム ComputeClass を作成する手順は、Ironwood（TPU7x）を使用するか、以前の TPU バージョンを使用するかによって異なります。

Ironwood（TPU7x）

ワークロードポリシーを作成します。この手順は、選択したトポロジに応じてマルチホストノードプールを作成する場合にのみ必要です。単一ホストのノードプールを使用する場合は、この手順をスキップします。
```
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）トポロジの詳細については、トポロジセクションをご覧ください。
- PROJECT_ID: 実際の Google Cloud プロジェクト ID。
- REGION: ワークロードポリシーのリージョン。ワークロードポリシーはリージョンリソースであり、ノードプール全体で使用できます。

次のマニフェストを tpu-compute-class.yaml として保存します。

apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: tpu-class
spec:
  priorities:
    - tpu:
        type: tpu7x
        topology: TPU_TOPOLOGY
        count: 4
      placement:
        policyName: WORKLOAD_POLICY_NAME
  nodePoolAutoCreation:
    enabled: true

（省略可）特定の予約またはサブブロックを消費できます。たとえば、ComputeClass マニフェストに次の specs を追加できます。
```
  reservations:
    affinity: Specific
    specific:
      - name: RESERVATION_NAME
        reservationBlock:
          name: RESERVATION_BLOCK_NAME
          reservationSubBlock:
            name: RESERVATION_SUB_BLOCK_NAME
```
次のように置き換えます。
- RESERVATION_NAME: Compute Engine の容量予約の名前。
- RESERVATION_BLOCK_NAME: Compute Engine の容量予約ブロックの名前。
- RESERVATION_SUB_BLOCK_NAME: Compute Engine の容量予約サブブロックの名前。
詳細については、予約済みゾーンリソースの使用をご覧ください。

その他の TPU バージョン

TPU 用に構成されたカスタム ComputeClass を使用して v3、v4、v5p、v5e、v6e（Trillium）TPU をプロビジョニングするには、次の操作を行います。

次のマニフェストを tpu-compute-class.yaml として保存します。

apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: tpu-class
spec:
  priorities:
  - tpu:
      type: TPU_TYPE
      count: NUMBER_OF_CHIPS
      topology: TOPOLOGY
  - spot: true
    tpu:
      type: {"<var>"}}TPU_TYPE
      count: NUMBER_OF_CHIPS
      topology: TOPOLOGY
  - flexStart:
      enabled: true
    tpu:
      type: {"<var>"}}TPU_TYPE
      count: NUMBER_OF_CHIPS
      topology: TOPOLOGY
  nodePoolAutoCreation:
    enabled: true

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

TPU_TYPE: 使用する TPU タイプ（tpu-v4-podslice など）。GKE でサポートされている値にする必要があります。
TOPOLOGY: スライス内の TPU チップの配置（2x2x4 など）。選択した TPU タイプでサポートされているトポロジである必要があります。
NUMBER_OF_CHIPS: 使用するコンテナの TPU チップの数。limits と requests の値は同じにする必要があります。

ComputeClass をデプロイします。
```
kubectl apply -f tpu-compute-class.yaml
```
カスタム ComputeClasses と TPU の詳細については、TPU の構成をご覧ください。

TPU を消費する Job を作成する

次のマニフェストを tpu-job.yaml として保存します。

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/compute-class: tpu-class
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS

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

NUMBER_OF_CHIPS: 使用するコンテナの TPU チップの数。limits と requests の値は同じにする必要があります。また、選択したカスタム ComputeClass の CHIP_COUNT 値と同じにする必要があります。
MEMORY_SIZE: TPU が使用するメモリの最大量。メモリ上限は、使用する TPU のバージョンとトポロジによって異なります。詳しくは、アクセラレータの最小値と最大値をご覧ください。
NUMBER_OF_CHIPS: 使用するコンテナの TPU チップの数。limits と requests の値は同じにする必要があります。

Job をデプロイします。
```
kubectl create -f tpu-job.yaml
```
この Job を作成すると、GKE は自動的に次の処理を行います。
- Pod を実行するノードをプロビジョニングします。指定した TPU タイプ、トポロジ、リソースリクエストに応じて、これらのノードは単一ホストスライスまたはマルチホストスライスのいずれかになります。GKE は最優先の TPU リソースの可用性に応じて優先度の低いリソースにフォールバックし、取得可能性を最大化できます。
- Pod に taint を追加してノードに toleration を追加し、他のワークロードが TPU ワークロードと同じノードで実行されないようにします。
詳しくは、カスタム ComputeClasses についてをご覧ください。
このセクションを終了した後、作成したリソースを削除すると、それ以上の請求は発生しません。
```
kubectl delete -f tpu-job.yaml
```

ワークロードで TPU をリクエストする

このセクションでは、Autopilot で TPU をリクエストする Job の作成方法について説明します。TPU を必要とするワークロードでは、次のことを指定する必要があります。

TPU のバージョンとトポロジのノードセレクタ
ワークロード内のコンテナの TPU チップの数

サポートされている TPU バージョン、トポロジ、スライス内の TPU チップとノードの数については、TPU のバージョンを選択するをご覧ください。

ワークロード内の TPU リクエストに関する考慮事項

TPU を使用できるのは、Pod 内の 1 つのコンテナのみです。コンテナがリクエストする TPU チップの数は、スライス内のノードに接続している TPU チップ数と同じにする必要があります。たとえば、2x4 トポロジで TPU v5e（tpu-v5-lite-podslice）をリクエストする場合は、次のいずれかのリクエストを実行できます。

4 チップ。これにより、それぞれ 4 個の TPU チップを持つ 2 個のマルチホストノードが作成されます。
8 チップ。8 個の TPU チップを持つ単一ホストノードが 1 つ作成されます。

費用対効果を最大にするためのベストプラクティスは、リクエストするスライス内のすべての TPU を常に使用することです。4 個の TPU チップを持つ 2 個のノードのマルチホストスライスをリクエストする場合は、両方のノードで実行され、スライスの 8 個の TPU チップをすべて消費するワークロードをデプロイする必要があります。

TPU をリクエストするワークロードを作成する

TPU をリクエストするジョブを次の手順で作成します。マルチホスト TPU スライスで実行されるワークロードがある場合は、名前でワークロードを選択するヘッドレス Service も作成する必要があります。このヘッドレス Service は、ワークロード内の Pod を指すように Kubernetes DNS 構成を更新して、マルチホストスライス内の異なるノード上の Pod が相互に通信できるようにします。

次のマニフェストを tpu-autopilot.yaml として保存します。

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      # Optional: Run in GKE Sandbox
      # runtimeClassName: gvisor
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_TYPE
        cloud.google.com/gke-tpu-topology: TOPOLOGY
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS

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

TPU_TYPE: 使用する TPU タイプ（tpu-v4-podslice など）。GKE でサポートされている値にする必要があります。
TOPOLOGY: スライス内の TPU チップの配置（2x2x4 など）。選択した TPU タイプでサポートされているトポロジである必要があります。
NUMBER_OF_CHIPS: 使用するコンテナの TPU チップの数。limits と requests の値は同じにする必要があります。
MEMORY_SIZE: TPU が使用するメモリの最大量。メモリ上限は、使用する TPU のバージョンとトポロジによって異なります。詳しくは、アクセラレータの最小値と最大値をご覧ください。

必要に応じて、次のフィールドを変更することもできます。

image: 使用する JAX AI イメージ。マニフェストの例では、このフィールドは最新の JAX AI イメージに設定されています。別のバージョンを設定するには、現行の JAX AI イメージのリストをご覧ください。
runtimeClassname: gvisor: この Pod を GKE Sandbox で実行できるようにする設定。使用する場合は、この行のコメント化を解除します。GKE Sandbox では、TPU バージョン v4 以降がサポートされています。詳しくは、GKE Sandbox をご覧ください。

Job をデプロイします。
```
kubectl create -f tpu-autopilot.yaml
```
この Job を作成すると、GKE は自動的に次の処理を行います。
1. Pod を実行するノードをプロビジョニングします。指定した TPU タイプ、トポロジ、リソースリクエストに応じて、これらのノードは単一ホストスライスまたはマルチホストスライスのいずれかになります。
2. Pod に taint を追加してノードに toleration を追加し、他のワークロードが TPU ワークロードと同じノードで実行されないようにします。
このセクションを完了したら、作成したワークロードを削除すれば、それ以上の請求は発生しません。
```
kubectl delete -f tpu-autopilot.yaml
```

TPU とコレクションスケジューリングをリクエストするワークロードを作成する

TPU Trillium では、コレクションスケジューリングを使用して TPU スライスノードをグループ化できます。これらの TPU スライスノードをグループ化すると、ワークロードの需要に合わせてレプリカの数を簡単に調整できます。 Google Cloud はソフトウェアアップデートを制御し、コレクション内に十分なスライスが常にあり、トラフィックの処理に使用できるようにします。

TPU Trillium は、推論ワークロードを実行する単一ホストノードプールとマルチホストノードプールのコレクションスケジューリングをサポートしています。次の表に示すのは、コレクションスケジューリング動作が使用する TPU スライスのタイプにどのように依存するかです。

マルチホスト TPU スライス: GKE はマルチホスト TPU スライスをグループ化してコレクションを形成します。各 GKE ノードプールは、このコレクション内のレプリカです。コレクションを定義するには、マルチホスト TPU スライスを作成し、コレクションに一意の名前を割り当てます。コレクションに TPU スライスを追加するには、同じコレクション名とワークロードタイプを持つ別のマルチホスト TPU スライスノードプールを作成します。
単一ホストの TPU スライス: GKE は、単一ホストの TPU スライスノードプール全体をコレクションと見なします。コレクションに TPU スライスを追加するには、単一ホストの TPU スライスノードプールのサイズを変更します。

コレクションスケジューリングの制限については、コレクションスケジューリングの仕組みをご覧ください。

マルチホスト TPU スライスを使用する

マルチホスト TPU スライスノードでのコレクションスケジューリングは、バージョン 1.31.2-gke.1537000 以降の Autopilot クラスタで使用できます。2x4 トポロジのマルチホスト TPU スライスノードは、1.31.2-gke.1115000 以降でのみサポートされます。マルチホスト TPU スライスノードを作成してコレクションとしてグループ化するには、ワークロード仕様に次の Kubernetes ラベルを追加します。

cloud.google.com/gke-nodepool-group-name: 各コレクションには、クラスタレベルで一意の名前が必要です。cloud.google.com/gke-nodepool-group-name ラベルの値は、クラスタラベルの要件に沿っている必要があります。

cloud.google.com/gke-workload-type: HIGH_AVAILABILITY

たとえば、次のコードブロックは、マルチホスト TPU スライスを含むコレクションを定義します。

  nodeSelector:
    cloud.google.com/gke-nodepool-group-name: ${COLLECTION_NAME}
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 4x4
...

単一ホストの TPU スライスを使用する

単一ホスト TPU スライスノードでのコレクションスケジューリングは、バージョン 1.31.2-gke.1088000 以降の Autopilot クラスタで使用できます。単一ホストの TPU スライスノードを作成してコレクションとしてグループ化するには、ワークロード仕様に cloud.google.com/gke-workload-type:HIGH_AVAILABILITY ラベルを追加します。

たとえば、次のコードブロックは、単一ホストの TPU スライスを含むコレクションを定義します。

  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 2x2
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
  ...

カスタムコンピューティングクラスを使用してコレクションをデプロイする

カスタムコンピューティングクラスを使用して TPU ワークロードとコレクションスケジューリングをリクエストするワークロードをデプロイする方法の詳細については、TPU マルチホストコレクションと TPU SLO のワークロードタイプを定義するをご覧ください。

例: マルチホストスライスで TPU チップの合計数を表示する

次のワークロードは、マルチホスト TPU スライス内のノード全体の TPU チップ数を返します。マルチホストスライスを作成する場合、ワークロードに次のパラメータを設定します。

TPU バージョン: TPU v4
トポロジ: 2x2x4

このバージョンとトポロジの選択により、マルチホストスライスが作成されます。

次のマニフェストを available-chips-multihost.yaml として保存します。

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice # Node selector to target TPU v4 slice nodes.
        cloud.google.com/gke-tpu-topology: 2x2x4 # Specifies the physical topology for the TPU slice.
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())' # Python command to count available TPU chips.
        resources:
          requests:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4 # Request 4 TPU chips for this workload.
          limits:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4 # Limit to 4 TPU chips for this workload.

マニフェストをデプロイします。
```
kubectl create -f available-chips-multihost.yaml
```
GKE は、4 つの VM（マルチホスト TPU スライス）を使用して TPU v4 スライスを実行します。スライスには、相互接続された 16 個の TPU チップがあります。

Job によって 4 つの Pod が作成されたことを確認します。

kubectl get pods

出力は次のようになります。

NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s

いずれかの Pod のログを取得します。
```
kubectl logs POD_NAME
```
POD_NAME は、作成した Pod の名前に置き換えます。例: tpu-job-podslice-0-5cd8r

出力は次のようになります。
```
TPU cores: 16
```

省略可: ワークロードを削除します。

kubectl delete -f available-chips-multihost.yaml

例: 単一ノードで TPU チップを表示する

次のワークロードは、特定のノードに関連付けられている TPU チップの数を表示する静的 Pod です。単一ホストノードを作成するには、ワークロードに次のパラメータを設定します。

TPU バージョン: TPU v5e
トポロジ: 2x4

このバージョンとトポロジを選択すると、単一ホストのスライスが作成されます。

次のマニフェストを available-chips-singlehost.yaml として保存します。

apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice # Node selector to target TPU v5e slice nodes.
    cloud.google.com/gke-tpu-topology: 2x4 # Specify the physical topology for the TPU slice.
  containers:
  - name: tpu-job
    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8 # Request 8 TPU chips for this container.
      limits:
        google.com/tpu: 8 # Limit to 8 TPU chips for this container.

マニフェストをデプロイします。
```
kubectl create -f available-chips-singlehost.yaml
```
GKE は、TPU v5e を使用する 8 つの単一ホスト TPU スライスを含むノードをプロビジョニングします。各 TPU ノードには 8 個の TPU チップがあります（単一ホストの TPU スライス）。
Pod のログを取得します。
```
kubectl logs tpu-job-jax-v5
```
出力は次のようになります。
```
Total TPU chips: 8
```

省略可: ワークロードを削除します。

  kubectl delete -f available-chips-singlehost.yaml

TPU を観察、モニタリングする

ダッシュボード

Google Cloud コンソールのノードプールオブザーバビリティが一般提供になりました。GKE の TPU マルチホストノードプールのステータスを表示するには、Cloud Monitoring が提供する [GKE TPU Node Pool Status] ダッシュボードに移動します。

[GKE TPU Node Pool Status] に移動

このダッシュボードでは、マルチホスト TPU ノードプールの健全性に関する包括的な分析情報を確認できます。詳細については、TPU ノードとノードプールの健全性指標をモニタリングするをご覧ください。

Google Cloud コンソールの [Kubernetes クラスタ] ページの [オブザーバビリティ] タブにも TPU オブザーバビリティ指標が表示されます（[アクセラレータ > TPU] 見出しの下の TPU 使用率など）。詳細については、オブザーバビリティ指標を表示するをご覧ください。

TPU ダッシュボードは、GKE クラスタでシステム指標が有効になっている場合にのみ表示されます。

ランタイム指標

GKE バージョン 1.27.4-gke.900 以降、JAX バージョン 0.4.14 以降を使用し containerPort: 8431 を指定する TPU ワークロードでは、TPU 使用率の指標を GKE システム指標としてエクスポートします。Cloud Monitoring では、TPU ワークロードのランタイムパフォーマンスをモニタリングするために、次の指標を使用できます。

デューティサイクル: 過去のサンプリング期間（60 秒）において、TensorCore が TPU チップでアクティブに処理していた時間の割合。割合が大きいほど、TPU 使用率が高くなります。
メモリ使用量: 割り当てられたアクセラレータメモリの量（バイト単位）。60 秒ごとにサンプリングされます。
メモリの総容量: アクセラレータの総メモリ（バイト単位）。60 秒ごとにサンプリングされます。

これらの指標は、Kubernetes ノード（k8s_node）と Kubernetes コンテナ（k8s_container）のスキーマにあります。

Kubernetes コンテナ:

kubernetes.io/container/accelerator/duty_cycle
kubernetes.io/container/accelerator/memory_used
kubernetes.io/container/accelerator/memory_total

Kubernetes ノード:

kubernetes.io/node/accelerator/duty_cycle
kubernetes.io/node/accelerator/memory_used
kubernetes.io/node/accelerator/memory_total

TPU ノードとノードプールの健全性指標をモニタリングする

トレーニングジョブでエラーが発生した場合や、トレーニングジョブが失敗して終了した場合は、基盤となるインフラストラクチャに関連する指標を確認して、中断の原因が基盤となるノードまたはノードプールの問題であるかどうかを判断できます。

ノードのステータス

GKE バージョン 1.32.1-gke.1357001 以降では、次の GKE システム指標が GKE ノードの状態を公開します。

kubernetes.io/node/status_condition

condition フィールドは、Ready、DiskPressure、MemoryPressure などのノードの状態を報告します。status フィールドには、その状態について報告されたステータス（True、False、Unknown）が表示されます。これは、k8s_node モニタリング対象リソースタイプの指標です。

この PromQL クエリは、特定のノードが Ready かどうかを示します。

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    node_name="NODE_NAME",
    condition="Ready",
    status="True"}

クラスタの問題のトラブルシューティングを行うには、他の状態を示しているノードを確認します。

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition!="Ready",
    status="True"}

Ready ではないノードを確認することもできます。

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition="Ready",
    status="False"}

データがない場合、ノードは準備完了です。ステータス状態は 60 秒ごとにサンプリングされます。

次のクエリを使用すると、フリート全体のノードのステータスを把握できます。

avg by (condition,status)(
  avg_over_time(
    kubernetes_io:node_status_condition{monitored_resource="k8s_node"}[${__interval}]))

ノードプールのステータス

k8s_node_pool モニタリング対象リソースの次の GKE システム指標は、GKE ノードプールのステータスを公開します。

kubernetes.io/node_pool/status

この指標は、マルチホスト TPU ノードプールに対してのみ報告されます。

status フィールドは、ノードプールのステータス（Provisioning、Running、Error、Reconciling、Stopping など）を報告します。ステータスの更新は、GKE API オペレーションの完了後に行われます。

特定のノードプールが Running ステータスかどうかを確認するには、次の PromQL クエリを使用します。

kubernetes_io:node_pool_status{
    monitored_resource="k8s_node_pool",
    cluster_name="CLUSTER_NAME",
    node_pool_name="NODE_POOL_NAME",
    status="Running"}

プロジェクト内のノードプールの数をステータス別にモニタリングするには、次の PromQL クエリを使用します。

count by (status)(
  count_over_time(
    kubernetes_io:node_pool_status{monitored_resource="k8s_node_pool"}[${__interval}]))

ノードプールの可用性

次の GKE システム指標は、マルチホスト TPU ノードプールが使用可能かどうかを示します。

kubernetes.io/node_pool/multi_host/available

この指標の値は、ノードプール内のすべてのノードが使用可能な場合は True になり、それ以外の場合は False になります。この指標は 60 秒ごとにサンプリングされます。

プロジェクトでマルチホスト TPU ノードプールが使用可能かどうかを確認するには、次の PromQL クエリを使用します。

avg by (node_pool_name)(
  avg_over_time(
    kubernetes_io:node_pool_multi_host_available{
      monitored_resource="k8s_node_pool",
      cluster_name="CLUSTER_NAME"}[${__interval}]))

ノードの中断回数

次の GKE システム指標は、最後のサンプリング以降の GKE ノードの中断回数を報告します（この指標は 60 秒ごとにサンプリングされます）。

kubernetes.io/node/interruption_count

interruption_type（TerminationEvent、MaintenanceEvent、PreemptionEvent など）フィールドと interruption_reason（HostError、Eviction、AutoRepair など）フィールドは、ノードが中断された理由を把握するのに役立ちます。

プロジェクトのクラスタ内の TPU ノードで発生した中断とその原因の内訳を取得するには、次の PromQL クエリを使用します。

  sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node"}[${__interval}]))

ホストメンテナンスイベントのみを表示するには、interruption_reason の HW/SW Maintenance 値をフィルタするようにクエリを更新します。次の PromQL クエリを使用します。

  sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node", interruption_reason="HW/SW Maintenance"}[${__interval}]))

ノードプールごとに集計された中断回数を表示するには、次の PromQL クエリを使用します。

  sum by (node_pool_name,interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_pool_interruption_count{monitored_resource="k8s_node_pool", interruption_reason="HW/SW Maintenance", node_pool_name=NODE_POOL_NAME }[${__interval}]))

ノードプールの復旧時間（TTR）

次の GKE システム指標は、GKE マルチホスト TPU ノードプールの復旧期間の分布を報告します。

kubernetes.io/node_pool/accelerator/times_to_recover

この指標に記録された各サンプルは、ダウンタイム期間からの単一のノードプール復旧イベントを示します。

この指標は、マルチホスト TPU ノードプールの復旧時間と中断間隔を追跡するのに役立ちます。

次の PromQL クエリを使用すると、クラスタの過去 7 日間の平均復旧時間（MTTR）を計算できます。

sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_sum{
    monitored_resource="k8s_node_pool", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_count{
    monitored_resource="k8s_node_pool",cluster_name="CLUSTER_NAME"}[7d]))

ノードプールの中断間隔（TBI）

ノードプールの中断間隔は、インフラストラクチャが中断するまでの稼働時間を測定します。これは、一定期間の平均として計算されます。分子はインフラストラクチャの稼働時間の合計を測定し、分母はインフラストラクチャの中断時間の合計を測定します。

次の PromQL の例は、特定のクラスタの 7 日間の平均中断間隔（MTBI）を示しています。

sum(count_over_time(
  kubernetes_io:node_memory_total_bytes{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_interruption_count{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))

ホスト指標

GKE バージョン 1.28.1-gke.1066000 以降では、TPU スライスの VM は TPU 使用率の指標を GKE システム指標としてエクスポートします。Cloud Monitoring では、次の指標を使用して TPU ホストのパフォーマンスをモニタリングできます。

TensorCore の使用率: 使用されている TensorCore の現在の割合。TensorCore の値は、マトリックス乗算ユニット（MXU）およびベクトルユニットの合計と等しくなります。TensorCore の使用率の値は、過去のサンプル期間（60 秒）に実行された TensorCore オペレーションを、同じ期間にサポートされている TensorCore オペレーションの数で割ったものです。値が大きいほど、使用率が高いことを意味します。
メモリ帯域幅の使用率: 現在使用されているアクセラレータメモリ帯域幅の割合。サンプル期間（60 秒）で使用されたメモリ帯域幅を、同じサンプル期間でサポートされる最大帯域幅で割って計算されます。

これらの指標は、Kubernetes ノード（k8s_node）と Kubernetes コンテナ（k8s_container）のスキーマにあります。

Kubernetes コンテナ:

kubernetes.io/container/accelerator/tensorcore_utilization
kubernetes.io/container/accelerator/memory_bandwidth_utilization

Kubernetes ノード:

kubernetes.io/node/accelerator/tensorcore_utilization
kubernetes.io/node/accelerator/memory_bandwidth_utilization

詳細については、Kubernetes の指標と GKE のシステム指標をご覧ください。

ロギング

TPU VM を含む GKE ノード上で実行されているコンテナによって出力されたログは、GKE ロギングエージェントが収集し、Logging に送信して、Logging に表示します。

Autopilot での TPU ワークロードに関する推奨事項

次の推奨事項により、TPU ワークロードの効率が向上する可能性があります。

スケールダウンまたはノードのアップグレードのため、GKE が Pod を終了するまでの最大 7 日間の猶予期間は、拡張ランタイム Pod を使用します。拡張ランタイム Pod でメンテナンスの時間枠と除外を使用すると、ノードの自動アップグレードをさらに遅らせることができます。
容量を予約すると、ワークロードは可用性のキューに配置されることなく、リクエストされた TPU を受け取ることができます。

GKE で Cloud TPU を設定する方法については、次の Google Cloud リソースをご覧ください。

GKE で TPU を計画して設定する
GKE Autopilot に TPU ワークロードをデプロイする
GKE Standard に TPU ワークロードをデプロイする
ML タスクで Cloud TPU を使用する際のベストプラクティスを確認する
動画: GKE を使用して Cloud TPU で大規模な ML を構築する
TPU で KubeRay を使用して大規模言語モデルを提供する
GKE Sandbox での GPU ワークロードのサンドボックス化について学習する

GKE Autopilot に TPU ワークロードをデプロイする コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

Autopilot での TPU の仕組み

TPU 構成を計画する

料金

始める前に

TPU とその他の GKE リソースの割り当てを確保する

追加の GKE リソースの割り当て

TPU アプリケーションを準備する

GKE で TPU をプロビジョニングする場合のオプション

カスタムのコンピューティング クラスを使用して TPU を一元的にプロビジョニングする

カスタム ComputeClass を作成する

Ironwood（TPU7x）

その他の TPU バージョン

TPU を消費する Job を作成する

ワークロードで TPU をリクエストする

ワークロード内の TPU リクエストに関する考慮事項

TPU をリクエストするワークロードを作成する

TPU とコレクション スケジューリングをリクエストするワークロードを作成する

マルチホスト TPU スライスを使用する

単一ホストの TPU スライスを使用する

カスタム コンピューティング クラスを使用してコレクションをデプロイする

例: マルチホスト スライスで TPU チップの合計数を表示する

例: 単一ノードで TPU チップを表示する

TPU を観察、モニタリングする

ダッシュボード

ランタイム指標

TPU ノードとノードプールの健全性指標をモニタリングする

ノードのステータス

ノードプールのステータス

ノードプールの可用性

ノードの中断回数

ノードプールの復旧時間（TTR）

ノードプールの中断間隔（TBI）

ホスト指標

ロギング

Autopilot での TPU ワークロードに関する推奨事項

GKE Autopilot に TPU ワークロードをデプロイする

カスタムのコンピューティングクラスを使用して TPU を一元的にプロビジョニングする

TPU とコレクションスケジューリングをリクエストするワークロードを作成する

カスタムコンピューティングクラスを使用してコレクションをデプロイする

例: マルチホストスライスで TPU チップの合計数を表示する