GKE で Cloud TPU アプリケーションを実行する

警告: この手順は、TPU ノード アーキテクチャを使用する TPU にのみ適用されます。TPU ノード アーキテクチャは、TPU v2 と v3 の GKE でのみサポートされています。これらの手順は TPU v2 と v3 にのみ適用されます。GKE での TPU v4、v5e、v5p などの TPU の使用については、TPU ワークロードをデプロイするをご覧ください。

このガイドでは、次の方法について説明します。

• Google Kubernetes Engine で実行する準備として Cloud TPU 構成を設定する
• Cloud TPU をサポートする GKE クラスタを作成する
• TensorBoard を使用して Cloud TPU 指標を可視化し、パフォーマンスを分析する
• Docker 内にモデルをビルドしてコンテナ化する

TPU VM アーキテクチャの詳細については、システム アーキテクチャをご覧ください。このガイドは、TPU ノード アーキテクチャにのみ使用できます。

GKE で Cloud TPU アプリケーションを実行するメリット

Cloud TPU トレーニング アプリケーションは、GKE Pod 内の GKE コンテナで実行するように構成できます。この場合のメリットは次のとおりです。

• ワークフローの設定と管理の改善: GKE が TPU ライフサイクルを管理します。GKE を使用して Cloud TPU の初期化とトレーニングをセットアップすると、GKE によるジョブ障害復旧などのワークロードの繰り返しと管理が可能になります。

• コストの最適化: ジョブがアクティブになっている間だけ、TPU の料金が発生します。GKE は、ジョブのリソース要件に従って、TPU を自動的に作成、削除します。

• 柔軟な使用: Pod 仕様で少し変更するだけで、別のハードウェア アクセラレータ（CPU、GPU、TPU）をリクエストできます。

  kind: Pod
  metadata:
    name: example-tpu
    annotations:
      # The Cloud TPUs that will be created for this Job will support
      # TensorFlow 2.12.1. This version MUST match the
      # TensorFlow version that your model is built on.
      tf-version.cloud-tpus.google.com: "2.12.1"
  spec:
    containers:
    - name: example-container
      resources:
        limits:
          cloud-tpus.google.com/v2: 8
          # See the line above for TPU, or below for CPU / GPU.
          # cpu: 2
          # nvidia.com/gpu: 1
  

• スケーラビリティ: GKE には、数百単位の GKE Pod や TPU ノードにスケーリングできる API（JobおよびDeployment）があります。

• フォールト トレランス: GKE のジョブ API と TensorFlow チェックポイント メカニズムにより、Run-to-Completion セマンティクスが提供されます。VM インスタンスや Cloud TPU ノードで障害が発生すると、チェックポイントから読み込まれた最新の状態を使用して、トレーニング ジョブが自動的に再実行されます。

Cloud TPU と GKE の構成の要件と制限事項

GKE の構成を定義する際は、次の点に注意してください。

• Cloud TPU は、Windows Server ノードプールではサポートされていません。
• GKE クラスタおよびノードプールは、Cloud TPU を使用できるゾーン内に作成する必要があります。トレーニング データとモデルを格納する Cloud Storage バケットも GKE クラスタと同じリージョン内に作成する必要があります。使用可能なゾーンのリストについては、タイプとゾーンに関するドキュメントをご覧ください。
• GKE クラスタには RFC 1918 準拠の IP アドレスを使用する必要があります。詳しくは、GKE ネットワーキングをご覧ください。
• 各コンテナがリクエストできる Cloud TPU は 1 つまでですが、1 つのポッド内の複数のコンテナのそれぞれが Cloud TPU をリクエストできます。

始める前に

1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

3. Make sure that billing is enabled for your Google Cloud project.

4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

5. Make sure that billing is enabled for your Google Cloud project.

GKE で Cloud TPU を使用する際に、プロジェクトは Google Cloud の課金対象となるコンポーネントを使用します。Cloud TPU の料金と GKE の料金を確認して費用を見積もり、使い終わったら手順に沿ってリソースのクリーンアップを行います。

6. Google Cloud コンソールで次の API を有効にします。

   • Cloud TPU API
   • Compute Engine API
   • GKE API

Cloud TPU をサポートする新しいクラスタを作成する

次の手順で環境を設定し、gcloud CLI を使用して Cloud TPU をサポートする GKE クラスタを作成します。

1. gcloud コンポーネントをインストールします。GKE を Cloud TPU で実行するには、このコンポーネントが必要です。

   $ gcloud components install kubectl 

2. gcloud を Google Cloud プロジェクト ID で構成します。

   $ gcloud config set project project-name

   project-name は、Google Cloud プロジェクトの名前で置き換えます。

   このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した Google Cloud API の呼び出しを許可します。

