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

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

[追記]
Docker 19.03 以降はコンテナの起動方法が違うので以下を参照。
NVIDIA Docker って今どうなってるの？ (19.11版) - Qiita
[追記おわり]

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 を実行する様子¹

しかしこれはどんなイメージでもできるというわけではない。
例えば普通の 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_DRIVER_CAPABILITIES を指定する必要がある。

値	用途
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 ↩