巧用Docker快速部署GPU環(huán)境
在 Linux 服務器上使用 GPU 跑深度學習的模型很正常不過。如果我們想用 Docker 實現(xiàn)同樣的需求，就需要做些額外的工作。本質(zhì)上就是我們要在容器里能看到并且使用宿主機上的顯卡。在這篇文章里我們就介紹一下 Docker 使用 GPU 的環(huán)境搭建。
Nvidia 驅(qū)動
某些命令以 Ubuntu 作為示例。首先宿主機上必現(xiàn)安裝 Nvidia 驅(qū)動。
這里推薦從 Nvidia 官網(wǎng)下載腳本安裝，安裝和卸載都比較方便并且適用于任何 Linux 發(fā)行版，包括 CentOS，Ubuntu 等。NVIDIA Telsa GPU 的 Linux 驅(qū)動在安裝過程中需要編譯 kernel module，系統(tǒng)需提前安裝 gcc 和編譯 Linux Kernel Module 所依賴的包，例如 kernel-devel-$(uname -r) 等。
安裝 gcc 和 kernel-dev
- $ sudo apt install gcc kernel-dev -y
安裝 Nvidia 驅(qū)動
- 訪問 https://www.nvidia.com/Download/Find.aspx
- 選擇對應操作系統(tǒng)和安裝包，并單擊 [SEARCH] 搜尋驅(qū)動，選擇要下載的驅(qū)動版本
3. 在宿主機上下載并執(zhí)行對應版本安裝腳本
- $ wget https://www.nvidia.com/content/DriverDownload-March2009/confirmation.php?url=/tesla/450.80.02/NVIDIA-Linux-x86_64-450.80.02.run&lang=us&type=Tesla
- $ chmod +x NVIDIA-Linux-x86_64-450.80.02.run && ./NVIDIA-Linux-x86_64-450.80.02.run
4. 驗證
使用 nvidia-smi 命令驗證是否安裝成功，如果輸出類似下圖則驅(qū)動安裝成功。
CUDA 驅(qū)動
CUDA（Compute Unified Device Architecture）是顯卡廠商 NVIDIA 推出的運算平臺。CUDA™是一種由 NVIDIA 推出的通用并行計算架構，該架構使 GPU 能夠解決復雜的計算問題。它包含了 CUDA 指令集架構（ISA）以及 GPU 內(nèi)部的并行計算引擎。這里安裝的方式和顯卡驅(qū)動安裝類似。
1. 訪問官網(wǎng)下載對應版本安裝包，https://developer.nvidia.com/cuda-toolkit-archive
2. 配置環(huán)境變量
- $ echo 'export PATH=/usr/local/cuda/bin:$PATH' | sudo tee /etc/profile.d/cuda.sh
- $ source /etc/profile
nvidia-docker2
Docker 的安裝這里就不展開了，具體查看官方文檔非常詳細。
這里我們就直接介紹安裝 nvidia-docker2.
既然叫 nvidia-docker2 就有 nvidia-docker1 就是它的 1.0 版本目前已經(jīng)廢棄了，所以注意不要裝錯。
這里先簡單說一下 nvidia-docker2 的原理，nvidia-docker2 的依賴由下幾部分組成.
- libnvidia-container
- nvidia-container-toolkit
- nvidia-container-runtime
nvidia-container-runtime 是在 runc 基礎上多實現(xiàn)了 nvidia-container-runime-hook (現(xiàn)在叫 nvidia-container-toolkit)，該 hook 是在容器啟動后（Namespace已創(chuàng)建完成），容器自定義命令(Entrypoint)啟動前執(zhí)行。當檢測到 NVIDIA_VISIBLE_DEVICES 環(huán)境變量時，會調(diào)用 libnvidia-container 掛載 GPU Device 和 CUDA Driver。如果沒有檢測到 NVIDIA_VISIBLE_DEVICES 就會執(zhí)行默認的 runc。
下面分兩步安裝
1. 設置 repository 和 GPG key
- $ 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
2. 安裝
- $ sudo apt-get update
- $ sudo apt-get install -y nvidia-docker2
- $ sudo systemctl restart docker
3. 驗證
執(zhí)行以下命令：
- $ docker run --rm --gpus all nvidia/cuda:10.2-base nvidia-smi
如果輸出跟直接在宿主機上執(zhí)行 nvidia-smi 一致則說明安裝成功。如果跑的深度學習模型使用的是 tensorflow 可以在容器里執(zhí)行:
- import tensorflow as tf
- tf.contrib.eager.num_gpus()
如果輸出了宿主機上的 Nvidia 顯卡數(shù)量，則模型能使用到顯卡加速。如果使用的是 pytorch 可以在容器里執(zhí)行：
- import torch
- torch.cuda.is_available()
如果輸出 True 證明環(huán)境也成功了，可以使用顯卡。
4. 使用示例
- 使用所有顯卡
- $ docker run --rm --gpus all nvidia/cuda nvidia-smi
- $ docker run --rm --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all nvidia/cuda nvidia-smi
- 指明使用哪幾張卡
- $ docker run --gpus '"device=1,2"' nvidia/cuda nvidia-smi
- $ docker run --rm --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=1,2 nvidia/cuda nvidia-smi
到這里在 Docker 下使用 Nvidia 顯卡加速計算的基礎環(huán)境搭建就介紹完了。后續(xù)我們可以繼續(xù)研究一下 K8S 下調(diào)度 GPU 的實現(xiàn)。
