Настройка CUDA для Docker

Часто при разработке и развертывании deep learning моделей требуется использовать GPU для ускорения вычислений. В этом случае требуется как минимум установить драйвер GPU, CUDA, а также cuDNN для некоторых фреймворков. Иногда это может занять продолжительное время из-за возникающих при установке проблем. Однако можно пойти по более простому пути и использовать Docker-контейнер от Nvidia, в котором уже установлена CUDA и cuDNN (если необходимо). В этом случае нам потребуется только установить Docker и драйвер GPU.

В этой статье мы рассмотрим настройку CUDA для Docker на примере сервера с установленной Ubuntu Server 20.04. Для подробного рассмотрения темы возьмем чистый сервер, на котором даже не установлен драйвер GPU. Описанные далее действия были успешно выполнены на сервере с RTX 2080 Ti и ноутбуке с RTX 3060. Однако в теории они могут быть выполнены для любой видеокарты, поддерживаемой вашим ML-фреймворком.

Установка драйвера GPU

Если вы используете Ubuntu Desktop, то вы можете легко установить драйвер с помощью утилиты Программы и обновления. Но мы рассмотрим способ установки с помощью терминала. Сначала обновим Ubuntu и добавим PPA репозиторий с последними версиями драйверов

sudo apt update && sudo apt upgrade -y
sudo add-apt-repository ppa:graphics-drivers/ppa -y
sudo apt update

Теперь мы можем посмотреть последнюю версию драйвера с помощью следующей команды

ubuntu-drivers devices

В моем случае драйвер имеет версию 510. Установим его и перезагрузим сервер

sudo apt install nvidia-driver-510    
sudo reboot

После перезагрузки можем проверить правильность установки с помощью команды

nvidia-smi

Она должна отобразить информацию о GPU

Установка Docker

Актуальный способ установки Docker для Ubuntu можно найти здесь. Нам необходимо выполнить команды

sudo apt-get update
sudo apt-get install ca-certificates curl gnupg lsb-release
echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu $  (lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
sudo apt-get update
sudo apt-get install docker-ce docker-ce-cli containerd.io

Docker установлен, однако при вводе каждой команды нам требуется указывать sudo. Чтобы этого не делать, выполним

sudo usermod -aG docker ${USER}   
su - ${USER}

Если при вводе второй команды возникли проблемы, то можно просто самостоятельно выйти с сервера и зайти снова.

Установка nvidia-container-toolkit

Также нам нужно установить дополнительный пакет, который позволит контейнерам использовать GPU. Выполним команды

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-container-toolkit
sudo systemctl restart docker

Альтернативный способ установки можно найти здесь

Тестируем

Поздравляю, мы завершили настройку CUDA для Docker!

Теперь можно приступать к настройке Docker своего образа. Его можно создать на основе образов от Nvidia

FROM nvidia/cuda:11.2.2-base

или разработчиков DL-фреймворков. Например,

FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime

При выборе необходимо обращать внимание на версию CUDA, т.к. определенные версии CUDA требует определенных версий драйвера GPU. С таблицей совместимости можно ознакомиться здесь.

Для тестирования можно воспользоваться небольшим приложением на Streamlit, в котором мы запустим Pytorch lightning-модель для классификации изображений. Проект можно найти в репозитории на Github. Сначала клонируем репозиторий и загрузим модель из релизов

git clone https://github.com/andBabaev/cuda_guide.git
cd cuda_guide
wget https://github.com/andBabaev/cuda_guide/releases/download/v0.1/model.ckpt -P models

Для запуска приложения в репозитории приведены примеры Dockerfile на основе образов от Nvidia и Pytorch. Образ от Pytorch уже содержит настроенный фреймворк, что ускоряет его сборку, однако итоговый размер образа больше на 1,7ГБ. Для запуска приложения на основе образа Nvidia выполним

docker build -t cuda_guide_nvidia -f Dockerfile.app.nvidia .     
docker run -p 8501:8501 --gpus all -it cuda_guide_nvidia:latest

а для запуска образа Pytorch

docker build -t cuda_guide_torch -f Dockerfile.app.torch .     
docker run -p 8501:8501 --gpus all -it cuda_guide_torch:latest

При запуске контейнера важно указать параметр –gpus, иначе GPU не будут доступны. Кроме all он также может принимать значения 0,1,2…, которые являются индексами GPU в окне nvidia-smi.

Также в репозитории можно найти ipynb-файл для обучения модели и отдельный dockerfile для его запуска. Для запуска jupyter notebook выполните

docker build -t cuda_guide_train -f Dockerfiles/Dockerfile.train .
docker run \
    -v ${PWD}/models:/app/models \
    -v ${PWD}/train:/app/train \
    -p 9999:9999 \
    --gpus all \
    --name cuda_guide_train -it cuda_guide_train:latest

На этом всё! Буду рад ответить на ваши вопросы и предложения в комментариях. Удачи!