Edited at
DockerDay 17

NVIDIA GPU が使える Docker イメージの作り方

(Docker Advent Calendar に空きがあったので埋めてみた)

Docker コンテナ内から NVIDIA の GPU にアクセスするためには NVIDIA Docker を使えばいい、というのはもはや言うまでもないと思う。

$ docker run --runtime=nvidia --rm nvidia/cuda:9.1-runtime nvidia-smi

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 390.77 Driver Version: 390.77 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
|===============================+======================+======================|
| 0 GeForce GTX 108... Off | 00000000:01:00.0 Off | N/A |
| 25% 43C P0 54W / 250W | 0MiB / 11178MiB | 0% Default |
+-------------------------------+----------------------+----------------------+
| 1 GeForce GTX 108... Off | 00000000:02:00.0 Off | N/A |
| 0% 42C P0 54W / 250W | 0MiB / 11178MiB | 3% Default |
+-------------------------------+----------------------+----------------------+

↑ NVIDIA ドライバ, Docker, NVIDIA Docker (version 2.0) がインストールされたマシンで Docker コンテナ内から nvidia-smi を実行する様子1

しかしこれはどんなイメージでもできるというわけではない。

例えば普通の debian イメージで同じことをやろうとしてもうまくいかない。

$ docker run --runtime=nvidia --rm debian nvidia-smi

docker: Error response from daemon: OCI runtime create failed: container_linux.go:348: starting container process caused "exec: \"nvidia-smi\": executable file not found in $PATH"
: unknown.

コンテナから GPU にアクセスするためには、ランタイムに NVIDIA Docker を使うだけでなく、コンテナ側でいくつか設定をする必要がある。


NVIDIA_VISIBLE_DEVICES

コンテナ内からどの GPU が使えるかを環境変数 NVIDIA_VISIBLE_DEVICES で指定する。

NVIDIA_VISIBLE_DEVICES=0,1 のように番号で指定する。

すべての GPU を利用する場合は NVIDIA_VISIBLE_DEVICES=all で良い。

先ほどの debian イメージの例で言うと、コンテナ起動時に環境変数 NVIDIA_VISIBLE_DEVICES を指定することで動く。

$ docker run --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=0 --rm debian nvidia-smi

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 390.77 Driver Version: 390.77 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
|===============================+======================+======================|
| 0 GeForce GTX 108... Off | 00000000:01:00.0 Off | N/A |
| 0% 43C P0 55W / 250W | 0MiB / 11178MiB | 0% Default |
+-------------------------------+----------------------+----------------------+

この例では NVIDIA_VISIBLE_DEVICES=0 を指定しているので、コンテナ内からはふたつある GPU のうちひとつ目だけが見えるようになっている。


NVIDIA_DRIVER_CAPABILITIES

NVIDIA_VISIBLE_DEVICES を指定しただけでは、nvidia-smi コマンドは使えても OpenGL や CUDA を利用することができない。

これらを利用するためには、用途に応じて環境変数 NVIDIA_VISIBLE_DEVICES を指定する必要がある。


用途

compute
CUDA / OpenCL アプリケーション

compat32
32 ビットアプリケーション

graphics
OpenGL / Vulkan アプリケーション

utility

nvidia-smi コマンドおよび NVML

video
Video Codec SDK

display
X11 ディスプレイに出力

デフォルトでは utility のみが有効になっている。

例えば CUDA アプリケーションを実行する Docker コンテナであれば、NVIDIA_DRIVER_CAPABILITIES=utility,cuda のようにカンマで繋げて指定する。(ただし、CUDA アプリケーションを実行する場合は別途コンテナ内に CUDA をインストールしておく必要がある)

また、 all を指定すればすべての値を指定したのと同じになる。


Anaconda イメージで TensorFlow GPU を実行する例

例えば、Anaconda イメージをベースにして GPU 対応版 TensorFlow を実行できる Docker イメージは、以下の Dockerfile で作れる。


Dockerfile

FROM continuumio/anaconda3

RUN conda install -y tensorflow-gpu

ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES utility,compute


※ Anaconda の場合 conda install cudatoolkit で CUDA をインストールできるが、それをしなくても conda install tensorflow-gpu する際の依存解決で cudatoolkit が入るので、この場合は Dockerfile には書かなくてもよい。


最後に

書いてから気づいたけど、だいたい NVIDIA/nvidia-container-runtime の README に書いてあることを日本語でまとめただけだった。





  1. インストールされている NVIDIA ドライバのバージョンによっては CUDA のバージョンを落とす必要がある: https://blog.sky-net.pw/article/103