Docker 部署 NVIDIA PyTorch 环境指南

摘要: 本文档详细介绍了如何在 GPU 节点上使用 Docker 拉取并运行 NVIDIA PyTorch 官方容器镜像。内容涵盖容器启动参数详解、SSH 服务配置、镜像保存分发以及 VSCode 远程调试环境搭建。

1. 拉取镜像

1.1 访问 NGC 仓库

NVIDIA GPU Cloud (NGC) 提供了经过高度优化的深度学习框架镜像，预装了 CUDA、cuDNN、NCCL 等组件。

• 地址: NVIDIA NGC Catalog - PyTorch

1.2 拉取命令

建议选择与宿主机驱动兼容的镜像版本（例如宿主机驱动支持 CUDA 12.x）。

# 拉取 24.07 版本 (包含 PyTorch 2.x, CUDA 12.x)
docker pull nvcr.io/nvidia/pytorch:24.07-py3

2. 启动容器 (关键步骤)

为了让容器充分利用硬件资源（GPU、IB 网络、内存），必须配置特定的启动参数。

docker run -it \
  --gpus all \
  --ipc=host \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  --privileged=true \
  --network host \
  --shm-size=20G \
  --name pytorch_dev \
  nvcr.io/nvidia/pytorch:24.07-py3

参数详解

参数	作用	必要性
--gpus all	透传所有 GPU	核心，否则无法使用 CUDA。
--ipc=host	共享宿主机 IPC	核心，多卡训练 (DDP) 通信和 DataLoader 必须。
--ulimit memlock=-1	解除内存锁定限制	核心，IB RDMA 通信必须 (Pin Memory)。
--ulimit stack=...	增加栈空间	防止深度递归或大数组导致 Stack Overflow。
--network host	使用宿主机网络	方便 SSH 直连，网络性能损耗最小。
--shm-size=20G	扩大共享内存	默认为 64M，太小会导致 PyTorch DataLoader 报错 Bus error。
--privileged=true	特权模式	允许访问设备驱动文件。

3. 容器内环境配置

进入容器后，通常需要安装 SSH 服务以便远程连接。

3.1 安装基础工具

apt-get update
apt-get install -y iproute2 net-tools inetutils-ping openssh-server vim

3.2 配置 SSH 服务

1. 修改配置: 编辑 /etc/ssh/sshd_config。

   # 更改端口 (防止与宿主机 22 冲突，前提是使用了 --network host)
   Port 20240
   # 允许 Root 登录
   PermitRootLogin yes

2. 设置密码:

   passwd root
   # 输入你的密码，例如 123456

3. 启动服务:

   service ssh restart
   # 或 /usr/sbin/sshd -D &

3.3 验证登录

在宿主机或其他机器上测试：

ssh root@<宿主机IP> -p 20240

4. 挂载数据与模型

开发过程中，代码和数据集通常存放在宿主机，通过挂载映射到容器。

docker run -it ... \
  -v /home/user/project:/workspace/project \
  -v /data/datasets:/workspace/data \
  nvcr.io/nvidia/pytorch:24.07-py3

• /home/user/project: 宿主机路径。
• /workspace/project: 容器内路径（NGC 镜像默认工作目录为 /workspace）。

5. 镜像保存与分发

当你在容器内安装了额外的 Python 库（pip install ...）后，需要保存环境。

5.1 提交为新镜像

# 查看容器 ID
docker ps

# 提交 (Commit)
docker commit <Container_ID> pytorch24:custom-v1

5.2 导出离线包 (Save)

用于在无外网环境分发。

docker save -o pytorch24-custom.tar pytorch24:custom-v1

5.3 导入 (Load)

在目标机器上：

docker load -i pytorch24-custom.tar

6. VSCode 远程调试 (Remote-SSH)

这是最高效的开发方式。

1. 安装插件: 在 VSCode 中安装 Remote - SSH 插件。
2. 配置 Host: 点击左下角 >< 图标 -> Open SSH Configuration File。

   Host gpu-container
       HostName <宿主机IP>
       Port 20240
       User root

3. 连接: 点击左下角 >< -> Connect to Host -> 选择 gpu-container。
4. 开发: 此时 VSCode 的文件资源管理器将直接显示容器内的文件，且可以使用容器内的 Python解释器进行 Debug。

7. 总结

步骤	命令/工具	关键点
拉取	docker pull	认准 nvcr.io 官方镜像
启动	docker run	务必带上 --ipc=host --shm-size
连接	ssh -p <port>	容器内需安装 openssh-server
开发	VSCode Remote	开发体验与本地一致
分发	commit / save	固化环境，并在集群分发