GKE Standard モードで GPU を使用してモデルをトレーニングする

サンプル リポジトリのクローンを作成する

Cloud Shell で、次のコマンドを実行します。

git clone https://github.com/GoogleCloudPlatform/ai-on-gke/ ai-on-gke
cd ai-on-gke/tutorials-and-examples/gpu-examples/training-single-gpu

Standard モードのクラスタと GPU ノードプールを作成する

Cloud Shell で以下の操作を行います。

1. Workload Identity Federation for GKE を使用する Standard クラスタを作成し、Cloud Storage FUSE ドライバをインストールします。

   gcloud container clusters create gke-gpu-cluster \
       --addons GcsFuseCsiDriver \
       --location=us-central1 \
       --num-nodes=1 \
       --workload-pool=PROJECT_ID.s3ns.svc.id.goog
   

   PROJECT_ID は、実際の Cloud de Confianceプロジェクト ID に置き換えます。

   クラスタの作成には数分かかることもあります。

2. GPU ノードプールを作成します。

   gcloud container node-pools create gke-gpu-pool-1 \
       --accelerator=type=nvidia-tesla-t4,count=1,gpu-driver-version=default \
       --machine-type=n1-standard-16 --num-nodes=1 \
       --location=us-central1 \
       --cluster=gke-gpu-cluster
   

Cloud Storage バケットを作成する

1. Cloud de Confiance コンソールで、[バケットの作成] ページに移動します。

   [バケットの作成] に移動

2. [バケットに名前を付ける] フィールドに、次の名前を入力します。

   PROJECT_ID-gke-gpu-bucket
   

3. [続行] をクリックします。

4. [ロケーション タイプ] で、[リージョン] を選択します。

5. [リージョン] リストで us-central1 (Iowa) を選択し、[続行] をクリックします。

6. [データのストレージ クラスを選択する] セクションで、[続行] をクリックします。

7. [オブジェクトへのアクセスを制御する方法を選択する] セクションで、[アクセス制御] には [均一] を選択します。

8. [作成] をクリックします。

9. [公開アクセスの防止] ダイアログで、[このバケットに対する公開アクセス禁止を適用する] チェックボックスをオンにして、[確認] をクリックします。

GKE 用 Workload Identity 連携を使用してバケットにアクセスするようにクラスタを構成する

クラスタから Cloud Storage バケットにアクセスできるようにする手順は次のとおりです。

1. Cloud de Confiance サービス アカウントを作成する。
2. クラスタに Kubernetes ServiceAccount を作成する。
3. Kubernetes ServiceAccount を Cloud de Confiance サービス アカウントにバインドする。

Cloud de Confiance サービス アカウントを作成する

1. Cloud de Confiance コンソールで、[サービス アカウントの作成] ページに移動します。

   [サービス アカウントの作成] に移動

2. [サービス アカウント ID] フィールドに「gke-ai-sa」と入力します。

3. [作成して続行] をクリックします。

4. [ロール] リストで、[Cloud Storage] > [Storage Insights コレクタ サービス] ロールを選択します。

5. [add 別のロールを追加] をクリックします。

6. [ロールを選択] リストで、[Cloud Storage] > [Storage Object Admin] ロールを選択します。

7. [続行] をクリックして [完了] をクリックします。

クラスタに Kubernetes ServiceAccount を作成する

Cloud Shell で、次の操作を行います。

1. Kubernetes Namespace を作成します。

   kubectl create namespace gke-ai-namespace
   

2. Namespace に Kubernetes ServiceAccount を作成します。

   kubectl create serviceaccount gpu-k8s-sa --namespace=gke-ai-namespace
   

Kubernetes ServiceAccount を Cloud de Confiance サービス アカウントにバインドする

Cloud Shell で、次のコマンドを実行します。

1. Cloud de Confiance サービス アカウントに IAM バインディングを追加します。

   gcloud iam service-accounts add-iam-policy-binding gke-ai-sa@PROJECT_ID.s3ns.iam.gserviceaccount.com \
       --role roles/iam.workloadIdentityUser \
       --member "serviceAccount:PROJECT_ID.s3ns.svc.id.goog[gke-ai-namespace/gpu-k8s-sa]"
   

   --member フラグは、 Cloud de Confiance上の Kubernetes ServiceAccount の完全な ID を指定します。

2. Kubernetes ServiceAccount にアノテーションを付けます。

   kubectl annotate serviceaccount gpu-k8s-sa \
       --namespace gke-ai-namespace \
       iam.gke.io/gcp-service-account=gke-ai-sa@PROJECT_ID.s3ns.iam.gserviceaccount.com
   

Pod が Cloud Storage バケットにアクセスできることを確認する

1. Cloud Shell で、次の環境変数を作成します。

   export K8S_SA_NAME=gpu-k8s-sa
   export BUCKET_NAME=PROJECT_ID-gke-gpu-bucket
   

   PROJECT_ID は、実際の Cloud de Confianceプロジェクト ID に置き換えます。

