【Docker】GPU が使える Jupyter Notebook 環境を最速で用意する

#目的
Docker コンテナで Jupyter Notebook を立ち上げて、GPU が使えるようにします。
NVIDIA 公式の CUDA がインストールされた Docker イメージを元にして、コンテナを作成します。
環境を準備すれば、あとは下記の手順に沿ってコピペでいけると思っています。

#事前準備

Ubuntu に NVIDIA ドライバ、nvidia-container-toolkit、Docker がインストールされていることが前提となります。
以下のサイトを参考にインストールしました。

NVIDIA Docker って今どうなってるの？ (19.11版)
https://qiita.com/ksasaki/items/b20a785e1a0f610efa08

#環境

ホスト OS は Ubuntu 18.04.5 LTS を使用しました。

```
$ cat /etc/os-release
NAME="Ubuntu"
VERSION="18.04.5 LTS (Bionic Beaver)"
ID=ubuntu
ID_LIKE=debian
PRETTY_NAME="Ubuntu 18.04.5 LTS"
VERSION_ID="18.04"
HOME_URL="https://www.ubuntu.com/"
SUPPORT_URL="https://help.ubuntu.com/"
BUG_REPORT_URL="https://bugs.launchpad.net/ubuntu/"
PRIVACY_POLICY_URL="https://www.ubuntu.com/legal/terms-and-policies/privacy-policy"
VERSION_CODENAME=bionic
UBUNTU_CODENAME=bionic
```

NVIDIA ドライバの確認。

```
$ dpkg -l cuda-drivers | grep cuda
ii  cuda-drivers   460.32.03-1  amd64        CUDA Driver meta-package, branch-agnostic
```

nvidia-container-toolkit の確認。

```
$ dpkg -l nvidia-container-toolkit | grep nvidia
ii  nvidia-container-toolkit 1.4.2-1      amd64        NVIDIA container runtime hook
```

Docker の確認。

```
$ docker --version
Docker version 20.10.5, build 55c4c88
```

GPU の確認。GeForce GTX 1070 が見えています。

```
$ nvidia-smi
Sun Mar 28 09:37:10 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| 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    On   | 00000000:01:00.0  On |                  N/A |
|  0%   31C    P8    15W / 200W |    308MiB /  8116MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1506      G   /usr/lib/xorg/Xorg                136MiB |
|    0   N/A  N/A      1717      G   /usr/bin/gnome-shell               65MiB |
|    0   N/A  N/A      6503      G   ...AAAAAAAA== --shared-files       59MiB |
|    0   N/A  N/A      6963      G   ...AAAAAAAAA= --shared-files       42MiB |
+-----------------------------------------------------------------------------+
```

#今回使用・作成する Docker イメージとコンテナ

- nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04 　→ NVIDIA 公式の Docker イメージ
- my-nvidia-cuda 　→ NVIDIA 公式の Docker イメージを元に、今回新規作成するイメージ（pip パッケージの追加などを行う）
- my-jupyter 　→ 今回作成するイメージから起動するコンテナ

#ディレクトリ構成とファイルの内容

ディレクトリ構成。
NVIDIA 公式イメージを元に今回の my-nvidia-cuda イメージを作成するため、Dockerfile などを配置します。
Docker ファイル内で、requirements.txt に記述したパッケージを pip でインストールが行われます。
src ディレクトリはコンテナと共有する用です。
test.ipynb は動作確認用に後で作る想定なので、無くて構いません。

```
$ tree my-nvidia-cuda/
my-nvidia-cuda/
├── Dockerfile
├── requirements.txt
└── src
    └── test.ipynb
```

Dockerfile の内容。

```
$ cat my-nvidia-cuda/Dockerfile
FROM nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04

USER root

COPY ./requirements.txt /tmp
WORKDIR /code

RUN apt-get update && apt-get -y upgrade
RUN apt install -y curl python3 python3-distutils
RUN curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py

RUN pip install -r /tmp/requirements.txt
```

