Docker Run可访问GPU但Build不行，如何在构建时使用PyTorch&CUDA？

环境与背景 #

• 正在为NVIDIA GPU构建DeepStream Docker镜像，已安装NVIDIA Container Toolkit，原始Dockerfile可正常完成构建。
• 构建完成后，通过以下命令可启动带GPU支持的容器，且容器内可正常使用PyTorch：

sudo docker run --runtime=nvidia --gpus all --name Test -it deepstream:dgpu

遇到的问题 #

尝试在docker build过程中安装并使用PyTorch，当构建序列中执行导入PyTorch并调用CUDA的代码时，出现如下错误：

#0 0.895 Traceback (most recent call last):
#0 0.895   File "./X.py", line 15, in <module>
#0 0.895     dummy_input = torch.randn([1, 3, 224, 224], device='cuda')
#0 0.895   File "/usr/local/lib/python3.8/dist-packages/torch/cuda/__init__.py", line 229, in _lazy_init
#0 0.895     torch._C._cuda_init()
#0 0.895 RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx

本地Docker环境的NVIDIA驱动正常：

NVIDIA-SMI 510.47.03    Driver Version: 510.47.03    CUDA Version: 11.6 

核心矛盾：docker build过程无法识别NVIDIA驱动或GPU，只有容器启动后才能正常使用GPU；且docker build没有类似--runtime=nvidia --gpus all的参数可以直接启用GPU支持。

更新信息 #

问题疑似与BuildKit版本相关，目前仅找到通过设置DOCKER_BUILDKIT=0禁用BuildKit的解决方法，寻求其他更合适的方案。

方法1：启用BuildKit的GPU支持（推荐） #

新版本BuildKit已支持在构建阶段暴露GPU，操作步骤如下：

1. 确保NVIDIA Container Toolkit配置正确，重启Docker服务：

sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

2. 使用docker buildx build配合--gpus参数执行构建：

DOCKER_BUILDKIT=1 docker buildx build --gpus all -t deepstream:dgpu .

方法2：避免构建阶段触发CUDA初始化 #

如果仅需在构建阶段安装PyTorch，不需要实际运行CUDA代码，可调整代码逻辑：

• 修改PyTorch相关代码，仅在容器运行时触发CUDA初始化：

import torch

if __name__ == "__main__":
    # 仅在容器启动后检查CUDA可用性并执行相关操作
    if torch.cuda.is_available():
        dummy_input = torch.randn([1, 3, 224, 224], device='cuda')

• 修改Dockerfile，确保构建阶段仅执行PyTorch安装操作，不运行需要CUDA的代码。

方法3：使用NVIDIA CUDA基础镜像作为构建底座 #

选择包含CUDA runtime的NVIDIA官方镜像作为构建阶段的基础镜像，确保构建环境具备完整的CUDA依赖：

# 构建阶段使用匹配CUDA版本的NVIDIA基础镜像
FROM nvidia/cuda:11.6.2-cudnn8-runtime-ubuntu20.04 AS builder

# 安装PyTorch（指定对应CUDA版本）
RUN pip3 install torch==1.13.1+cu116 torchvision==0.14.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116

# 后续构建步骤...

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