GCEでコンテナからGPUを使いたい。 #Docker

Google Compute Engine でGPUを使うコンテナアプリケーションを動かそうとしてハマったので記録しておく。

GCEにはcontainer optimzed os というコンテナを動かすときに使える軽量なイメージがあるらしい。
https://cloud.google.com/container-optimized-os/docs

しかしこのイメージはコンテナを動かすのに最低限必要なものしか入っていないようなので、nvidia-smi とかも入っていない。
COSイメージでGPUを動かすのに必要なGPUドライバをいいかんじにインストールしてくれるツールが存在するので、これを使う。

$ sudo cos-extensions install gpu

とやると、必要なものがインストールされるらしい。¹

cos-optimized imageでGPUをセットしてVMインスタンスを起動する。ディスクは40GBあれば十分動いた。
起動したらsshする

とりあえずGPUがコンテナから見えているか確認する。nvidia/cuda のlatest タグは存在しないのでちゃんとバージョンを指定する必要がある。

$ docker run --rm --gpus all nvidia/cuda:11.0.3-base-ubuntu20.04 nvidia-smi -q
Unable to find image 'nvidia/cuda:11.0.3-base-ubuntu20.04' locally
11.0.3-base-ubuntu20.04: Pulling from nvidia/cuda
d7bfe07ed847: Pull complete 
75eccf561042: Pull complete 
191419884744: Pull complete 
a17a942db7e1: Pull complete 
16156c70987f: Pull complete 
Digest: sha256:57455121f3393b7ed9e5a0bc2b046f57ee7187ea9ec562a7d17bf8c97174040d
Status: Downloaded newer image for nvidia/cuda:11.0.3-base-ubuntu20.04
docker: Error response from daemon: could not select device driver "" with capabilities: [[gpu]].
ERRO[0007] error waiting for container: context canceled

見えていない。GPUドライバをインストールする

$ sudo cos-extensions install gpu

インストールしたら、コンテナからGPUを認識できているか確認する。

$ docker run --rm --gpus all nvidia/cuda:11.0.3-base-ubuntu20.04 nvidia-smi -q
docker: Error response from daemon: could not select device driver "" with capabilities: [[gpu]].
ERRO[0007] error waiting for container: context canceled

結果は同じ。なんでやねん。
調べると同じことにはまってる人がいた。

このissueを追った限りでは、cos-extensions install gpu でいい感じにGPUドライバがインストールされてdocker run --gpus all とかやればGPUを認識できるようになってほしいところが、そうなっていなくて、下記のような感じでいろいろやらないと動くようになっていないみたいである。

#remount the driver folder as executable
sudo mount --bind /var/lib/nvidia /var/lib/nvidia 
sudo mount -o remount,exec /var/lib/nvidia 

docker run \
  --volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64 \
  --volume /var/lib/nvidia/bin:/usr/local/nvidia/bin \
  --device /dev/nvidia0:/dev/nvidia0 \
  --device /dev/nvidia-uvm:/dev/nvidia-uvm \
  --device /dev/nvidiactl:/dev/nvidiactl \
  nvidia/cuda:11.0.3-base-ubuntu20.04 nvidia-smi -q

先人のおかげで、とりあえずコンテナからGPUを認識できるようなった。公式の対応してほしいが、現状対応される気配はなさそうである。
あとは上記設定をスタートアップスクリプトとして設定してしまえば、あとは特に気にすることなく動かすことができそうである。

ソースコードの中身はこのあたりにおいてある。 ↩