2. TensorFlow コンテナを含む Pod を作成します。

   envsubst < src/gke-config/standard-tensorflow-bash.yaml | kubectl --namespace=gke-ai-namespace apply -f -
   

   このコマンドで、作成した環境変数がマニフェスト内の対応する参照に代入されます。マニフェストをテキスト エディタで開いて、$K8S_SA_NAME と $BUCKET_NAME を対応する値に置き換えることもできます。

3. バケットにサンプル ファイルを作成します。

   touch sample-file
   gcloud storage cp sample-file gs://PROJECT_ID-gke-gpu-bucket
   

4. Pod の準備ができるまで待ちます。

   kubectl wait --for=condition=Ready pod/test-tensorflow-pod -n=gke-ai-namespace --timeout=180s
   

   Pod の準備ができると、出力は次のようになります。

   pod/test-tensorflow-pod condition met
   

5. TensorFlow コンテナでシェルを開きます。

   kubectl -n gke-ai-namespace exec --stdin --tty test-tensorflow-pod --container tensorflow -- /bin/bash
   

6. 作成したサンプル ファイルを閲覧してみます。

   ls /data
   

   出力にサンプル ファイルが表示されます。

7. ログを確認して Pod に接続されている GPU を特定します。

   python3 -m pip install 'tensorflow[and-cuda]'
   python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"
   

   出力には、Pod に接続された GPU が次のように表示されます。

   ...
   PhysicalDevice(name='/physical_device:GPU:0',device_type='GPU')
   

8. コンテナを終了します。

   exit
   

9. サンプル Pod を削除します。

   kubectl delete -f src/gke-config/standard-tensorflow-bash.yaml \
       --namespace=gke-ai-namespace
   

MNIST データセットを使用してトレーニングと予測を行う

このセクションでは、サンプル データセット MNIST でトレーニング ワークロードを実行します。

1. サンプルデータを Cloud Storage バケットにコピーします。

   gcloud storage cp src/tensorflow-mnist-example gs://PROJECT_ID-gke-gpu-bucket/ --recursive
   

2. 次の環境変数を作成します。

   export K8S_SA_NAME=gpu-k8s-sa
   export BUCKET_NAME=PROJECT_ID-gke-gpu-bucket
   

3. トレーニング Job を確認します。

   # Copyright 2023 Google LLC
   #
   # Licensed under the Apache License, Version 2.0 (the "License");
   # you may not use this file except in compliance with the License.
   # You may obtain a copy of the License at
   #
   #      http://www.apache.org/licenses/LICENSE-2.0
   #
   # Unless required by applicable law or agreed to in writing, software
   # distributed under the License is distributed on an "AS IS" BASIS,
   # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   # See the License for the specific language governing permissions and
   # limitations under the License.
   
   apiVersion: batch/v1
   kind: Job
   metadata:
     name: mnist-training-job
   spec:
     template:
       metadata:
         name: mnist
         annotations:
           gke-gcsfuse/volumes: "true"
       spec:
         nodeSelector:
           cloud.google.com/gke-accelerator: nvidia-tesla-t4
         tolerations:
         - key: "nvidia.com/gpu"
           operator: "Exists"
           effect: "NoSchedule"
         containers:
         - name: tensorflow
           image: tensorflow/tensorflow:latest-gpu 
           command: ["/bin/bash", "-c", "--"]
           args: ["cd /data/tensorflow-mnist-example; pip install -r requirements.txt; python tensorflow_mnist_train_distributed.py"]
           resources:
             limits:
               nvidia.com/gpu: 1
               cpu: 1
               memory: 3Gi
           volumeMounts:
           - name: gcs-fuse-csi-vol
             mountPath: /data
             readOnly: false
         serviceAccountName: $K8S_SA_NAME
         volumes:
         - name: gcs-fuse-csi-vol
           csi:
             driver: gcsfuse.csi.storage.gke.io
             readOnly: false
             volumeAttributes:
               bucketName: $BUCKET_NAME
               mountOptions: "implicit-dirs"
         restartPolicy: "Never"

4. トレーニング Job をデプロイします。

   envsubst < src/gke-config/standard-tf-mnist-train.yaml | kubectl -n gke-ai-namespace apply -f -
   

   このコマンドで、作成した環境変数がマニフェスト内の対応する参照に代入されます。マニフェストをテキスト エディタで開いて、$K8S_SA_NAME と $BUCKET_NAME を対応する値に置き換えることもできます。

5. Job のステータスが Completed になるまで待ちます。

   kubectl wait -n gke-ai-namespace --for=condition=Complete job/mnist-training-job --timeout=180s
   

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

   job.batch/mnist-training-job condition met
   