requirements.txt の内容。Dockerfile の pip install で読み出されます。
jupyter ほか、必要なパッケージを記述します。

```
$ cat my-nvidia-cuda/requirements.txt 
jupyter
jupyterlab
numpy
pandas
matplotlib
scikit-learn
scikit-image
scipy
torch
torchvision
tqdm
albumentations
tensorflow-gpu
Pillow
opencv-python
```

「my-nvidia-cuda」の名前で Docker イメージをビルドします。
しばらく時間がかかります。

```
$ cd my-nvidia-cuda
$ docker build . -t my-nvidia-cuda
```

# Docker コンテナの起動

先程作成した Docker イメージからコンテナを起動します。
「--gpus all」オプションをつけて GPU を使えるようにします。
「sh -c 'jupyter-lab ....」オプションで jupyter-lab が立ち上がるようにします。

```
$ docker run -it --rm --gpus all -v `pwd`/src:/code -p 8888:8888 --name my-jupyter my-nvidia-cuda sh -c 'jupyter-lab --allow-root --ip=*'
/usr/local/lib/python3.8/dist-packages/jupyter_server/transutils.py:13: FutureWarning: The alias `_()` will be deprecated. Use `_i18n()` instead.
  warnings.warn(warn_msg, FutureWarning)
[I 2021-03-28 01:02:20.035 ServerApp] jupyterlab | extension was successfully linked.
[I 2021-03-28 01:02:20.047 ServerApp] Writing notebook server cookie secret to /root/.local/share/jupyter/runtime/jupyter_cookie_secret
[I 2021-03-28 01:02:20.220 ServerApp] nbclassic | extension was successfully linked.
[W 2021-03-28 01:02:20.247 ServerApp] WARNING: The Jupyter server is listening on all IP addresses and not using encryption. This is not recommended.
[I 2021-03-28 01:02:20.250 LabApp] JupyterLab extension loaded from /usr/local/lib/python3.8/dist-packages/jupyterlab
[I 2021-03-28 01:02:20.250 LabApp] JupyterLab application directory is /usr/local/share/jupyter/lab
[I 2021-03-28 01:02:20.253 ServerApp] jupyterlab | extension was successfully loaded.
[I 2021-03-28 01:02:20.257 ServerApp] nbclassic | extension was successfully loaded.
[I 2021-03-28 01:02:20.257 ServerApp] Serving notebooks from local directory: /code
[I 2021-03-28 01:02:20.257 ServerApp] Jupyter Server 1.5.1 is running at:
[I 2021-03-28 01:02:20.257 ServerApp] http://ea30cc7e2196:8888/lab?token=e982672b6ea939de710fdc5742141e6025356026af5a329f
[I 2021-03-28 01:02:20.257 ServerApp]  or http://127.0.0.1:8888/lab?token=e982672b6ea939de710fdc5742141e6025356026af5a329f
[I 2021-03-28 01:02:20.258 ServerApp] Use Control-C to stop this server and shut down all kernels (twice to skip confirmation).
[W 2021-03-28 01:02:20.261 ServerApp] No web browser found: could not locate runnable browser.
[C 2021-03-28 01:02:20.261 ServerApp] 
    
    To access the server, open this file in a browser:
        file:///root/.local/share/jupyter/runtime/jpserver-6-open.html
    Or copy and paste one of these URLs:
        http://ea30cc7e2196:8888/lab?token=e982672b6ea939de710fdc5742141e6025356026af5a329f
     or http://127.0.0.1:8888/lab?token=e982672b6ea939de710fdc5742141e6025356026af5a329f
```

上記コマンドで表示された http://127.0.0.1:8888/lab?token=..... のアドレスをブラウザで開き、Jupyter Lab にアクセスします。

# Jupyter Notebook で GPU が使えることを確認する

「!nvidia-smi」で GPU の情報が表示されれば OK です。
PyTorch から CUDA が利用可能なことも確認できています。

![jupyter_notebook_gpu.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/591407/1d3755bf-65b3-8fce-e7ad-69789391ac0d.png)