Google Kubernetes Engine にて GPU を利用可能にする

## はじめに

機械学習用途で Kubernetes を利用していると、どこかのタイミングで GPU を使いたくなってくると思います。

GCP の提供する GKE はマネージドで Kubernetes を利用できるサービスです。GCP で用意してくれるものに乗っかると CPU ベースの処理はもちろんですが、GPU ベースの処理も比較的容易に可能になります。一方で[公式ドキュメント](https://cloud.google.com/kubernetes-engine/docs/how-to/gpus)を参考に進めましたが、私自身の経験と知識不足によりいくつかの躓きポイントもありました。

この記事では GKE 上で GPU を利用できるようになるまでの一連の手続きを紹介します。これから GKE にて GPU を使いたいと思っている方にとって参考になれば幸いです。

なお手元に GKE クラスタが既にあり、既に CPU ベースの運用経験があることを前提に進めます。


## GPU 割当の確認

デフォルトだと GPU の上限は、各プロジェクトで1台です。正確には「各 GCP プロジェクトの」「各リージョンにおける」「各 GPU モデル」の上限が1台です。例えばある GCP プロジェクトにて、asia-northeast1 で利用できる NVIDIA T4 GPUs という GPU モデルは 1台に制限されています。

実際に稼働させるとなると、タスクの性質やチーム状況によっては GPU を複数台利用できることが望ましいかもしれません。割当のページから上限緩和のリクエストを送信できるので、予め分かっている場合は申請するといいかもしれません。

ちなみに私は 2台の追加利用申請を行いましたがリジェクトされてしまいました。何かしらの実績が必要なのかもしれません。自動返信のメールに従い、[こちら](https://cloud.google.com/contact/) から詳細な申請をお送りしました。


ちなみに注意点ですが、asia-northeast1 では NVIDIA Tesla T4 という GPU モデルのみ利用可能です（2020年9月時点）。加えて asia-northeast1 は a〜c の ゾーンを含みますが、ドキュメントによると asia-northeast1-a と asia-northeast1-c のみ利用可能です。

参考 : https://cloud.google.com/compute/docs/gpus?hl=ja#gpus-list


## GPU ノードプールの作成

GKE クラスタはノード（VM=GCE）の論理的集合体である「ノードプール」を複数個持ちます。このあたりの話は割愛しますが、重要なことは GPU ノード専用のノードプールを作成する必要があるということです。

コンソールにせよ、gcloud にせよ、Terraform にせよ GPU ノードプールを作成すると、そのノードプール内のノードには以下の taint が付与されます。

```YAML
key: nvidia.com/gpu
effect: NoSchedule
```

通常 taint が付与されているノードに対して pod をスケジュールするためには toleration を用います。ただし GKE で GPU を利用する場合は toleration ではなく、pod に resource 指定をすることで GPU ノードプールへのスケジューリングを可能にします。

現在我々の環境では Kubeflow Pipelines を利用しており、基本的にスケジュールの単位は Deployment や Job ではなく Pod になるため、このように記述しているという背景もあります。

```YAML
apiVersion: v1
kind: Pod
metadata:
 name: my-gpu-pod
spec:
 containers:
 - name: my-gpu-container
   image: nvidia/cuda:10.0-runtime-ubuntu18.04
   command: ["/bin/bash"]
   resources:
     limits:
      nvidia.com/gpu: 1
```

`nvidia.com/gpu: 1` の記述ですが、ドキュメントには以下のように説明されています。

```YAML
key: nvidia.com/gpu
value: 使用する GPU の数
```

taint と toleration で話が逸れてしまいましたが、Terraform で GPU ノードプールを作成します。以下のように記述します。

```tf
resource "google_container_node_pool" "np-n1-highmem-4-tesla-t4" {
 provider           = google-beta
 name               = "np-n1-highmem-4-tesla-t4"
 cluster            = google_container_cluster.hoge.name
 location           = "asia-northeast1-c"
 initial_node_count = 1
 
 autoscaling {
   min_node_count = 0
   max_node_count = 1
 }

 management {
   auto_repair   = true
   auto_upgrade  = true
 }

 node_locations = [
   "asia-northeast1-c",
 ]

 node_config {
   disk_size_gb = 100
   machine_type = "n1-highmem-4"

   guest_accelerator {
     type  = "nvidia-tesla-t4"
     count = 1
   }

   image_type  = "Ubuntu"
   
   preemptible = true
   
   metadata = {
     disable-legacy-endpoints = true
   } 

   taint {
     key    = "nvidia.com/gpu"
     value  = "present"
     effect = "NO_SCHEDULE"
   }

   shielded_instance_config {
     enable_secure_boot          = false
     enable_integrity_monitoring = true
   }

   workload_metadata_config {
     node_metadata = "GKE_METADATA_SERVER"
   }

   oauth_scopes = [
     "https://www.googleapis.com/auth/cloud-platform"
   ]
 }
}
```

いくつかピックアップします。まず肝心の GPU ですが、guest_accelerator ブロック内にモデルと台数を記述します。そしてそれを動かすインスタンスとして今回は n1-highmem-4 を採用しています。またコストカットを目的にプリエンプティブルインスタンスを採用しています（開発環境であるため）。

```tf
 node_config {
   disk_size_gb = 100
   machine_type = "n1-highmem-4"

   guest_accelerator {
     type  = "nvidia-tesla-t4"
     count = 1
   }

   image_type  = "Ubuntu"
   
   preemptible = true
 }
```

次にオートスケーリングの設定です。最小として0台、最大として1台指定しています。そのため GPU を使用する時は 1台のみ立ち上がり、全く利用しない時は 0台の状態を維持するのでこれもコストカットに寄与します。

```tf
 autoscaling {
   min_node_count = 0
   max_node_count = 1
 }
```

GPU は現在、汎用 N1 マシンタイプでのみサポートされています（2020年9月時点）。そのため N2 や E2 などの他のマシンタイプを使って GPU を利用することはできません。

参考 : https://cloud.google.com/compute/docs/gpus?hl=ja#restrictions

## NVIDIA GPU デバイスドライバのインストール

これで Pod をスケジュールするための GPU ノードは容易できました。

GPU ノードプールを作成したら NVIDIA のデバイスドライバを DaemonSet 経由でインストールします。これによってオートスケーリング機能の良さを殺さずにデバイスドライバのインストールを実現できます。

DaemonSet 自体は Kubernetes の機能で、全てのノードで特定の Pod を実行する仕組みです。そのため、オートスケーリングによって新しく GPU ノードが立ち上がったとしても、そのノードにデバイスドライバが自動的にインストールされることになり、人の手を介すことがなくなります。

ということでクラスタに DaemonSet を apply します。前項で作成したノードプールのOS が Ubuntu なら以下のコマンドを実行します。

```sh

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nvidia-driver-installer/ubuntu/daemonset-preloaded.yaml
```

Ubuntu ではなく Container-Optimized OS（COS）の場合は以下のコマンドになります。COS は GCP 上でのコンテナのホスティングと実行を目的とした、Google が推奨する OS です。Ubuntu などの指定をしなければデフォルトはこの OS になっているはずです。


```sh

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nvidia-driver-installer/cos/daemonset-preloaded.yaml
```

## 挙動確認

ここまで GKE に GPU ノードプールを作成し、デバイスをインストールしました。最後に GPU に Pod をスケジュール出来るかどうかを確認してみます。

以下の Pod を apply します。しばらくして Pod が Completed になったことを確認したらイベントとログを見てみます。

```YAML
apiVersion: v1
kind: Pod
metadata:
name: my-gpu-pod
spec:
 containers:
 - name: my-gpu-container
   image: nvidia/cuda:10.0-runtime-ubuntu18.04
   command: ["nvidia-smi"]
   resources:
     limits:
       nvidia.com/gpu: 1
 restartPolicy: Never
```


今回 Pod を作成するにあたり restartPolicy: Never にしています。自分はドキュメントに従って restartPolicy を設定せずに Pod を apply しましたが、Pod が CrashLoopBackOff になり「何が悪いんだろう…？」とハマってしまいました。

実は Pod はデフォルトで restartPolicy: Always となります（知らなかった）。そのため nvidia-smi で役割を終えたコンテナが無限に再実行されることになります。Kubernetes の挙動としてはある意味正しいのですが、nvidia-smi コマンドで挙動を確認するだけなら restart させる必要はありません。



イベントを見てみると、オートスケーリングが発動して GPU ノードが立ち上がりスケジュールされていることがわかります。はじめに失敗しているのは、スケジュール先の GPU ノードがないためです。何もない時は 0台になるように設定しましたね。

```sh

$ kubectl describe pod my-gpu-pod

Events:
 Type     Reason            Age                    From                                                          Message
 ----     ------            ----                   ----                                                          -------
 Normal   TriggeredScaleUp  5m34s                  cluster-autoscaler                                            pod triggered scale-up:
 Warning  FailedScheduling  3m20s (x3 over 6m11s)  default-scheduler                                             0/1 nodes are available: 1 Insufficient nvidia.com/gpu.
 Warning  FailedScheduling  2m45s (x5 over 3m16s)  default-scheduler                                             0/2 nodes are available: 2 Insufficient nvidia.com/gpu.
 Normal   Scheduled         2m34s                  default-scheduler                                             Successfully assigned default/my-gpu-pod to gke-dev-cluster-np-n1-highmem-4-tesla-11e538ee-8mj2
 Normal   Pulling           2m34s                  kubelet, gke-dev-cluster-np-n1-highmem-4-tesla-11e538ee-8mj2  Pulling image "nvidia/cuda:10.0-runtime-ubuntu18.04"
 Normal   Pulled            2m19s                  kubelet, gke-dev-cluster-np-n1-highmem-4-tesla-11e538ee-8mj2  Successfully pulled image "nvidia/cuda:10.0-runtime-ubuntu18.04"
 Normal   Created           2m12s                  kubelet, gke-dev-cluster-np-n1-highmem-4-tesla-11e538ee-8mj2  Created container my-gpu-container
 Normal   Started           2m12s                  kubelet, gke-dev-cluster-np-n1-highmem-4-tesla-11e538ee-8mj2  Started container my-gpu-container
```

次に Pod のログを見てみましょう。nvidia-smi コマンドを使用するとNVIDIA GPU デバイスドライバが正常に稼働していることを確認できます。以下のようなログが確認できたら GPU ノードプールの作成と、DaemonSet 経由のドライバインストールが無事成功していると考えてOKです。



```sh

$ kubectl logs my-gpu-pod

Tue Sep 22 08:23:26 2020       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 418.152.00   Driver Version: 418.152.00   CUDA Version: 10.1     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   55C    P8    10W /  70W |      0MiB / 15079MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
                                                                              
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+
```

下地は整いました。素敵な GKE x GPU ライフを！