January 15, 2018

KubernetesでGPUを使う

一般的なWebアプリケーションと比較してMachine Leaning（ML）は複雑なインフラを要求する．Data processingを行う環境やModelのTraining/Validationを行う環境，実際にサービスからModelを利用するためのServingの環境といった複数の異なる環境が必要であり，WorkloadによってはCPUだけではなくGPUも必要になる．これらを効率的に扱うためのインフラを構築・運用するのは容易でなくGoogle and Uber’s Best Practices for Deep Learningにあるようにこれまで培われてきたDevOpsの知見を結集していく必要がある．

このような複雑なMLのインフラとしてContainerとKubernetesが利用されることが多くなってきている．特に複数の環境間のPortabilityやTrainingとServingでのScalability，各種ハードウェアやCloud Providerの抽象化においてその利点を多く享受できるからである．実際KubeCon2017感想: Kubernetes in 2018にも書いたようにKubeCon 2017ではML on kubernetesのセッションが多く見られたし，メルカリでもKubernetesを使ったMLモデルのServingを始めている（参考）．

ML on Kubernetesで特に気になるのはGPUをいかに使うかだろう．Kubernetesでは既にv1.6から実験的にNVIDIA GPUへのPodのスケジューリングが可能になっていた．そしてv1.8以降は新たに導入されたDevice PluginとExtended Resourcesという機構によってより容易にGPUが使えるようになっている．

ただしネット上でKubernetes GPUについて検索すると多くの試行錯誤が散らばっておりどの情報が最適な解かわからない状態になっている．本記事では現時点（2018年1月）において最適な方法を簡単に整理する．具体的にはGKEからGPUをつかう最適な方法についてまとめる．なおこの分野は今後も変化がある部分であるので定期的に更新を行う．

KubernetesでGPUを利用するのに必要なこと

KubernetesでNVIDIA GPUを利用するには以下の3つが必要である．

GPUを積んだインスタンスを準備する
NVIDIA GPUのDriverを準備する
NVIDIA Device Pluginを有効にする

以下ではこれらを順番に見ていく．

GPUを積んだインスタンスを準備する

まずはGPUを準備しなければ何も始まらない．現時点（2018年1月）でGKEでGPUを積んだインスタンス（NodePool）を準備するにはAlpha clusterを使う必要がある．Alpha Clusterは全てのKubernetes APIと機能が有効になった実験用のClusterであり30日間しか利用できない．つまり現時点ではGKEでGPUはProduction readyでない．（追記）GKE 1.9.2-gke.1よりAlphaクラスタでなくてもGPUを積んだインスタンス（NodePool）が使えるようになった．またGPUを使えるRegionは限られておりアジアだと現時点（2018年2月）でasia-east1-a（Taiwan）のみが利用可能である（参考）．

GPUを積んだNode Poolを立てるには--acceleratorオプションを利用する．以下の例ではNvidia Tesla k80を利用する．

gcloud alpha container node-pools create “${NODE_POOL}” \
       --cluster="${CLUSTER}" \
       --project="${PROJECT_ID}" \
       --zone="asia-east1-a" \
       --machine-type=n1-standard-2 \
       --accelerator type=nvidia-tesla-k80,count=2

NVIDIA GPUのDriverを準備する

次に各インスタンスにNVIDIA GPUのDriverを準備する必要がある．全てのインスタンスにログインしてインストールを実行するわけにはいかないのでDaemonsetを使う．

GKEではContainer-Optimized OS（COS）がデフォルトのOSとして使われるがそのためのDriverのインストーラーはGoogleCloudPlatform/cos-gpu-installerプロジェクトにある．GoogleCloudPlatform/container-engine-acceleratorsプロジェクトで準備されているDaemonsetを使えば全NodePoolに対してDriverのインストールが行える．GKE 1.8の場合は以下を実行すれば良い．

$ kubectl create -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/k8s-1.8/device-plugin-daemonset.yaml

GKE 1.9の場合は以下を実行する．

$ kubectl create -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/k8s-1.9/daemonset.yaml

なおCOSだけではなく実験的にUbuntu用のDaemonsetインストーラーも準備されている（参考）．

NVIDIA Device Pluginを有効にする

最後にNVIDIA Device Pluginを有効にする．Device PluginはKubernetes 1.8から導入された機能である．Device PluginによりGPUやFPGAsといったVendor specificな初期化や設定が必要なリソースをKubernetesのコアのコードを変更なしに利用できるようにする．

Daemonsetで各NodePoolにインストールするとはまずDevice PluginはExtended Resource（これも1.8で導入された）を利用して新しいDeviceリソースを登録・利用可能にする．例えばNvidia GPUの場合はnvidia.com/gpuというリソース名でそのリソースをPodから利用できるようにする．

次にDevice Pluginは各ノードのDeviceの状態を監視し変更があればKubeletに通知を行う．そしてContainerが作成されるときにはDevice specificなオペレーションを実行しContainerでそのDeviceを利用するためのステップをKubeletに通達する．例えばNvidia GPUの場合はノードにインストールされた必要なDriverのライブラリをContainerにMountする（Device Plugin以前は自分で必要なホストディレクトリをMountする必要があった）．

GKE 1.8の場合は先のDaemonsetを有効にすればDevice Pluginも同時に有効になる．GKE1.9の場合はGPUが積まれた場合にAddonとして自動で有効になる（参考）．

GCP以外の環境であればNVIDIA公式が提供するDevice Plugin．NVIDIA/k8s-device-pluginが利用できる（これにはnvidia-docker 2.0が必要である）．Kubernetes 1.8の場合は以下を実行すれば良い．

$ kubectl create -f https://raw.githubusercontent.com/NVIDIA/k8s-device-plugin/v1.8/nvidia-device-plugin.yml

PodをGPUにスケジューリングする

最後に以下のようなYAMLを書けばPodをGPUにスケジューリングできる．

apiVersion: v1
kind: Pod
metadata:
  name: cuda-vector-add
spec:
  restartPolicy: OnFailure
  containers:
    - name: cuda-vector-add
      # https://github.com/kubernetes/kubernetes/blob/v1.7.11/test/images/nvidia-cuda/Dockerfile
      image: "k8s.gcr.io/cuda-vector-add:v0.1"
      resources:
        limits:
          nvidia.com/gpu: 1 # requesting 1 GPU

まとめ

本記事ではKubernetesからGPUを使う方法の現状を整理した．GKEではAlpha Clusterのみで利用可能でありProduction用途にはまだ向いていないがすぐに利用可能になるだろう．GPUを使う事自体は始まりに過ぎずこの上でどのようなWorkflowを実現していくかがSRE的に/DevOps的に重要になるだろう．