Ubuntu で nvidia-container-toolkit を使って PyTorch の環境を作りたい

Docker上でGPUを使用する際には専用のパッケージをインストールする必要があります．以前はnvidia-docker2とnvidia-container-runtimeの2つをインストールしていましたが，現在はnvidia-container-toolkitというパッケージに統合されています

nvidia-container-toolkitの詳細はこちらの記事が詳しいです

Docker Engineをインストールする

インストール済みの方は飛ばしてください

経験上，nvidia-dockerを使用する際には Docker Desktop ではなく Docker Engine を使用したほうが良いです

インストールガイドの通りに Docker Engine をインストールします

• https://docs.docker.com/engine/install/ubuntu/#set-up-the-repository

リポジトリを追加する

sudo apt update
sudo apt install ca-certificates curl
sudo install -m 0755 -d /etc/apt/keyrings
sudo curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
sudo chmod a+r /etc/apt/keyrings/docker.asc

# Add the repository to Apt sources:
echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu \
  $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \
  sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
sudo apt update

Docker Engine をインストールする

sudo apt-get install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

dockerコマンドを一般ユーザーでも実行できるようにするため，ユーザーをdockerグループに所属させる

sudo groupadd docker # 作成済みの場合は不要
sudo usermod -aG docker <ユーザー名>

インストール確認

docker ps

NVIDIAドライバをインストールする

インストール済みの方は飛ばしてください

ハードウェアに対応したNVIDIAドライバを検索する

ubuntu-drivers devices

driver   : nvidia-driver-555-open - third-party non-free
driver   : nvidia-driver-535 - distro non-free
driver   : nvidia-driver-550-open - third-party non-free
driver   : nvidia-driver-535-open - distro non-free
driver   : nvidia-driver-545-open - distro non-free
driver   : nvidia-driver-550 - third-party non-free
driver   : nvidia-driver-535-server - distro non-free
driver   : nvidia-driver-470-server - distro non-free
driver   : nvidia-driver-555 - third-party non-free recommended
driver   : nvidia-driver-545 - distro non-free
driver   : nvidia-driver-535-server-open - distro non-free
driver   : nvidia-driver-470 - distro non-free
driver   : xserver-xorg-video-nouveau - distro free builtin

基本的には recommended と書いてあるドライバをインストールする

sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update
sudo apt install nvidia-driver-555

インストールできたら再起動する

sudo reboot

インストール確認

nvidia-smi

nvidia-container-toolkit を使えるようにする

インストールガイドの通りにインストールします

• https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html

リポジトリを追加する

curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
  && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

nvidia-container-toolkitをインストールする

sudo apt install -y nvidia-container-toolkit

docker側の設定を追加する

sudo nvidia-ctk runtime configure --runtime=docker

/etc/docker/daemon.jsonに次の設定が追加されます

/etc/docker/daemon.json

{
    "runtimes": {
        "nvidia": {
            "args": [],
            "path": "nvidia-container-runtime"
        }
    }
}

dockerデーモンを再起動する

sudo systemctl restart docker

これでOKです

PyTorch環境を作る

今回は pip x ruff x torch の環境を作ります．Poetry, Ryeの環境も作っているので興味があれば

Dockerfileを用意してあります

用意したビルドスクリプトdocker.shを使ってコンテナを立ち上げます

bash docker.sh build
bash docker.sh shell

VSCodeに付属しているターミナルでコンテナを立ち上げると，ウィンドウを閉じた際にコンテナが落ちます．ターミナルアプリからコンテナを立ち上げましょう

PyTorchのインストール確認

python

Python 3.8.10 (default, Jul 29 2024, 17:02:10) 
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> print(torch.cuda.get_device_name())
NVIDIA GeForce RTX 3060 Ti

ruffも動きます

make format

ruff format src
2 files left unchanged
ruff check --fix src
All checks passed!

VSCodeをコンテナにアタッチする

VSCodeを使用している場合，Docker拡張機能を使ってコンテナ内で作業できます

devcontainer.jsonを書くのは面倒なので Dockerタブ > Attach Visual Studio Codeを選び，起動済みのコンテナをVSCodeにアタッチします

(初回のみ) Open Folderから/appを開きます．拡張機能タブで @recommendedと検索すると，私の環境で使用している拡張機能をインストールできます

---

以上です

備考: NVIDIA関連パッケージを全て削除する

NVIDIAドライバのバージョンとnvidia-container-toolkitの相性が悪いと上手くいかなかったりします

その際はNVIDIA関連パッケージを全て削除してからインストールし直すと上手くいったりします

sudo apt --purge remove nvidia*
sudo apt --purge remove cuda*
sudo apt --purge remove cudnn*
sudo apt --purge remove libnvidia*
sudo apt --purge remove libcuda*
sudo apt --purge remove libcudnn*
sudo apt autoremove
sudo apt autoclean
sudo apt update
sudo rm -rf /usr/local/cuda*