如何在同一个 docker 映像上使用不同类型的 gpus（例如 1080Ti 与 2080Ti）而不需要重新运行“python setup.py develop”？

Posted 2023-04-15

【中文标题】如何在同一个 docker 映像上使用不同类型的 gpus（例如 1080Ti 与 2080Ti）而不需要重新运行“python setup.py develop”？

【英文标题】：How do I use separate types of gpus (e.g. 1080Ti vs 2080Ti) on the same docker image without needing to re-run `python setup.py develop`?

【发布时间】：2021-11-05 17:07:28

【问题描述】：

我正在使用基于 pytorch 的 repository，其中安装步骤指定使用此 setup.py file 运行 python setup.py develop。我一直在使用 docker 镜像在 1080Ti 和 1080 GPU 上运行存储库，该镜像克隆了 repo 并在构建过程中运行 setup.py 脚本。以下是从我的 Dockerfile 复制的文件。

RUN git clone https://github.com/CVMI-Lab/ST3D.git
WORKDIR /ST3D
RUN nvidia-smi
RUN pip install -r requirements.txt
RUN python setup.py develop

进入容器后，我只在 repo 中挂载特定文件夹，如下所示：

GPU_ID = 0

ENVS="  --env=NVIDIA_VISIBLE_DEVICES=$GPU_ID
        --env=CUDA_VISIBLE_DEVICES=$GPU_ID
        --env=NVIDIA_DRIVER_CAPABILITIES=all"

VOLUMES="       --volume=$DATA_PATH:/ST3D/data
                --volume=$CODE_PATH/pcdet:/ST3D/pcdet
                --volume=$CODE_PATH/tools:/ST3D/tools
                --volume=$CODE_PATH/output:/ST3D/output"

docker  run -d -it --rm \
        $VOLUMES \
        $ENVS \
        --runtime=nvidia \
        --gpus $GPU_ID \
        --privileged \
        --net=host \
        --workdir=/ST3D \
        darrenjkt/st3d:v0.3.0

最近我们在同一台计算机上安装了 2080Ti。当我使用相同的 python 脚本进入仅带有 2080Ti gpu 的同一个 docker 容器时，我收到以下错误：

RuntimeError: CUDA error: no kernel image is available for execution on the device

此错误与安装在 setup.py 中的 cpp 模块之一有关。

我可以通过再次运行python setup.py develop 来解决这个问题，这将使它能够与 2080Ti 一起工作。然后我尝试将 docker 容器提交到 2080Ti 特定的 docker 映像，并将 1080 docker 容器提交到 1080 特定的映像。但是我注意到，一旦我在 2080Ti 容器上运行 python setup.py develop，它就会给我 1080 图像的 CUDA 错误。如果我再次在 1080 gpu 上运行 setup.py，它会在 2080Ti 图像上返回 CUDA 错误。这让我感到困惑，因为我没有挂载构建文件，而是将它们单独保存在容器中并将其提交给新映像。

所以我的问题是，如何设置我的环境/docker 映像，使其不需要每次都重新构建 setup.py？

【问题讨论】：

添加参数以构建您需要支持的 GPU：pytorch.org/docs/stable/cpp_extension.html

哦，原来如此！谢谢你。我在我的 Dockerfile 中添加了RUN TORCH_CUDA_ARCH_LIST="6.1 7.5" python setup.py develop，它成功地在两个 GPU 上运行。 “默认情况下，扩展程序将编译为在扩展程序构建过程中可见的所有卡片拱门上运行，加上 PTX”根据您发送的文档链接上的这句话，我认为扩展程序会自行找到拱门因此我认为没有必要覆盖它。

如果您没有明确列出要构建的一组 GPU 架构，我认为 torch 构建系统（尝试）自动检测基于已安装 GPU 的构建设置。

【参考方案1】：

通过使用以下内容构建 docker 映像解决了该问题：

RUN git clone https://github.com/CVMI-Lab/ST3D.git
WORKDIR /ST3D
RUN nvidia-smi
RUN pip install -r requirements.txt
RUN TORCH_CUDA_ARCH_LIST="6.1 7.5" python setup.py develop

TORCH_CUDA_ARCH_LIST="6.1 7.5" 是 1080/1080Ti 的 6.1 和 2080Ti 的 7.5。这会覆盖默认的 cpp_extension 行为，以明确指定要支持的计算能力。