概要

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をインストールするには以下のコマンド。

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の生成方法

リポジトリをクローン

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

ビルド済みのDockerfileを削除

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

rm -rf ./cuda/dist/*

仮想環境を作成

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

conda create -n build_cuda python=3.10.6
conda activate build_cuda

poetryをインストール

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

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

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

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

cd cuda

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

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

~/.local/bin/poetry install

Dockerfileを生成

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

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と追加する。

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

            globber = f"*"
            ....

ファイルを作成

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

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

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を生成する。

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がインストールされているとわかる。

...
.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