CUDA入りDockerfileを作る

## 概要

nvidiaが提供しているDockerfileを生成するツールを使って、CUDAのDockerfileを生成する方法。

> nvidia/cuda の Dockerfile を生成するツール
> https://gitlab.com/nvidia/container-images/cuda

こんなときに便利

- Dockerで動くGPUを使ったディープラーニングの環境を作りたい時。
- ホスト端末の環境を汚したくない時。
- 複数のCUDAを使いたい時。

## 環境

- OS: Ubuntu20.04

#### 必要な設定

- Dockerが動作すること。

#### 上記に加えて生成したDockerfileを使ってDockerイメージ動作させるために必要な設定

- NVIDIAのドライバが正しく動作すること。
- `nvidia-docker2`がインストールされていること。

ubuntu20.04に`nvidia-docker2`をインストールするには以下のコマンド。

```bash
distribution=$(. /etc/os-release;echo $ID$VERSION_ID) && curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - && curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
sudo apt-get install -y nvidia-docker2
sudo systemctl restart docker
```

## Dockerfileの生成方法

### リポジトリをクローン

```bash
git clone https://gitlab.com/nvidia/container-images/cuda.git
```

### ビルド済みのDockerfileを削除

クローンしてきたリポジトリの中に、既にビルド済みの`Dockerfile`が存在する`dist`ディレクトリがあるので、これを削除する。

```bash
rm -rf ./cuda/dist/*
```

### 仮想環境を作成

あまり環境を汚したくないので、ビルド用の仮想環境を作成する。今回はconda（miniconda）を使って作成している。

```bash
conda create -n build_cuda python=3.10.6
conda activate build_cuda
```

### poetryをインストール

`poetry`というパッケージマネージャが使われていたので、これを使って依存関係をインストールする。

そのために、まずは`poetry`をインストールする。

```bash
curl -sSL https://install.python-poetry.org | python3 -
```

### 作業ディレクトリをクローンしてきたリポジトリに移動

```bash
cd cuda
```

### poetryを使って依存関係をインストール

自分の環境では、`~/.local/bin/poetry`にインストールされたので以下を実行して依存関係をインストールする。

```bash
~/.local/bin/poetry install
```

### Dockerfileを生成

試しにubuntu20.04でCUDA11.2.2のDockerfileを生成してみる。

```bash
python manager.py --manifest="manifests/cuda.yaml" generate --cuda-version 11.2.2 --os-name ubuntu --os-version 20.04
python manager.py generate --help
```

実行すると、`リポジトリのディレクトリ/dist/CUDAのバージョン/OS名`ディレクトリに`base`、`devel`、`runtime`のディレクトリが生成され`Dockerfile`が生成される。

##### 生成されたDockerfileの種類について

> https://hub.docker.com/r/nvidia/cuda

base: CUDAランタイム(cudart)を含む
runtime: baseの上にCUDA数学ライブラリとNCCLを含む。cuDNNも含むランタイムイメージもある。
devel: runtimeの上にヘッダー、CUDAイメージのビルドに必要な開発ツールを含む。これらのイメージは、マルチステージビルドに特に便利。


## カスタマイズしたDockerfileを作成してみる。

jupyter/minimal-notebookにcudaを足してみる。

#### manager.pyを修正

`manager.py`の1041行目（バージョンによって変動するかもしれない）あたりにある、リストにファイル名のパターンを追加する。

例えば、`custom-dockerfile.j2`というファイルを追加したい場合は、以下のように`custom`と追加する。

```python
....
    def generate_containerscripts(self):
        for img in ["base", "devel", "runtime", "custom"]:
            self.cuda["target"] = img

            globber = f"*"
            ....
```

#### ファイルを作成

`templates/ディストリビューション名`ディレクトリに`リストで追加した名前-dockerfile.j2`というファイルを作成する。

今回の場合、`custom-dockerfile.j2`というファイルを追加した。

