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

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

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

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

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 は、次の機能に対応しています。

  1. Spot Pod
  2. 特定の容量予約
  3. 拡張ランタイム Pod
  4. Flex Start

TPU 構成を計画する

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

料金

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

始める前に

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

  • Google Kubernetes Engine API を有効にする。
    • Google Kubernetes Engine API の有効化
  • このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。すでに gcloud CLI をインストールしている場合は、gcloud components update を実行して最新のバージョンを取得する。
  • 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 の割り当ての上限と現在の使用量を確認するには、次の操作を行います。

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

    [割り当て] に移動

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

    1. 次の表を使用して、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 オンデマンド インスタンスの割り当てのプロパティと名前 Spot2 インスタンスの割り当てのプロパティと名前
      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

    2. [項目（ロケーションなど）] プロパティを選択し、「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 ノードの最大数以上の値に設定する必要があります。

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

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

GKE Autopilot では、Kubernetes nodeSelectors を使用して個々のワークロードで TPU を直接使用できます。

また、カスタム コンピューティング クラスを使用して TPU をリクエストすることもできます。カスタム コンピューティング クラスを使用すると、プラットフォーム管理者は、ノード スケーリングの決定時に GKE が優先順位を付けるノードの構成の階層を定義して、ワークロードが選択したハードウェアで実行されるようにできます。

手順については、カスタムのコンピューティング クラスを使用して TPU を一元的にプロビジョニングするをご覧ください。

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

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

  1. JAX、PyTorch、TensorFlow などのフレームワークは、libtpu 共有ライブラリを使用して TPU VM にアクセスします。libtpu には、XLA コンパイラ、TPU ランタイム ソフトウェア、TPU ドライバが含まれています。PyTorch と JAX の各リリースには、特定の libtpu.so バージョンが必要です。GKE で TPU を使用するには、次のバージョンを使用してください。
    TPU タイプ libtpu.so のバージョン
    TPU Trillium（v6e）
    tpu-v6e-slice
    TPU v5e
    tpu-v5-lite-podslice
    TPU v5p
    tpu-v5p-slice
    • 推奨される jax[tpu] バージョン: 0.4.19 以降。
    • 推奨される torchxla[tpuvm] バージョン: 毎晩更新されるバージョンの 2023 年 10 月 23 日付のビルドを使用することをおすすめします。
    TPU v4
    tpu-v4-podslice
    TPU v3
    tpu-v3-slice
    tpu-v3-device
  2. TPU リソースをリクエストするコンテナに、次の環境変数を設定します。
    • TPU_WORKER_ID: 各 Pod の一意の整数。この ID は、TPU スライス内の一意のワーカー ID を示します。このフィールドでサポートされる値の範囲は、0 から Pod 数から 1 を引いた値までです。
    • TPU_WORKER_HOSTNAMES: スライス内で相互に通信する必要がある TPU VM ホスト名または IP アドレスのカンマ区切りのリスト。スライス内の TPU VM ごとにホスト名または IP アドレスが必要です。IP アドレスまたはホスト名のリストは、TPU_WORKER_ID によって順序付けされ、ゼロのインデックスが付けられます。

      • GKE は、completionMode: Indexed、subdomain、parallelism > 1 で Job が作成され、google.com/tpu プロパティをリクエストしたときに、変更用 Webhook を使用してこれらの環境変数を自動的に挿入します。また、Service をバックアップする Pod に DNS レコードが追加されるように、ヘッドレス Service が追加されます。

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

ワークロードで 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 が相互に通信できるようにします。

  1. 次のマニフェストを 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: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          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 のバージョンとトポロジによって異なります。詳しくは、アクセラレータの最小値と最大値をご覧ください。
    • ** 省略可能な runtimeClassname: gvisor*: この Pod を GKE Sandbox で実行できるようにする設定。使用する場合は、この行のコメント化を解除します。GKE Sandbox では、TPU バージョン v4 以降がサポートされています。詳しくは、GKE Sandbox をご覧ください。

  2. Job をデプロイします。

    kubectl create -f tpu-autopilot.yaml

    この Job を作成すると、GKE は自動的に次の処理を行います。

    1. Pod を実行するノードをプロビジョニングします。指定した TPU タイプ、トポロジ、リソース リクエストに応じて、これらのノードは単一ホストスライスまたはマルチホスト スライスのいずれかになります。
    2. Pod に taint を追加してノードに toleration を追加し、他のワークロードが TPU ワークロードと同じノードで実行されないようにします。

  3. このセクションを完了したら、作成したワークロードを削除すれば、それ以上の請求は発生しません。

    kubectl delete -f tpu-autopilot.yaml

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