3. Cloud TPU リソースを使用する予定のゾーンで gcloud を構成します。 この例では us-central1-b を使用していますが、サポートされている任意のゾーンで TPU を使用できます。

   $ gcloud config set compute/zone us-central1-b

4. Cloud TPU をサポートするクラスタを GKE 上に作成するには、gcloud container clusters create コマンドを使用します。

   注: 共有 VPC で既存の CIDR 範囲を指定する必要がある場合は、共有 VPC で既存の CIDR 範囲を使用するに示すコマンドを使用します。

   $ gcloud container clusters create cluster-name \
     --release-channel=stable \
     --scopes=cloud-platform \
     --enable-ip-alias \
     --enable-tpu

   コマンドフラグの説明

   リリース チャンネル

   リリース チャンネルは、クラスタの自動アップグレードを管理する手段を提供します。新しいクラスタを作成するときに、そのリリース チャンネルを選択できます。クラスタのアップグレード先は、そのチャネルで提供されるバージョンに限定されます。

   scopes

   これで、クラスタ内のすべてのノードで Cloud Storage バケットにアクセスできます。この設定を機能させるためには、クラスタとストレージ バケットが同じプロジェクト内に存在していなければなりません。デフォルトでは、Kubernetes Pod はデプロイされているノードのスコープを継承します。このため、scopes=cloud-platform によりクラスタ内で実行中のすべての Kubernetes Pod に cloud-platform スコープが指定されます。ポッド単位でアクセスを制限する場合は、GKE のサービス アカウントで認証するためのガイドをご覧ください。

   enable-ip-alias

   クラスタでエイリアス IP 範囲が使用されるように指定します。 これは、GKE 上で Cloud TPU を使用する場合に必要です。

   enable-tpu

   クラスタで Cloud TPU がサポートされるように指定します。

   tpu-ipv4-cidr（任意。前述のチュートリアルでは指定されていません）

   Cloud TPU に使用する CIDR 範囲を指定します。IP_RANGE を IP/20 の形式で指定します（例: 10.100.0.0/20）。このフラグを指定しなければ、/20 サイズの CIDR 範囲が自動的に割り当てられます。

クラスタが作成されると、次のようなメッセージが表示されます。

NAME             LOCATION       MASTER_VERSION    MASTER_IP     MACHINE_TYPE   NODE_VERSION      NUM_NODES  STATUS
cluster-name  us-central1-b  1.16.15-gke.4901  34.71.245.25  n1-standard-1  1.16.15-gke.4901  3          RUNNING

Cloud TPU を Kubernetes ポッド仕様でリクエストする

Kubernetes Pod の仕様:

• コンテナ内のすべてのモデルは、同じ TensorFlow バージョンを使用して構築する必要があります。サポートされているバージョンをご覧ください。

• コンテナ仕様内の resource フィールドの limits セクションで、Cloud TPU リソースを指定します。

  Cloud TPU リソースの単位は Cloud TPU コア数になります。次の表に、有効なリソース リクエストの例を示します。有効な TPU リソースの完全なリストについては、TPU タイプとゾーンをご覧ください。

  使用する予定のリソースが Cloud TPU Pod の場合、Cloud TPU Pod のデフォルトの割り当てはゼロであるため、割り当てをリクエストしてください。

  リソース リクエスト	Cloud TPU タイプ
  cloud-tpus.google.com/v2: 8	Cloud TPU v2 デバイス（8 コア）
  cloud-tpus.google.com/preemptible-v2: 8	プリエンプティブル Cloud TPU v2 デバイス（8 コア）
  cloud-tpus.google.com/v3: 8	Cloud TPU v3 デバイス（8 コア）
  cloud-tpus.google.com/preemptible-v3: 8	プリエンプティブル Cloud TPU v3 デバイス（8 コア）
  cloud-tpus.google.com/v2: 32	v2-32 Cloud TPU Pod（32 コア）
  cloud-tpus.google.com/v3: 32	v3-32 Cloud TPU Pod（32 コア）

  ポッド仕様でのリソースと制限の指定の詳細については、Kubernetes ドキュメントをご覧ください。

以下に示すサンプル Pod 仕様では、TensorFlow 2.12.0 を使用したプリエンプティブル Cloud TPU v2-8 TPU が 1 つリクエストされています。

Cloud TPU ノードの存続期間は、それらのノードをリクエストする Kubernetes Pod にバインドされます。Cloud TPU は、Kubernetes Pod がスケジュールされるとオンデマンドで作成され、Kubernetes Pod が削除されるとリサイクルされます。

