GCPで動かすNvidia-docker2

かねてから開発が進められていたNvidia-docker2がインストールできるようになっていたので、使ってみました。
https://github.com/NVIDIA/nvidia-docker

Nvidia-docker1との最大の違いは起動コマンドがnvidia-docker runでなくなったことです。
Nvidia-docker2の起動はdocker runで済みます。

GCEのUbuntu16.04インスタンスにGPU K80でNvidia-docker2をインストールして動かしてみたので、その手順を記載します。

ひとまずインストールしてみる

Nvidia-docker2は標準でdocker-ceサポートのようなので、まずはdocker-ceをインストールします。
コマンドは以下ですが、詳しくは公式ドキュメントをご参照ください。

sudo su -

apt-get -y install apt-transport-https ca-certificates curl software-properties-common
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -
add-apt-repository \
   "deb [arch=amd64] https://download.docker.com/linux/ubuntu \
   $(lsb_release -cs) \
   stable"

apt-get -y update
apt-get -y install docker-ce

これでdocker-ceがインストールされます。つづいてnvidia-docker2をインストールします。
すでにnvidia-docker1がインストールされている場合、apt-get purgeする必要がありますので、ご注意ください。
https://github.com/NVIDIA/nvidia-docker

sudo su -
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey |   apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/ubuntu16.04/amd64/nvidia-docker.list |   tee /etc/apt/sources.list.d/nvidia-docker.list
apt-get -y update

apt-get -y install nvidia-docker2

さっそく使ってみましょう。

systemctl start docker
systemctl enable docker
systemctl status docker

docker run --runtime=nvidia --rm nvidia/cuda nvidia-smi

あら、エラーが出ました。
nvidiaオプションなぞ知らん、とのことです。

CUDAをインストールしてNvidia-docker2をdockerdに登録する

Nvidia-docker2を使うためにはCUDAのインストールと、nvidia-container-runtimeのdockerdへの登録が必要になります。
まずはCUDAのインストールです。
https://developer.nvidia.com/cuda-downloads

ひとまずCUDA8.0を入れておきます。

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/cuda-repo-ubuntu1604_8.0.61-1_amd64.deb
dpkg -i cuda-repo-ubuntu1604_8.0.61-1_amd64.deb
apt-get -y update
apt-get -y install --allow-unauthenticated cuda-8-0

echo 'export CUDA_HOME=/usr/local/cuda' >> /etc/profile
echo 'export PATH=/usr/local/cuda-8.0/bin:${PATH}' >> /etc/profile
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-8.0/lib64:/usr/lib/nvidia-384:/usr/lib/nvidia:${CUDA_HOME}:${LD_LIBRARY_PATH}' >> /etc/profile
source /etc/profile

続いてnvidia-container-runtimeです。
nvidia-container-runtimeはあまり解説がありませんが、nvidia-docker2をdockerで使うためのランタイムだと思います。名前的に。
https://github.com/NVIDIA/nvidia-container-runtime

インストールと設定は以下です。

apt-get -y install nvidia-container-runtime

# Setup a rootfs based on Ubuntu 16.04
cd $(mktemp -d) && mkdir rootfs
curl -sS http://cdimage.ubuntu.com/ubuntu-base/releases/16.04/release/ubuntu-base-16.04-core-amd64.tar.gz | tar --exclude 'dev/*' -C rootfs -xz

# Create an OCI runtime spec
nvidia-container-runtime spec
sed -i 's;"sh";"nvidia-smi";' config.json
sed -i 's;\("TERM=xterm"\);\1, "NVIDIA_VISIBLE_DEVICES=0";' config.json

# Run the container
nvidia-container-runtime run nvidia_smi

nvidia-container-runtime run nvidia_smiを実行すると、nvidia-smiでGPUの稼働情報が表示されます。
最後にdockerdのオーバーライドです。

mkdir -p /etc/systemd/system/docker.service.d
tee /etc/systemd/system/docker.service.d/override.conf <<EOF
[Service]
ExecStart=
ExecStart=/usr/bin/dockerd
EOF

systemctl daemon-reload
systemctl restart docker
systemctl status docker

公式ドキュメントを読むと、上記のヒアドキュメントしている部分は以下の用になっていますが、これだとdocker.serviceが起動しませんでした。

tee /etc/systemd/system/docker.service.d/override.conf <<EOF
[Service]
ExecStart=
ExecStart=/usr/bin/dockerd --host=fd:// --add-runtime=nvidia=/usr/bin/nvidia-container-runtime
EOF

これで準備完了です。
nvidia-docker2をdocker runしてみましょう。

docker run --runtime=nvidia --rm nvidia/cuda nvidia-smi

--runtime=nvidiaを指定し、nvidia/cudaイメージでnvidia-smiします。
実行結果は以下のとおりです。

ためしにKerasでCifar10のCNNを流してみます。
ちゃんとGPUを使ってくれているようです。