AWS EC2 p2.xlarge Docker运行PyTorch提示无NVIDIA驱动无法识别CUDA问题

1. 先确认EC2主机侧基础配置正确 #

容器要读取GPU资源，必须依赖宿主机预装的NVIDIA驱动和容器GPU工具链，缺少任一组件都会触发找不到驱动的报错：

• 先验证宿主机驱动状态：在EC2实例终端执行nvidia-smi，能正常输出GPU型号、驱动版本、CUDA版本信息才算驱动安装正常
• 安装nvidia-container-toolkit工具链并重启Docker服务：

执行以下命令完成工具链安装：

curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/ubuntu18.04/libnvidia-container.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt update && sudo apt install -y nvidia-container-toolkit
sudo systemctl restart docker

2. 补全docker-compose.yml的GPU配置 #

Docker默认不会把GPU设备透传到容器内，必须在compose配置中显式指定运行时和GPU资源：

• 在你的docker-compose.yml中添加GPU相关配置，参考示例如下：

version: '3.8'
services:
  # 替换成你自己的服务名
  impaction-detect:
    build: .
    # 必须指定nvidia运行时才能加载GPU驱动
    runtime: nvidia
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

• 配置完成后使用docker compose up启动服务，不要使用旧版docker-compose命令

3. 修正Dockerfile中的已知问题 #

当前Dockerfile存在三处会影响CUDA识别的错误：

• pip安装命令笔误：RUN pip install pip install nvidia-ml-py3多写了一次pip install，修正为RUN pip install nvidia-ml-py3
• CUDA版本不匹配：基础镜像用的是CUDA11.2版本，但安装的PyTorch是cu113版本，建议将基础镜像替换为nvidia/cuda:11.3.1-cudnn8-devel-ubuntu18.04，版本对齐可避免兼容问题
• 无效的conda逻辑：你没有在镜像中安装conda，后续magma的条件安装逻辑完全无效，不需要的话可以直接删除对应的ARG和RUN代码段，避免多余报错

4. 结果验证 #

容器启动后先在容器内执行nvidia-smi，如果能正常输出GPU信息说明透传配置正常，再执行python -c "import torch; print(torch.cuda.is_available())"，输出True即说明问题解决。如果nvidia-smi正常但PyTorch仍识别不到CUDA，可重新安装和基础镜像CUDA版本完全匹配的PyTorch版本即可。

内容的提问来源于stack exchange，提问作者James Faure