apiVersion: v1
kind: Pod
metadata:
  name: gke-tpu-pod
  annotations:
     # The Cloud TPUs that will be created for this Job will support
     # TensorFlow 2.12.1. This version MUST match the
     # TensorFlow version that your model is built on.
     tf-version.cloud-tpus.google.com: "2.12.1"
spec:
  restartPolicy: Never
  containers:
  - name: gke-tpu-container
    # The official TensorFlow 2.12.1 image.
    # https://hub.docker.com/r/tensorflow/tensorflow
    image: tensorflow/tensorflow:2.12.1
    command:
    - python
    - -c
    - |
      import tensorflow as tf
      print("Tensorflow version " + tf.__version__)

      tpu = tf.distribute.cluster_resolver.TPUClusterResolver('$(KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS)')
      print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

      tf.config.experimental_connect_to_cluster(tpu)
      tf.tpu.experimental.initialize_tpu_system(tpu)
      strategy = tf.distribute.TPUStrategy(tpu)

      @tf.function
      def add_fn(x,y):
          z = x + y
          return z

      x = tf.constant(1.)
      y = tf.constant(1.)
      z = strategy.run(add_fn, args=(x,y))
      print(z)
    resources:
      limits:
        # Request a single Preemptible v2-8 Cloud TPU device to train the model.
        cloud-tpus.google.com/preemptible-v2: 8

Job の作成

GKE クラスタ内にジョブを作成して、kubectl をインストールするには、次の手順に従います。

1. テキスト エディタを使用して Pod 仕様 example-job.yaml を作成し、前に示した Pod 仕様にコピーして貼り付けます。

2. ジョブを実行します。

   $ kubectl create -f example-job.yaml

   pod "gke-tpu-pod" created

   このコマンドによって作成されるジョブが、ポッドを自動的にスケジュールします。

3. GKE Pod がスケジュールされて、Cloud TPU ノードがプロビジョニングされたことを確認します。Cloud TPU ノードをリクエストする GKE Pod は、実行までに 5 分ほど保留されることがあります。GKE Pod がスケジュールされるまでは、次のような出力が表示されます。

   $ kubectl get pods -w
   
   NAME          READY     STATUS    RESTARTS   AGE
   gke-tpu-pod   0/1       Pending   0          1m
   

   約 5 分後に、次のような出力が表示されます。

   NAME          READY     STATUS              RESTARTS   AGE
   gke-tpu-pod   0/1       Pending             0          21s
   gke-tpu-pod   0/1       Pending             0          2m18s
   gke-tpu-pod   0/1       Pending             0          2m18s
   gke-tpu-pod   0/1       ContainerCreating   0          2m18s
   gke-tpu-pod   1/1       Running             0          2m48s
   gke-tpu-pod   0/1       Completed           0          3m8s
   

   Ctrl+C キーを押して「kubectl get」コマンドを終了する必要があります。

   次の kubectl コマンドを使用すると、各 GKE Pod に関して、ログ情報を出力してより詳細な情報を取得できます。たとえば、GKE Pod のログ出力を表示するには、次のように使用します。

   $ kubectl logs gke-tpu-pod

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

   2021-09-24 18:55:25.400699: I tensorflow/core/platform/cpu_feature_guard.cc:142]
   This TensorFlow binary is optimized with oneAPI Deep Neural Network Library
   (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
   To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
   2021-09-24 18:55:25.405947: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
   Initialize GrpcChannelCache for job worker -> {0 -> 10.0.16.2:8470}
   2021-09-24 18:55:25.406058: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
   Initialize GrpcChannelCache for job localhost -> {0 -> localhost:32769}
   2021-09-24 18:55:28.091729: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
   Initialize GrpcChannelCache for job worker -> {0 -> 10.0.16.2:8470}
   2021-09-24 18:55:28.091896: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
   Initialize GrpcChannelCache for job localhost -> {0 -> localhost:32769}
   2021-09-24 18:55:28.092579: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:427]
   Started server with target: grpc://localhost:32769
   Tensorflow version 2.12.1
   Running on TPU  ['10.0.16.2:8470']
   PerReplica:{
     0: tf.Tensor(2.0, shape=(), dtype=float32),
     1: tf.Tensor(2.0, shape=(), dtype=float32),
     2: tf.Tensor(2.0, shape=(), dtype=float32),
     3: tf.Tensor(2.0, shape=(), dtype=float32),
     4: tf.Tensor(2.0, shape=(), dtype=float32),
     5: tf.Tensor(2.0, shape=(), dtype=float32),
     6: tf.Tensor(2.0, shape=(), dtype=float32),
     7: tf.Tensor(2.0, shape=(), dtype=float32)
   }
   

   GKE Pod の詳細な説明を表示するには、次のように使用します。

   $ kubectl describe pod gke-tpu-pod

   詳しくは、アプリケーションのイントロスペクションとデバッグをご覧ください。

