docker-composeでGPUが使えるコンテナを立てる

概要

• docker-composeでコンテナの管理をしたい
• docker-composeでGPUを利用するのにちょっと詰まったので、手順をメモしておきます

先人が丁寧に記事化してくれていたり、公式ドキュメントでも言及されているので、そっちを見てね。

環境

実験した環境です。

• OS: Ubuntu 20.04 LTS
• docker: Docker version 19.03.12, build 48a66213fe
• docker-compose: version 1.26.0, build d4451659

nVidiaドライバがインストールされており、ローカル環境でnvidia-smiなどのコマンドが使えることを前提とします。

手順

• nvidia-container-runtimeをインストール
• /etc/docker/daemon.jsonにruntimeの追加
• dockerの再起動
• docker-compose.ymlを書く

nvidia-container-runtimeのインストール

まずは公式ドキュメントに書かれている通り、NVIDIA-CONTAINER-RUNTIMEをインストールします。
これはドキュメントに書かれている通りです。インストール手順については、公式にこちらで案内があります。

まずはapt-getするためにリポジトリを追加します。

$ curl -s -L https://nvidia.github.io/nvidia-container-runtime/gpgkey | \
  sudo apt-key add -
$ distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
$ curl -s -L https://nvidia.github.io/nvidia-container-runtime/$distribution/nvidia-container-runtime.list | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-runtime.list
$ sudo apt-get update

apt-getする

$ apt-get install nvidia-container-runtime

パスを確認する

$ which nvidia-container-runtime
/usr/bin/nvidia-container-runtime

ここまでで、GPUを利用できるコンテナを立てるだけなら、docker runのオプションで--gpus allなどを指定すれば良いはずです。
ちなみに私の環境では、nvidia-container-runtimeを明示的にインストールしてなかったのですが、なぜか--gpusオプションでコンテナを立てることができてました。なぜかは記憶がありません。
(nvidiaドライバとCUDA toolkitはインストールした記憶があるので、たぶんこれだと思いますが、仕組みが理解できていませんです)

/etc/docker/daemon.jsonにruntimeの追加

docker-composeを使ってコンテナを立てたいのですが、docker-composeには現状で--gpusオプションみたいなものは無いようです。
そこで、runtimeパラメータを指定しますが、nvidiaランタイムを手動で追加します。
（daemon.jsonに自動で追加されるみたいに書いている方もいましたが、僕の環境では手動で書かないとダメでした）

$ cat /etc/docker/daemon.json
{
    "data-root": "/mnt/reservoir/docker",
    "runtimes": {
        "nvidia": {
            "path": "/usr/bin/nvidia-container-runtime",
            "runtimeArgs": []
        }
    }
}

こんな感じです。

dockerの再起動

言わずもがなですが

$ sudo systemctl restart docker

docker-compose.ymlを書いてコンテナを確認

こんな感じでruntimeパラメータと環境変数を指定します。

version: '2.3'

services:
    chainer_test:
        build: 
            context: .
            dockerfile: ./docker/Dockerfile
        image: test_gpu:0.1
        runtime: nvidia
        environment:
            NVIDIA_VISIBLE_DEVICES: all
            NVIDIA_DRIVER_CAPABILITIES: all
        container_name: test_gpu.1
        volumes:
            - $PWD:/workspace
        tty: true

あとは通常通りにbuildとupすればGPUが利用できるコンテナが立ち上がります。

$ docker-compose build
$ docker-compose up

こんな感じ

# nvidia-smi 
Tue Mar  2 13:47:04 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.102.04   Driver Version: 450.102.04   CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce GTX 1070    Off  | 00000000:01:00.0 Off |                  N/A |
|  0%   36C    P8     9W / 230W |    792MiB /  8118MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

(補足)docker-composeのversion 3を使う場合

version3ではruntimeパラメータが利用できないようです。
version3を使いたい場合には、dockerのdefault-runtimeをnvidiaにすれば良いみたいです。

$ cat /etc/docker/daemon.json
{
    "default-runtime": "nvidia",
    ~~以下略~~
}

これで、次のようruntimeを指定しなくてもGPUが利用できるコンテナを立てることができます。

version: '3'

services:
    chainer_test:
        build: 
            context: .
            dockerfile: ./docker/Dockerfile
        image: test_gpu:0.2
        environment:
            NVIDIA_VISIBLE_DEVICES: all
            NVIDIA_DRIVER_CAPABILITIES: all
        container_name: test_gpu.2
        volumes:
            - $PWD:/workspace
        tty: true

(余談)詰まってたところ

公式ドキュメントでは、nvidia-container-runtime-hookのパスを確認しとけと書かれています。
で、上記の通り、僕の環境ではnvidia-container-runtimeを明示的にインストールしてなかったのに--gpusオプションが有効だったのと、nvidia-container-runtime-hookが存在していたので、nvidia-container-runtimeはインストール済みだったのだろうと思い込んでました。
なので、インストールから先の手順を試行錯誤してしまってました。ほんと恥ずかしい。