Google Kubernetes Engine にて GPU を利用可能にする
機械学習用途で Kubernetes を利用していると、どこかのタイミングで GPU を使いたくなってくるはずです。実際、アトラエでも GKE（Google Kubernetes Engine） の上で GPU を利用しBERT 等の重めの処理を行っています。
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: 1nvidia.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.wevox.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
これで Pod をスケジュールするための GPU ノードは容易できました。
NVIDIA 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.yamlUbuntu ではなく 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 経由のドライバインストールが無事成功していると考えて良いでしょう。
$ 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 |
+-----------------------------------------------------------------------------+実際にどのように利用するか
アトラエではデータサイエンティストが活動する機械学習環境として、Kubeflow を利用しています。Kubeflow Pipelines というツールを使うことで、Jupyter Notebook から GKE 上に Pod をスケジューリングします。
以前解説ブログを書いたので、興味がある方はご覧ください。
GCP のレイヤで説明すると、GCE インスタンス 上のJupyter Notebook の中でコードを書き、そこから各 pod を GKE のノードへとスケジュールします。
具体的にはKubeflow Pipelines SDK という Python の SDK を用いて記述します。
余談
Kubeflow Pipelines SDK はちょっと癖のある SDK です。以前、簡単な使い方を記事にしたのでよかったら参考にしてみて下さい。
https://qiita.com/oguogura/items/32fcaaa7ece2ab868e81
裏では 先程の YAML が叩かれているとは思いますが、この SDK では以下のように記述して GPU の起動と Pod のスケジューリングを行っています。
import kfp
import kfp.dsl as dsl
from kfp.compiler import Compiler
from kfp.components import func_to_container_op
@kfp.dsl.component
def test_tensorflow_gpu() -> int:
import subprocess
from tensorflow.python.client import device_lib
cmd = ["nvidia-smi"]
proc = subprocess.run(cmd ,stdout = subprocess.PIPE, stderr = subprocess.PIPE)
print(proc.stdout.decode("utf8"))
print(device_lib.list_local_devices())
return 0
if __name__=="__main__":
gpu_image = "tensorflow/tensorflow:1.15.2-gpu-py3"
test_tensorflow_gpu_op = kfp.components.func_to_container_op(func = test_tensorflow_gpu, base_image=gpu_image)
client = kfp.Client(host='YOUR HOST')
pipeline_name = 'test-gpu'
@kfp.dsl.pipeline(
name=pipeline_name,
description='test gpu component'
)
def test_gpu():
test_tensorflow_gpu_task = test_tensorflow_gpu_op()
test_tensorflow_gpu_task.set_cpu_request("1")
test_tensorflow_gpu_task.set_memory_request("3G")
test_tensorflow_gpu_task.set_cpu_limit("1")
test_tensorflow_gpu_task.set_memory_limit("3G")
test_tensorflow_gpu_task.set_gpu_limit("1")
zip_filename = './zip_test/{}.zip'.format(pipeline_name)
kfp.compiler.Compiler().compile(estimate_methods, zip_filename)
experiment = client.create_experiment(name=pipeline_name)
run = client.run_pipeline(
experiment_id = experiment.id,
job_name = "test-gpu",
pipeline_package_path = zip_filename
)具体的には以下の部分で GPU のリソース指定をします。実質的に YAML でリソース指定しているのをここで Python として渡しているだけになります。このあたりの書き方はドキュメントにあります。
test_tensorflow_gpu_task = test_tensorflow_gpu_op()
test_tensorflow_gpu_task.set_cpu_request("1")
test_tensorflow_gpu_task.set_memory_request("3G")
test_tensorflow_gpu_task.set_cpu_limit("1")
test_tensorflow_gpu_task.set_memory_limit("3G")
test_tensorflow_gpu_task.set_gpu_limit("1")これによって GPU を用いるジョブが発生した時に GPU が起動し、完了し次第0台にスケールダウンさせることが出来ます。ちなみに GPU が立ち上がりジョブがはじまるまで、5分弱かかります。もう少し早くなるといいですね。
最後に
最後まで読んでくださりありがとうございます。株式会社アトラエでは一緒に夢に向かって突き進んでくれる仲間を募集しています。Twitter からも気軽にどうぞ！
新卒採用の応募はこちら （LINE@に登録いただいております）