6. TensorFlow コンテナのログを確認します。

   kubectl logs -f jobs/mnist-training-job -c tensorflow -n gke-ai-namespace
   

   出力には、以下のイベントの発生が示されます。

   • 必要な Python パッケージのインストール
   • MNIST データセットのダウンロード
   • GPU を使用したモデルのトレーニング
   • モデルを保存する
   • モデルを評価する

   ...
   Epoch 12/12
   927/938 [============================>.] - ETA: 0s - loss: 0.0188 - accuracy: 0.9954
   Learning rate for epoch 12 is 9.999999747378752e-06
   938/938 [==============================] - 5s 6ms/step - loss: 0.0187 - accuracy: 0.9954 - lr: 1.0000e-05
   157/157 [==============================] - 1s 4ms/step - loss: 0.0424 - accuracy: 0.9861
   Eval loss: 0.04236088693141937, Eval accuracy: 0.9861000180244446
   Training finished. Model saved
   

7. トレーニング ワークロードを削除します。

   kubectl -n gke-ai-namespace delete -f src/gke-config/standard-tf-mnist-train.yaml
   

推論ワークロードをデプロイする

このセクションでは、サンプル データセットを入力として受け取り、予測を返す推論ワークロードをデプロイします。

1. 予測に使用する画像をバケットにコピーします。

   gcloud storage cp data/mnist_predict gs://PROJECT_ID-gke-gpu-bucket/ --recursive
   

2. 推論ワークロードを確認します。

   # Copyright 2023 Google LLC
   #
   # Licensed under the Apache License, Version 2.0 (the "License");
   # you may not use this file except in compliance with the License.
   # You may obtain a copy of the License at
   #
   #      http://www.apache.org/licenses/LICENSE-2.0
   #
   # Unless required by applicable law or agreed to in writing, software
   # distributed under the License is distributed on an "AS IS" BASIS,
   # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   # See the License for the specific language governing permissions and
   # limitations under the License.
   
   apiVersion: batch/v1
   kind: Job
   metadata:
     name: mnist-batch-prediction-job
   spec:
     template:
       metadata:
         name: mnist
         annotations:
           gke-gcsfuse/volumes: "true"
       spec:
         nodeSelector:
           cloud.google.com/gke-accelerator: nvidia-tesla-t4
         tolerations:
         - key: "nvidia.com/gpu"
           operator: "Exists"
           effect: "NoSchedule"
         containers:
         - name: tensorflow
           image: tensorflow/tensorflow:latest-gpu 
           command: ["/bin/bash", "-c", "--"]
           args: ["cd /data/tensorflow-mnist-example; pip install -r requirements.txt; python tensorflow_mnist_batch_predict.py"]
           resources:
             limits:
               nvidia.com/gpu: 1
               cpu: 1
               memory: 3Gi
           volumeMounts:
           - name: gcs-fuse-csi-vol
             mountPath: /data
             readOnly: false
         serviceAccountName: $K8S_SA_NAME
         volumes:
         - name: gcs-fuse-csi-vol
           csi:
             driver: gcsfuse.csi.storage.gke.io
             readOnly: false
             volumeAttributes:
               bucketName: $BUCKET_NAME
               mountOptions: "implicit-dirs"
         restartPolicy: "Never"

3. 推論ワークロードをデプロイします。

   envsubst < src/gke-config/standard-tf-mnist-batch-predict.yaml | kubectl -n gke-ai-namespace apply -f -
   

   このコマンドで、作成した環境変数がマニフェスト内の対応する参照に代入されます。マニフェストをテキスト エディタで開いて、$K8S_SA_NAME と $BUCKET_NAME を対応する値に置き換えることもできます。

4. Job のステータスが Completed になるまで待ちます。

   kubectl wait -n gke-ai-namespace --for=condition=Complete job/mnist-batch-prediction-job --timeout=180s
   

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

   job.batch/mnist-batch-prediction-job condition met
   

5. TensorFlow コンテナのログを確認します。

   kubectl logs -f jobs/mnist-batch-prediction-job -c tensorflow -n gke-ai-namespace
   

   出力は、次のような各画像の予測とモデルの信頼度です。

   Found 10 files belonging to 1 classes.
   1/1 [==============================] - 2s 2s/step
   The image /data/mnist_predict/0.png is the number 0 with a 100.00 percent confidence.
   The image /data/mnist_predict/1.png is the number 1 with a 99.99 percent confidence.
   The image /data/mnist_predict/2.png is the number 2 with a 100.00 percent confidence.
   The image /data/mnist_predict/3.png is the number 3 with a 99.95 percent confidence.
   The image /data/mnist_predict/4.png is the number 4 with a 100.00 percent confidence.
   The image /data/mnist_predict/5.png is the number 5 with a 100.00 percent confidence.
   The image /data/mnist_predict/6.png is the number 6 with a 99.97 percent confidence.
   The image /data/mnist_predict/7.png is the number 7 with a 100.00 percent confidence.
   The image /data/mnist_predict/8.png is the number 8 with a 100.00 percent confidence.
   The image /data/mnist_predict/9.png is the number 9 with a 99.65 percent confidence.