在 Docker 中配置所有容器默认使用 GPU
引言
在当今的开发环境中，GPU（图形处理单元）已经成为加速计算的重要工具，尤其在深度学习和数据科学领域。为了充分利用 GPU，我们需要确保在 Docker 容器中可以轻松地使用 GPU。本文将带你了解如何在 Docker 中配置所有的容器默认使用 GPU。
整体流程
下面是实现此功能的基本步骤：
| 步骤 | 描述 |
|---|---|
| 1 | 确保安装 NVIDIA 驱动和 NVIDIA Container Toolkit |
| 2 | 编辑 Docker 配置以启用 GPU |
| 3 | 使用环境变量设置默认 GPU 选项 |
| 4 | 验证配置是否成功 |
步骤详解
步骤 1: 安装 NVIDIA 驱动和 NVIDIA Container Toolkit
首先，确保你的机器上已经安装了 NVIDIA 驱动。然后安装 NVIDIA Container Toolkit，它允许 Docker 容器访问 GPU。
# 添加 NVIDIA 的官方包源
curl -s -L | sudo apt-key add -
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
# 更新包源并安装 nvidia-docker2
sudo apt-get update && sudo apt-get install -y nvidia-docker2
# 重启 Docker 守护进程
sudo systemctl restart docker
curl ...：添加 NVIDIA 的官方源。apt-get install -y nvidia-docker2：安装 NVIDIA Docker 组件。systemctl restart docker：重启 Docker 守护进程以应用新配置。
步骤 2: 编辑 Docker 配置以启用 GPU
接下来，编辑 Docker 的配置文件，以使所有容器都能使用 GPU。
# 打开或创建 Docker 的 daemon.json 文件
sudo nano /etc/docker/daemon.json
# 在文件中添加以下内容
{
"default-runtime": "nvidia",
"runtimes": {
"nvidia": {
"path": "nvidia-container-runtime",
"runtimeArgs": []
}
}
}
default-runtime: 将默认运行时设置为 NVIDIA。runtimes: 定义 NVIDIA 运行时的路径及其他参数。
保存文件后，重启 Docker 守护进程：
sudo systemctl restart docker
步骤 3: 使用环境变量设置默认 GPU 选项
可以使用环境变量设置默认 GPU 选项，确保在启动容器时自动启用 GPU。
# 设置默认环境变量
export NVIDIA_VISIBLE_DEVICES=all
NVIDIA_VISIBLE_DEVICES=all: 指定访问所有 GPU。
你可以把这条命令写入 ~/.bashrc 文件，这样每次启动终端时都会自动设置。
echo 'export NVIDIA_VISIBLE_DEVICES=all' >> ~/.bashrc
source ~/.bashrc
步骤 4: 验证配置是否成功
最后，通过运行一个简单的 Docker 容器来验证配置是否成功。
# 使用 nvidia-smi 检查 GPU 使用情况
docker run --rm nvidia/cuda:11.0-base nvidia-smi
docker run --rm nvidia/cuda:11.0-base nvidia-smi: 运行一个包含nvidia-smi的容器，查看 GPU 状态。
如果配置成功，你应该能看到 GPU 的详细信息。
相关图示
序列图
以下是整个过程的序列图，描述了步骤之间的交互关系：
sequenceDiagram
participant User
participant NvidiaDriver
participant Docker
participant Container
User->>NvidiaDriver: 安装 NVIDIA 驱动
User->>Docker: 安装 NVIDIA Container Toolkit
User->>Docker: 编辑 daemon.json
Docker->>Docker: 重启服务
User->>Container: 启动容器
Container->>NvidiaDriver: 检查 GPU 状态
类图
以下是相关类的组成：
classDiagram
class Docker {
+installNvidiaDriver()
+editDaemonJson()
+restartService()
+runContainer()
}
class NvidiaDriver {
+install()
+checkStatus()
}
class Container {
+accessGPUs()
+executeCommand()
}
Docker --> NvidiaDriver : Utilizes
Docker --> Container : Runs
结尾
通过遵循上述步骤，你已经成功配置了 Docker，以确保所有的容器默认使用 GPU。这不仅加速了计算过程，也为深度学习和大数据分析提供了强大的支持。如果在实施过程中遇到问题，请随时查阅相应的文档或寻求社区的帮助。希望这篇文章能帮助你顺利入门 GPU 计算！