```Dockerfile
FROM jupyter/minimal-notebook:python-3.8.8
USER root
WORKDIR /tmp

ENV OS ubuntu2004
{% if "x86_64" in cuda %}
ENV NVARCH x86_64
{% if "requires" in cuda.x86_64 %}
ENV NVIDIA_REQUIRE_CUDA "{{ cuda.x86_64.requires }}"
{% endif %}

ENV NV_CUDA_CUDART_VERSION {{ cuda.x86_64.components.cudart.version }}
ENV NV_CUDA_COMPAT_PACKAGE cuda-compat-{{ cuda.version.major }}-{{ cuda.version.minor }}
{% endif %}
ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility

ENV CUDA_VERSION {{ cuda.version.release_label }}

ENV NV_CUDA_LIB_VERSION "{{ cuda.version.release_label + "-1" }}"
{% if "x86_64" in cuda %}

ENV NV_CUDA_CUDART_DEV_VERSION {{ cuda.x86_64.components.cudart_dev.version }}
ENV NV_NVML_DEV_VERSION {{ cuda.x86_64.components.nvml_dev.version }}
ENV NV_LIBCUSPARSE_DEV_VERSION {{ cuda.x86_64.components.libcusparse_dev.version }}
    {% if "libnpp_dev" in cuda.x86_64.components %}
        {% set has_libnpp_dev_package = true %}
ENV NV_LIBNPP_DEV_VERSION {{ cuda.x86_64.components.libnpp_dev.version }}
ENV NV_LIBNPP_DEV_PACKAGE libnpp-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBNPP_DEV_VERSION}
    {% endif %}

ENV NV_LIBCUBLAS_DEV_VERSION {{ cuda.x86_64.components.libcublas_dev.version }}
ENV NV_LIBCUBLAS_DEV_PACKAGE_NAME libcublas-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}
ENV NV_LIBCUBLAS_DEV_PACKAGE ${NV_LIBCUBLAS_DEV_PACKAGE_NAME}=${NV_LIBCUBLAS_DEV_VERSION}
    {% if "nvprof" in cuda.x86_64.components %}
        {% set has_nvprof_package = true %}

ENV NV_NVPROF_VERSION {{ cuda.x86_64.components.nvprof.version }}
ENV NV_NVPROF_DEV_PACKAGE cuda-nvprof-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_NVPROF_VERSION}
    {% endif %}
    {% if "libnccl2" in cuda.x86_64.components and cuda.x86_64.components.libnccl2 and "libnccl2_dev" in cuda.x86_64.components and cuda.x86_64.components.libnccl2_dev %}
        {% set has_libnccl_dev_package = true %}

ENV NV_LIBNCCL_DEV_PACKAGE_NAME libnccl-dev
ENV NV_LIBNCCL_DEV_PACKAGE_VERSION {{ cuda.x86_64.components.libnccl2_dev.version }}
ENV NCCL_VERSION {{ cuda.x86_64.components.libnccl2_dev.version }}
ENV NV_LIBNCCL_DEV_PACKAGE ${NV_LIBNCCL_DEV_PACKAGE_NAME}=${NV_LIBNCCL_DEV_PACKAGE_VERSION}+cuda{{ cuda.version.major }}.{{ cuda.version.minor }}
        {% if "libnccl2" in cuda.x86_64.components and "source" in cuda.x86_64.components.libnccl2 %}
ENV NV_LIBNCCL_PACKAGE_SHA256SUM {{ cuda.x86_64.components.libnccl2.sha256sum }}
ENV NV_LIBNCCL_PACKAGE_SOURCE {{ cuda.x86_64.components.libnccl2.source }}
ENV NV_LIBNCCL_PACKAGE_SOURCE_NAME {{ cuda.x86_64.components.libnccl2.basename }}

RUN apt-get update && apt-get install -y --no-install-recommends wget

RUN wget -q ${NV_LIBNCCL_PACKAGE_SOURCE} \
    && echo "$NV_LIBNCCL_PACKAGE_SHA256SUM  ${NV_LIBNCCL_PACKAGE_SOURCE_NAME}" | sha256sum -c --strict - \
    && dpkg -i ${NV_LIBNCCL_PACKAGE_SOURCE_NAME} \
    && rm -f ${NV_LIBNCCL_PACKAGE_SOURCE_NAME} \
    && apt-get purge --autoremove -y wget \
    && rm -rf /var/lib/apt/lists/*

        {% endif %}
        {% if "libnccl2_dev" in cuda.x86_64.components and "source" in cuda.x86_64.components.libnccl2_dev %}
ENV NV_LIBNCCL_DEV_PACKAGE_SHA256SUM {{ cuda.x86_64.components.libnccl2_dev.sha256sum }}
ENV NV_LIBNCCL_DEV_PACKAGE_SOURCE {{ cuda.x86_64.components.libnccl2_dev.source }}
ENV NV_LIBNCCL_DEV_PACKAGE_SOURCE_NAME {{ cuda.x86_64.components.libnccl2_dev.basename }}
RUN apt-get update && apt-get install -y --no-install-recommends wget

RUN wget -q ${NV_LIBNCCL_DEV_PACKAGE_SOURCE} \
    && echo "$NV_LIBNCCL_DEV_PACKAGE_SHA256SUM  ${NV_LIBNCCL_DEV_PACKAGE_SOURCE_NAME}" | sha256sum -c --strict - \
    && dpkg -i ${NV_LIBNCCL_DEV_PACKAGE_SOURCE_NAME} \
    && rm -f ${NV_LIBNCCL_DEV_PACKAGE_SOURCE_NAME} \
    && apt-get purge --autoremove -y wget \
    && rm -rf /var/lib/apt/lists/*
        {% endif %}
    {% endif %}
{% endif %}
ENV CUDNN_VERSION {{ cuda.x86_64.components.cudnn8.version }}

RUN sed -i 's@archive.ubuntu.com@ftp.jaist.ac.jp/pub/Linux@g' /etc/apt/sources.list && \
    apt-get update && apt-get install -y --no-install-recommends \
    wget software-properties-common tzdata dirmngr gpg-agent \
{% if cuda.os.version == "22.04" and not (cuda.image_tag_suffix | length) %}
    gnupg2 curl ca-certificates && \
    curl -fsSLO {{ cuda.repo_url }}/${NVARCH}/cuda-keyring_1.0-1_all.deb && \
    dpkg -i cuda-keyring_1.0-1_all.deb && \
{% elif cuda.os.version != "16.04" or cuda.flavor == "jetson" %}
    {# Still used for internal 22.04 images until https://jirasw.nvidia.com/browse/CUDAINST-1240 #}
    gnupg2 curl ca-certificates && \
    curl -fsSL {{ cuda.repo_url }}/${NVARCH}/3bf863cc.pub | apt-key add - && \
    echo "deb {{ cuda.repo_url }}/${NVARCH} /" > /etc/apt/sources.list.d/cuda.list && \
{% else %}
    ca-certificates apt-transport-https gnupg-curl && \
    NVIDIA_GPGKEY_SUM=a21c1a0b18a4196fa901b833e13c4fa64f094d7d9e8a6495318e7255f0ef23d1 && \
    NVIDIA_GPGKEY_FPR=eb693b3035cd5710e231e123a4b469963bf863cc && \
    apt-key adv --fetch-keys {{ cuda.repo_url }}/${NVARCH}/3bf863cc.pub && \
    apt-key adv --export --no-emit-version -a $NVIDIA_GPGKEY_FPR | tail -n +5 > cudasign.pub && \
    echo "$NVIDIA_GPGKEY_SUM  cudasign.pub" | sha256sum -c --strict - && rm cudasign.pub && \
    echo "deb {{ cuda.repo_url }}/${NVARCH} /" > /etc/apt/sources.list.d/cuda.list && \
{% endif %}
    wget {{ cuda.repo_url }}/${NVARCH}/cuda-$OS.pin && \
    mv cuda-$OS.pin /etc/apt/preferences.d/cuda-repository-pin-600 && \
    apt-key adv --fetch-keys {{ cuda.repo_url }}/${NVARCH}/7fa2af80.pub && \
    add-apt-repository "deb {{ cuda.repo_url }}/${NVARCH}/ /" && \
    apt-get update && \
    apt-get install -y --no-install-recommends \
    cuda-cudart-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_CUDART_DEV_VERSION} \
    ${NV_CUDA_COMPAT_PACKAGE} \
    cuda-command-line-tools-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_LIB_VERSION} \
    cuda-minimal-build-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_LIB_VERSION} \
    cuda-libraries-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_LIB_VERSION} \
    cuda-nvml-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_NVML_DEV_VERSION} \
{% if has_nvprof_package %}
    ${NV_NVPROF_DEV_PACKAGE} \
{% endif %}
{% if has_libnpp_dev_package %}
    ${NV_LIBNPP_DEV_PACKAGE} \
{% endif %}
    libcusparse-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBCUSPARSE_DEV_VERSION} \
    ${NV_LIBCUBLAS_DEV_PACKAGE} \
{% if has_libnccl_dev_package %}
    libnccl2=${NCCL_VERSION}+cuda{{ cuda.version.major }}.{{ cuda.version.minor }} \
    ${NV_LIBNCCL_DEV_PACKAGE} \
{% endif %}
{% if has_cuda_profiler_api_package %}
    ${NV_CUDA_PROFILER_API_PACKAGE} \
{% endif %}
    libcudnn8=${CUDNN_VERSION}+cuda{{ cuda.version.major }}.{{ cuda.version.minor }} && \
    apt-mark hold libcudnn8 \
{% if ( cuda.version.major | int ) == 11 and ( cuda.version.minor | int ) <= 2 %}
    && ln -s cuda-{{ cuda.version.major }}.{{ cuda.version.minor }} /usr/local/cuda \
{% endif %}
    && echo "/usr/local/nvidia/lib" >> /etc/ld.so.conf.d/nvidia.conf \
    && echo "/usr/local/nvidia/lib64" >> /etc/ld.so.conf.d/nvidia.conf
#
# cleanup
#
RUN npm cache clean --force && \
    conda clean --all -f -y && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

WORKDIR /home/jovyan
EXPOSE 8888
USER $NB_UID
CMD jupyter notebook --no-browser --port 8888 --ip=*
```

