Docker 需要将宿主机的 gpu 映射到容器内部，并且在容器内部需要调用 CUDA 库。
实现思路
- Nvidia docker 在容器内支持 nvidia 库使用
- 将 cuda 的库映射到容器内部
- 在容器内部配置环境变量
操作步骤
Nvidia docker
安装方法：https://www.zywvvd.com/notes/tools/docker/nvidia-docker-install/
容器部署配置
使用 docker compose 部署容器
1
2
3
4
5
6
7
8
9
10
11
| services:
foobar:
container_name: test-container
image: ubuntu:latest
deploy:
resources:
reservations:
devices:
- driver: "nvidia"
count: "all"
capabilities: ["gpu"]
|
也可以指定部分 gpu 映射进容器
1
2
3
4
5
6
7
8
9
10
11
| services:
foobar:
container_name: test-container
image: ubuntu:latest
deploy:
resources:
reservations:
devices:
- driver: nvidia
device_ids: ["0", "3"]
capabilities: ["gpu"]
|
Cuda 库映射
在宿主机安装好 cuda 与 CuDNN 库后，在创建容器时将库路径映射进去。
需要映射 dev, cuda, gnu 文件夹：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
| services:
foobar:
container_name: test-container
image: ubuntu:latest
deploy:
resources:
reservations:
devices:
- driver: "nvidia"
count: "all"
capabilities: ["gpu"]
volumes:
- /usr/local/cuda:/usr/local/cuda
- /dev:/dev
- /usr/lib/x86_64-linux-gnu:/usr/lib/x86_64-linux-gnu
|
环境变量
增加环境变量到系统：
1
2
| export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
|
python 中可以自动添加：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
| import os
def update_env(): # 获取当前的环境变量
current_path = os.environ.get('PATH', '')
# 设置新的PATH，将CUDA的bin目录添加到当前PATH的前面
new_path = '/usr/local/cuda/bin:' + current_path
# 更新环境变量
os.environ['PATH'] = new_path
print('Updated PATH:', os.environ['PATH']) # 获取当前的环境变量
current_path = os.environ.get('LD_LIBRARY_PATH', '')
# 设置新的PATH，将CUDA的bin目录添加到当前PATH的前面
new_path = '/usr/local/cuda/lib64:' + current_path
# 更新环境变量
os.environ['LD_LIBRARY_PATH'] = new_path
print('Updated LD_LIBRARY_PATH:', os.environ['LD_LIBRARY_PATH'])
|
参考资料
文章链接：
https://www.zywvvd.com/notes/tools/docker/docker-gpu-config/
