大模型 | 使用 Docker 部署 AI 环境

AI女王

已于 2024-09-12 11:15:58 修改

阅读量854

点赞数 5

文章标签： docker 人工智能容器 python 大模型 Ai大模型

于 2024-08-11 16:15:00 首次发布

本文链接：https://blog.csdn.net/m0_56255097/article/details/141087787

版权

目标

制定一个小目标：

在 Docker 容器中可以正常执行nvidia-smi命令
在 Docker 容器中可以正常执行python -c "import torch; print(torch.cuda.is_available())"命令

好奇的同学可能会问，为什么不在 Docker 容器中执行nvcc --version命令呢？这个问题留到后面再解释。

预安装工作

我们还是以 Ubuntu 22.04 操作系统为例，首先我们要安装 Docker，这里就不再赘述了，可以参考官方文档[2]进行安装。

Docker 安装完成后，如果你想以非 root 用户执行docker命令的话，还需要将当前用户添加到docker用户组中，命令如下：

# 先确保 docker 用户组已存在  
sudo usermod -aG docker <user-name>

安装 NVIDIA 驱动

安装 NVIDIA 驱动的步骤可以参考这个

安装 NVIDIA Container Toolkit

NVIDIA Container Toolkit 允许开发者和用户将 NVIDIA GPU 的能力无缝地加入到 Docker 容器中。通过简单地安装一个插件，用户就可以运行为 GPU 优化的容器，无需进行任何修改。该工具集包括 NVIDIA 驱动、CUDA 等必要组件，确保 GPU 在容器环境中的高性能执行，对 AI、数据分析和高性能计算场景特别有用。

我们可以通过以下命令来安装 NVIDIA Container Toolkit：

# 先切换到 root 用户  
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | apt-key add -  
# 添加源  
curl -s -L https://nvidia.github.io/nvidia-docker/ubuntu22.04/nvidia-docker.list > /etc/apt/sources.list.d/nvidia-docker.list  
# 更新源  
apt update  
# 安装 NVIDIA Container Toolkit  
apt -y install nvidia-container-toolkit  
# 重启 Docker 服务  
systemctl restart docker

安装完成后，系统会一并安装 NVIDIA Container Toolkit 的 CLI 命令（nvidia-ctk），我们可以运行该命令确认安装是否成功：

$ nvidia-ctk --version  
NVIDIA Container Toolkit CLI version 1.13.5  
commit: 6b8589dcb4dead72ab64f14a5912886e6165c079

然后我们就可以在 Docker 容器中运行nvidia-smi命令，验证是否可以在 Docker 容器中正常使用 NVIDIA 显卡了，我们下载一个简单的镜像来进行验证：

$ docker run --rm --gpus all nvidia/cuda:11.7.1-base-ubuntu22.04 nvidia-smi  
  
Wed Aug 16 03:04:19 2023  
+-----------------------------------------------------------------------------+  
| NVIDIA-SMI 470.199.02   Driver Version: 470.199.02   CUDA Version: 11.7     |  
|-------------------------------+----------------------+----------------------+  
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |  
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |  
|                               |                      |               MIG M. |  
|===============================+======================+======================|  
|   0  Quadro M6000        Off  | 00000000:03:00.0 Off |                  Off |  
| 28%   38C    P8    13W / 250W |     15MiB / 12210MiB |      0%      Default |  
|                               |                      |                  N/A |  
+-------------------------------+----------------------+----------------------+  
  
+-----------------------------------------------------------------------------+  
| Processes:                                                                  |  
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |  
|        ID   ID                                                   Usage      |  
|=============================================================================|  
+-----------------------------------------------------------------------------+

可以看到，我们在 Docker 容器中可以正常使用 NVIDIA 显卡了。

下载 Docker 镜像

要运行 AI 环境，一般需要安装 CUDA 和 PyTorch，之前我们是在主机上安装这 2 个程序，但使用 Docker 的方式，我们可以直接下载已经安装好 CUDA 和 PyTorch 的镜像，这里推荐使用这个镜像：anibali/pytorch[4]，这个镜像中包含了 CUDA 和 PyTorch，我们可以通过以下命令来运行镜像：

$ docker run --rm --gpus all anibali/pytorch:1.13.0-cuda11.8-ubuntu22.04 nvidia-smi  
  
Wed Aug 16 03:09:33 2023  
+-----------------------------------------------------------------------------+  
| NVIDIA-SMI 470.199.02   Driver Version: 470.199.02   CUDA Version: 11.8     |  
|-------------------------------+----------------------+----------------------+  
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |  
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |  
|                               |                      |               MIG M. |  
|===============================+======================+======================|  
|   0  Quadro M6000        Off  | 00000000:03:00.0 Off |                  Off |  
| 28%   38C    P8    13W / 250W |     15MiB / 12210MiB |      0%      Default |  
|                               |                      |                  N/A |  
+-------------------------------+----------------------+----------------------+  
  
+-----------------------------------------------------------------------------+  
| Processes:                                                                  |  
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |  
|        ID   ID                                                   Usage      |  
|=============================================================================|  
+-----------------------------------------------------------------------------+

可以看到 CUDA 版本显示的是 11.8，跟 Docker 镜像中的 CUDA 版本是一致的。

我们再来看在 Docker 容器中是否可以正常执行python -c "import torch; print(torch.cuda.is_available())"命令：

# 进入 Docker 容器  
$ docker run -it --rm --gpus all anibali/pytorch:1.13.0-cuda11.8-ubuntu22.04 bash  
$ user@32e8c83f88c3:/app$ python -c "import torch; print(torch.cuda.is_available())"  
True

通过结果可以证明，容器中的 CUDA 和 PyTorch 程序可以正常使用。

为什么不在 Docker 容器中执行 nvcc 命令

回到原先那个问题，为什么不在 Docker 容器中执行**nvcc --version命令呢？我们可以在 Docker 容器中执行nvcc --version**命令，看看会发生什么：

user@32e8c83f88c3:/app$ nvcc --version  
bash: nvcc: command not found

发现容器中找不到nvcc命令，但是 CUDA 又是可以正常访问，再看 CUDA 的安装目录：

user@32e8c83f88c3:/app$ ls /usr/local/cuda*  
/usr/local/cuda:  
compat  lib64  targets  
  
/usr/local/cuda-11:  
compat  lib64  targets  
  
/usr/local/cuda-11.8:  
compat  lib64  targets

可以看到在 CUDA 安装目录中并没有bin文件夹（一般nvcc命令会放到这个文件夹里面），这是因为有些 Docker 镜像为了节省资源，会将一些不需要的文件去掉，只保留最核心的文件，已达到减小 Docker 镜像大小的目的。

总结

Docker 安装与主机安装的方式相比，我们少安装了 CUDA 和 PyTorch 程序，多安装了NVIDIA Container Toolkit和下载 Docker 镜像，但整体花费的时间其实是减少了的（因为 CUDA 和 PyToch 安装时间比较长）。Docker 安装最大的好处就是可以随时切换不同的环境，而且不会影响到主机的环境。比如我们今天安装了 CUDA 11.x 的版本，后面程序升级了可能需要安装 CUDA 12.x 的版本，如果是主机安装的话，就要卸载老的 CUDA 再重新安装新的 CUDA，这样就比较麻烦了，但是如果是 Docker 安装的话，我们只需要下载新的 Docker 镜像就可以了，非常方便。