もう一度以下を実行して`Dockerfile`を生成する。

```bash
python manager.py --manifest="manifests/cuda.yaml" generate --cuda-version 11.2.2 --os-name ubuntu --os-version 20.04
```

今回は、`./dist/11.2.2/ubuntu2004/custom/Dockerfile`に生成された。

今回は以下のような仕様になっている。

- `x86_64`のみを対象。他のアーキテクチャは別の`./templates/ディストリビューション名/****-dockerfile.j2`ファイルを見て適宜修正する。
- エラー回避のため、いくつか加筆している。
- `devel`イメージをベースにしている。
- 他のOSに変更する場合は`ENV OS ubuntu2004`の箇所を修正する。

## 説明

### manager.py

今回の方法は、`manager.py`というファイルを使ってDockerfileのみを生成する方法を使っている。

スクリプト実行時は、`python manager.py --manifest="manifests/cuda.yaml" generate オプション`の順番で実行する。

### manifests/cuda.yaml

`manifests/cuda.yaml`には、Dockerfileを生成するための各種ライブラリのバージョン情報が記載されている。

例えば、`cuda`のバージョンは`11.2.2`、`cudnn`のバージョンは以下のように記載されているので逆引きすれば`8.1.1.33-1`がインストールされているとわかる。

```yaml
...
.components_v11.2.2: &cuda11_2_2_components
  ...
  cudnn8:
    version: 8.1.1.33-1
...
cuda_v11.2.2:
  ...
  ubuntu20.04:
    ...
    x86_64:
      <<: *cuda11_2_2_requires
      components:
        <<: *cuda11_2_2_components
    arm64:
      requires: "cuda>=11.2"
      components:
        <<: *cuda11_2_2_components
```

## 参考

> https://github.com/jupyter/docker-stacks/blob/main/base-notebook/Dockerfile