はじめに
機械学習用途で Kubernetes を利用していると、どこかのタイミングで GPU を使いたくなってくると思います。
GCP の提供する GKE はマネージドで Kubernetes を利用できるサービスです。GCP で用意してくれるものに乗っかると CPU ベースの処理はもちろんですが、GPU ベースの処理も比較的容易に可能になります。一方で公式ドキュメントを参考に進めましたが、私自身の経験と知識不足によりいくつかの躓きポイントもありました。
この記事では GKE 上で GPU を利用できるようになるまでの一連の手続きを紹介します。これから GKE にて GPU を使いたいと思っている方にとって参考になれば幸いです。
なお手元に GKE クラスタが既にあり、既に CPU ベースの運用経験があることを前提に進めます。
GPU 割当の確認
デフォルトだと GPU の上限は、各プロジェクトで1台です。正確には「各 GCP プロジェクトの」「各リージョンにおける」「各 GPU モデル」の上限が1台です。例えばある GCP プロジェクトにて、asia-northeast1 で利用できる NVIDIA T4 GPUs という GPU モデルは 1台に制限されています。
実際に稼働させるとなると、タスクの性質やチーム状況によっては GPU を複数台利用できることが望ましいかもしれません。割当のページから上限緩和のリクエストを送信できるので、予め分かっている場合は申請するといいかもしれません。
ちなみに私は 2台の追加利用申請を行いましたがリジェクトされてしまいました。何かしらの実績が必要なのかもしれません。自動返信のメールに従い、こちら から詳細な申請をお送りしました。
ちなみに注意点ですが、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 が付与されます。
key: nvidia.com/gpu
effect: NoSchedule
通常 taint が付与されているノードに対して pod をスケジュールするためには toleration を用います。ただし GKE で GPU を利用する場合は toleration ではなく、pod に resource 指定をすることで GPU ノードプールへのスケジューリングを可能にします。
現在我々の環境では Kubeflow Pipelines を利用しており、基本的にスケジュールの単位は Deployment や Job ではなく Pod になるため、このように記述しているという背景もあります。
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 の記述ですが、ドキュメントには以下のように説明されています。
key: nvidia.com/gpu
value: 使用する GPU の数
taint と toleration で話が逸れてしまいましたが、Terraform で GPU ノードプールを作成します。以下のように記述します。
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 を採用しています。またコストカットを目的にプリエンプティブルインスタンスを採用しています（開発環境であるため）。
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台の状態を維持するのでこれもコストカットに寄与します。
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 なら以下のコマンドを実行します。
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 になっているはずです。
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 になったことを確認したらイベントとログを見てみます。
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台になるように設定しましたね。
$ 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です。
$ 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 ライフを！