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

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

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. Google Cloud アカウントにログインします。Google Cloud を初めて使用する場合は、アカウントを作成して、実際のシナリオでの Google プロダクトのパフォーマンスを評価してください。新規のお客様には、ワークロードの実行、テスト、デプロイができる無料クレジット $300 分を差し上げます。

  2. Google Cloud Console の [プロジェクト セレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

    プロジェクト セレクタに移動

  3. Google Cloud プロジェクトで課金が有効になっていることを確認します。

  4. Google Cloud Console の [プロジェクト セレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

    プロジェクト セレクタに移動

  5. Google Cloud プロジェクトで課金が有効になっていることを確認します。

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

  6. Google Cloud コンソールで次の 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 コマンドを使用します。

    $ 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 サポートを有効化する

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

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

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

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

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

    gcloud container clusters get-credentials cluster-name

カスタム CIDR 範囲の設定

デフォルトでは、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 を使用して GKE クラスタの TPU のガイドに沿って、共有 VPC の正しい構成を確認します。

クラスタ内の 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. データの確認が終了したら、gsutil コマンドを使用して、作成した Cloud Storage バケットを削除します。bucket-name を該当する Cloud Storage バケットの名前に置き換えます。

    $ gsutil rm -r gs://bucket-name