Docker イメージ内にモデルをビルドしてコンテナ化する

このプロセスの詳細については、独自のモデルのビルドおよびコンテナ化をご覧ください。

既存のクラスタでの Cloud TPU サポートを有効化する

この機能には、サービス固有の利用規約の「一般的なサービス規約」の「pre-GA サービス規約」が適用されます。 pre-GA の機能は「現状有姿」で利用できますが、サポートが制限される場合があります。 詳しくは、リリース ステージの説明をご覧ください。

既存の GKE クラスタで Cloud TPU のサポートを有効にするには、Google Cloud CLI で次の手順を行います。

1. Cloud TPU のサポートを有効にします。

   gcloud beta container clusters update cluster-name --enable-tpu
   

   cluster-name は、使用するクラスタの名前に置き換えます。

   重要: クラスタの更新後に、元の API サーバー エンドポイントが使用できなくなる可能性があります。

2. kubeconfig エントリを更新します。

   gcloud container clusters get-credentials cluster-name
   

カスタム CIDR 範囲の設定

この機能には、サービス固有の利用規約の「一般的なサービス規約」の「pre-GA サービス規約」が適用されます。 pre-GA の機能は「現状有姿」で利用できますが、サポートが制限される場合があります。 詳しくは、リリース ステージの説明をご覧ください。

デフォルトでは、GKE はクラスタによってプロビジョニングされた TPU に /20 のサイズの CIDR ブロックを割り当てます。次のコマンドを実行して、Cloud TPU のカスタム CIDR 範囲を指定できます。

gcloud beta container clusters update cluster-name \
  --enable-tpu \
  --tpu-ipv4-cidr 10.100.0.0/20

以下を置き換えます。

• cluster-name: 既存のクラスタの名前。
• 10.100.0.0/20: カスタム CIDR 範囲。

共有 VPC で既存の CIDR 範囲を使用する

注: このセクションは、共有 VPC ネットワークにクラスタをデプロイし、既存の CIDR 範囲を使用する必要があるユーザーを対象としています。これに該当しない場合は、このセクションをスキップします。

共有 VPC を使用する GKE クラスタの TPU に関するガイドに沿って、共有 VPC の正しい構成を確認します。

• 共有 VPC を使用した Cloud TPU を使用する GKE クラスタ

クラスタ内の Cloud TPU の無効化

重要: Cloud TPU サポートを無効にすると、GKE クラスタによってプロビジョニングされたすべての TPU が削除されます。オペレーションが開始されると、Cloud TPU を必要とするワークロードが使用できなくなります。

既存の GKE クラスタで Cloud TPU のサポートを無効にするには、Google Cloud CLI で次の手順を行います。

1. ワークロードのいずれも Cloud TPU を使用していないことを確認します。

   $ kubectl get tpu

2. クラスタで Cloud TPU サポートを無効にします。

   $ gcloud beta container clusters update cluster-name --no-enable-tpu

   cluster-name は、使用するクラスタの名前に置き換えます。

   このオペレーションには、ゾーンクラスタの場合は 5 分ほどかかり、リージョン クラスタの場合は 15 分ほどかかります（クラスタのリージョンに応じて異なります）。

3. エラーなしでオペレーションが完了したら、クラスタによってプロビジョニングされた TPU が削除されたことを確認できます。

   $ gcloud compute tpus list

   Cloud TPU によって作成される TPU の名前の形式は次のとおりです。

   $ gke-cluster-name-cluster-id-tpu-tpu-id

   以下を置き換えます。

   • cluster-name: 既存のクラスタの名前。
   • cluster-id: 既存のクラスタの ID。
   • tpu-id: Cloud TPU の ID。

   TPU が表示された場合は、次のコマンドを実行して手動で削除できます。

   $ gcloud compute tpus delete gke-cluster-name-cluster-id-tpu-tpu-id

クリーンアップ

GKE 上で Cloud TPU を使い終わったら、リソースをクリーンアップして、Cloud 請求先アカウントに余分に課金されないようにしてください。

1. 次のコマンドを実行して GKE クラスタを削除します。cluster-name をクラスタ名に、project-name を Google Cloud プロジェクト名に置き換えます。

   $ gcloud container clusters delete cluster-name \
   --project=project-name --zone=us-central1-b

2. データの確認が終了したら、gcloud CLI コマンドを使用して作成した Cloud Storage バケットを削除します。bucket-name を該当する Cloud Storage バケットの名前に置き換えます。

   $ gcloud storage rm gs://bucket-name --recursive