在nvidia/cuda Docker容器无法安装CUDA版PyTorch：wheel不兼容

我定义了如下镜像（./app/Dockerfile）：

FROM nvidia/cuda:11.3.1-base-ubuntu20.04
ENV TZ="Asia/Kuala_Lumpur"
################################################################################
#
#   Ubuntu 20.04, Cuda 11.3, python3.10
#
################################################################################

SHELL ["/bin/bash", "-c"]

# apt update
RUN apt update

# install python 3.10
RUN apt install -y software-properties-common && \
    add-apt-repository -y ppa:deadsnakes/ppa && \
    apt install -y python3.10 && \
    apt-get install -y \
    python3-pip python3.10-dev python3.10-distutils python3.10-venv && \
    cd /usr/local/bin && \
    ln -s /usr/bin/python3.10 python

然后在docker compose文件中定义了如下服务：

services:
  app:
    runtime: nvidia
    build:
      context: ./app
    image: ubuntu-cuda

当我在容器内尝试pip安装特定版本的PyTorch时：

pip install --upgrade pip

pip install https://download.pytorch.org/whl/cu113/torch-1.12.1%2Bcu113-cp310-cp310-linux_x86_64.whl

出现错误：

ERROR: torch-1.12.1+cu113-cp310-cp310-linux_x86_64.whl is not a supported wheel on this platform.

我已确认容器内成功安装了Python3.10：

python --version

输出：

Python 3.10.13

但无法在容器内运行nvidia-smi，不过执行以下命令时可正常输出：

docker run --rm --runtime=nvidia --gpus all ubuntu-cuda  nvidia-smi

输出为：

Tue Sep  5 12:03:02 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.161.03   Driver Version: 470.161.03   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |

容器内执行uname -a的输出：

Linux 0fe451225d84 5.4.0-135-generic #152-Ubuntu SMP Wed Nov 23 20:19:22 UTC 2022 x86_64 x86_64 x86_64 GNU/Linux

看起来不是Python版本或平台的问题，那可能是什么原因？我在配置相似的主机上可以正常执行该pip安装命令。

主机上的uname -r输出：

Linux opus 5.4.0-135-generic #152-Ubuntu SMP Wed Nov 23 20:19:22 UTC 2022 x86_64 x86_64 x86_64 GNU/Linux

nvidia-smi输出：

Tue Sep  5 12:08:22 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.161.03   Driver Version: 470.161.03   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|

python --version输出：

Python 3.10.9

核心原因排查 #

1. pip与Python版本不匹配：虽然容器内python指向3.10，但系统默认的pip可能仍绑定Ubuntu 20.04自带的Python3.8，导致安装时用错解释器。
2. 基础镜像缺少CUDA依赖：使用的nvidia/cuda:11.3.1-base仅包含最精简的CUDA基础文件，缺少PyTorch识别CUDA环境所需的完整runtime库。
3. 容器GPU资源未正确加载：docker-compose的runtime: nvidia配置可能未生效，导致容器内无法识别CUDA设备，pip因此判定wheel包不兼容。

解决步骤 #

1. 对齐pip与Python3.10版本 #

在容器内执行以下命令，确认pip绑定的Python版本：

pip --version

若输出显示不是Python3.10，重新安装对应版本的pip：

curl https://bootstrap.pypa.io/get-pip.py | python3.10

之后用python -m pip执行安装，避免版本混淆：

python -m pip install --upgrade pip
python -m pip install https://download.pytorch.org/whl/cu113/torch-1.12.1%2Bcu113-cp310-cp310-linux_x86_64.whl

2. 更换完整CUDA runtime镜像 #

修改Dockerfile的基础镜像为包含完整CUDA runtime的版本，确保PyTorch能识别CUDA环境：

FROM nvidia/cuda:11.3.1-runtime-ubuntu20.04
ENV TZ="Asia/Kuala_Lumpur"

重新构建镜像：

docker-compose build

3. 确保容器正确加载GPU资源 #

如果使用Docker Compose V2，改用更规范的资源配置：

services:
  app:
    build:
      context: ./app
    image: ubuntu-cuda
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

若使用旧版本Docker，启动容器时手动指定GPU参数：

docker-compose run --rm --gpus all app

4. 验证环境配置 #

在容器内执行以下命令，确认CUDA环境可用：

python -c "import torch; print(torch.cuda.is_available())"

返回True则说明配置成功。

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