GKE で Cloud 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 をリクエストできます。
始める前に
- 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.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Make sure that billing is enabled for your Google Cloud project.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Make sure that billing is enabled for your Google Cloud project.
-
Google Cloud コンソールで次の API を有効にします。
GKE で Cloud TPU を使用する際に、プロジェクトは Google Cloud の課金対象となるコンポーネントを使用します。Cloud TPU の料金と GKE の料金を確認して費用を見積もり、使い終わったら手順に沿ってリソースのクリーンアップを行います。
Cloud TPU をサポートする新しいクラスタを作成する
gcloud CLI を使用して環境を設定し、Cloud TPU をサポートする GKE クラスタを作成する手順は以下のとおりです。
gcloud
コンポーネントをインストールします。GKE を Cloud TPU で実行するには、このコンポーネントが必要です。$ gcloud components install kubectl
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 の呼び出しを許可します。Cloud TPU リソースを使用する予定のゾーンで
gcloud
を構成します。 この例ではus-central1-b
を使用していますが、サポートされている任意のゾーンで TPU を使用できます。$ gcloud config set compute/zone us-central1-b
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 をインストールするには、次の手順に従います。
テキスト エディタを使用して、Pod 仕様
example-job.yaml
を作成してコピーし、上記の Pod 仕様に貼り付けます。ジョブを実行します。
$ kubectl create -f example-job.yaml
pod "gke-tpu-pod" created
このコマンドによって作成されるジョブが、ポッドを自動的にスケジュールします。
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 で次の手順を行います。
Cloud TPU のサポートを有効にします。
gcloud beta container clusters update cluster-name --enable-tpu
cluster-name は、使用するクラスタの名前に置き換えます。
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 で次の手順を行います。
ワークロードのいずれも Cloud TPU を使用していないことを確認します。
$ kubectl get tpu
クラスタで Cloud TPU サポートを無効にします。
$ gcloud beta container clusters update cluster-name --no-enable-tpu
cluster-name は、使用するクラスタの名前に置き換えます。
このオペレーションには、ゾーンクラスタの場合は 5 分ほどかかり、リージョン クラスタの場合は 15 分ほどかかります(クラスタのリージョンに応じて異なります)。
エラーなしでオペレーションが完了したら、クラスタによってプロビジョニングされた 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 請求先アカウントに余分に課金されないようにしてください。
次のコマンドを実行して GKE クラスタを削除します。
cluster-name
をクラスタ名に、project-name
を Google Cloud プロジェクト名に置き換えます。$ gcloud container clusters delete cluster-name \ --project=project-name --zone=us-central1-b
データの確認が終了したら、
gsutil
コマンドを使用して、作成した Cloud Storage バケットを削除します。bucket-name
を該当する Cloud Storage バケットの名前に置き換えます。$ gsutil rm -r gs://bucket-name