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

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

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

https://qiita.com/tkusumi/items/f275f0737fb5b261a868

https://medium.com/nvidiajapan/nvidia-docker-%E3%81%A3%E3%81%A6%E4%BB%8A%E3%81%A9%E3%81%86%E3%81%AA%E3%81%A3%E3%81%A6%E3%82%8B%E3%81%AE-20-09-%E7%89%88-558fae883f44

### Docker Engineをインストールする
:::note warn
インストール済みの方は飛ばしてください
:::
経験上，`nvidia-docker`を使用する際には **Docker Desktop** ではなく **Docker Engine** を使用したほうが良いです

インストールガイドの通りに **Docker Engine** をインストールします
- https://docs.docker.com/engine/install/ubuntu/#set-up-the-repository

リポジトリを追加する
```shell
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** をインストールする
```shell
sudo apt-get install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
```

dockerコマンドを一般ユーザーでも実行できるようにするため，ユーザーを`docker`グループに所属させる
```shell
sudo groupadd docker # 作成済みの場合は不要
sudo usermod -aG docker <ユーザー名>
```

インストール確認
```shell
docker ps
```

### NVIDIAドライバをインストールする
:::note warn
インストール済みの方は飛ばしてください
:::

ハードウェアに対応したNVIDIAドライバを検索する
```shell
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 と書いてあるドライバをインストールする
```shell
sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update
sudo apt install nvidia-driver-555
```

インストールできたら再起動する
```shell
sudo reboot
```

インストール確認
```shell
nvidia-smi
```

### `nvidia-container-toolkit` を使えるようにする
インストールガイドの通りにインストールします
- https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html

リポジトリを追加する
```shell
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`をインストールする
```shell
sudo apt install -y nvidia-container-toolkit
```

docker側の設定を追加する
```shell
sudo nvidia-ctk runtime configure --runtime=docker
```

`/etc/docker/daemon.json`に次の設定が追加されます
```/etc/docker/daemon.json 
{
    "runtimes": {
        "nvidia": {
            "args": [],
            "path": "nvidia-container-runtime"
        }
    }
}
```

dockerデーモンを再起動する
```shell
sudo systemctl restart docker
```


これでOKです

## PyTorch環境を作る

今回は `pip` x `ruff` x `torch` の環境を作ります．[Poetry](https://github.com/tf63/docker-images/tree/master/pytorch-poetry), [Rye](https://github.com/tf63/docker-images/tree/master/pytorch-rye)の環境も作っているので興味があれば

`Dockerfile`を用意してあります

https://github.com/tf63/docker-images/blob/master/pytorch/docker/Dockerfile.cu111

用意したビルドスクリプト[`docker.sh`](https://github.com/tf63/docker-images/blob/master/pytorch/docker.sh)を使ってコンテナを立ち上げます

```shell
bash docker.sh build
bash docker.sh shell
```

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


`PyTorch`のインストール確認
```shell
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も動きます
```shell
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にアタッチします

<img src=https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3291419/78b0d8aa-3aa4-9452-74fa-9f38e4ef00b0.png width=300px/>

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

---

以上です

## 備考: NVIDIA関連パッケージを全て削除する
NVIDIAドライバのバージョンと`nvidia-container-toolkit`の相性が悪いと上手くいかなかったりします

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

```shell
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*
```