TPU Trillium では、コレクション スケジューリングを使用して TPU スライスノードをグループ化できます。これらの TPU スライスノードをグループ化すると、ワークロードの需要に合わせてレプリカの数を簡単に調整できます。 Trusted 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 をプロビジョニングし、ワークロードをデプロイするには、次の手順で操作します。

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

    apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: tpu-class
spec:
  priorities:
  - tpu:
      type: tpu-v5-lite-podslice
      count: 4
      topology: 2x4
  - spot: true
    tpu:
      type: tpu-v5-lite-podslice
      count: 4
      topology: 2x4
  - flexStart:
      enabled: true
    tpu:
      type: tpu-v6e-slice
      count: 4
      topology: 2x4
  nodePoolAutoCreation:
    enabled: true

  2. コンピューティング クラスをデプロイします。

    kubectl apply -f tpu-compute-class.yaml

    カスタム コンピューティング クラスと TPU の詳細については、TPU の構成をご覧ください。

  3. 次のマニフェストを 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: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          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 の値は同じにする必要があります。また、選択したカスタムのコンピューティング クラスの tpu.count フィールドの値と同じにする必要があります。
    • MEMORY_SIZE: TPU が使用するメモリの最大量。メモリ上限は、使用する TPU のバージョンとトポロジによって異なります。詳しくは、アクセラレータの最小値と最大値をご覧ください。
    • NUMBER_OF_CHIPS: 使用するコンテナの TPU チップの数。limits と requests の値は同じにする必要があります。

  4. Job をデプロイします。

    kubectl create -f tpu-job.yaml

    この Job を作成すると、GKE は自動的に次の処理を行います。

    • Pod を実行するノードをプロビジョニングします。指定した TPU タイプ、トポロジ、リソース リクエストに応じて、これらのノードは単一ホストスライスまたはマルチホスト スライスのいずれかになります。GKE は最優先の TPU リソースの可用性に応じて優先度の低いリソースにフォールバックし、取得可能性を最大化できます。
    • Pod に taint を追加してノードに toleration を追加し、他のワークロードが TPU ワークロードと同じノードで実行されないようにします。

    詳しくは、カスタムのコンピューティング クラスについてをご覧ください。

  5. このトピックの終了後に作成したリソースを削除すれば、それ以上の請求は発生しません。

    kubectl delete -f tpu-job.yaml

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

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

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

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

  1. 次のマニフェストを 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
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4
  2. マニフェストをデプロイします。
    kubectl create -f available-chips-multihost.yaml

    GKE は、4 つの VM（マルチホスト TPU スライス）を使用して TPU v4 スライスを実行します。スライスには、相互接続された 16 個の TPU チップがあります。

  3. 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
  4. いずれかの Pod のログを取得します。
    kubectl logs POD_NAME

    POD_NAME は、作成した Pod の名前に置き換えます。例: tpu-job-podslice-0-5cd8r

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

    TPU cores: 16
  5. 省略可: ワークロードを削除します。
    kubectl delete -f available-chips-multihost.yaml

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

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

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

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

  1. 次のマニフェストを 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
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8
  2. マニフェストをデプロイします。
    kubectl create -f available-chips-singlehost.yaml

    GKE は、TPU v5e を使用する 8 つの単一ホスト TPU スライスを含むノードをプロビジョニングします。各 TPU ノードには 8 個の TPU チップがあります（単一ホストの TPU スライス）。

  3. Pod のログを取得します。
    kubectl logs tpu-job-jax-v5

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

    Total TPU chips: 8
  4. 省略可: ワークロードを削除します。
      kubectl delete -f available-chips-singlehost.yaml

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

ダッシュボード

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

[GKE TPU Node Pool Status] に移動

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

Trusted 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 クエリを使用します。

```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 クエリを使用します。

```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 を設定する方法については、次の Trusted Cloud by S3NS リソースをご覧ください。