构建CPU运行PyTorch模型的高效Dockerfile：两阶段构建问题排查

我需要构建一个仅用于CPU运行PyTorch模型的最小体积Docker镜像，单阶段构建能正常拉取CPU版PyTorch，但使用以下两阶段构建代码时，Docker会自动下载PyTorch的CUDA/GPU版本：

FROM python:3.11-slim as builder

WORKDIR /app

# Set environment variables.
ENV PYTHONDONTWRITEBYTECODE 1
ENV PYTHONUNBUFFERED 1

RUN apt-get update && \
    apt-get install -y --no-install-recommends gcc

# Copy local code to the container image.
COPY requirements.txt .

# Install dependencies & model files
RUN pip install --no-cache-dir torch==2.0.1+cpu --index-url https://download.pytorch.org/whl/cpu && \
    pip wheel --no-cache-dir --no-deps --wheel-dir /app/wheels -r requirements.txt

FROM python:3.11-slim

WORKDIR /app

# Set environment variables.
ENV PORT 8080
ENV HOST 0.0.0.0

COPY --from=builder /app/wheels /wheels
COPY --from=builder /app/requirements.txt .

RUN pip install --no-cache /wheels/* && \
    huggingface-cli login --token xxx && \
    python -c 'from sentence_transformers import SentenceTransformer; SentenceTransformer("xxx", cache_folder="./app/artefacts")'

# Start the container
CMD python -m uvicorn app.main:app --host $HOST --port $PORT --workers 1

后来我修改了构建步骤，把PyTorch也改成用pip wheel处理，成功避免了CUDA版本的安装，但镜像大小从单阶段的1.5GB涨到了2.5GB，修改的代码如下：

RUN pip wheel --no-cache-dir --no-deps --wheel-dir /app/wheels torch==2.0.1 --index-url https://download.pytorch.org/whl/cpu && \
    pip wheel --no-cache-dir --no-deps --wheel-dir /app/wheels -r requirements.txt

问题分析 #

1. 初始两阶段构建拉取CUDA版本的原因：第一阶段先pip install了CPU版PyTorch，但后续pip wheel处理requirements.txt时，如果文件中仅写了torch==2.0.1（未指定+cpu），pip会默认从PyPI拉取CUDA版本的wheel——因为pip wheel是独立解析依赖，不会继承之前pip install的环境配置，必须显式指定CPU版的索引地址。
2. 镜像臃肿的原因：修改后的构建虽然解决了CUDA问题，但存在优化遗漏：
   • 第一阶段安装的编译依赖（如gcc）未清理，增加了临时镜像的体积；
   • 第二阶段安装依赖后未删除wheel文件和pip缓存；
   • 模型下载在第二阶段完成，未做冗余文件清理，且可能存在重复存储问题。

优化后的Dockerfile #

FROM python:3.11-slim as builder

WORKDIR /app

ENV PYTHONDONTWRITEBYTECODE 1
ENV PYTHONUNBUFFERED 1

# 安装编译依赖，打完wheel后立即清理
RUN apt-get update && \
    apt-get install -y --no-install-recommends gcc && \
    # 合并打包所有依赖（含CPU版PyTorch），指定CPU索引
    pip wheel --no-cache-dir --no-deps --wheel-dir /app/wheels -r requirements.txt --index-url https://download.pytorch.org/whl/cpu && \
    # 清理编译依赖与apt缓存
    apt-get purge -y --auto-remove gcc && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

# 在builder阶段预下载模型，避免第二阶段重复操作
RUN huggingface-cli login --token xxx && \
    python -c 'from sentence_transformers import SentenceTransformer; SentenceTransformer("xxx", cache_folder="./artefacts")'

FROM python:3.11-slim

WORKDIR /app

ENV PORT 8080
ENV HOST 0.0.0.0

# 仅拷贝必要文件：wheel包和预下载的模型
COPY --from=builder /app/wheels /wheels
COPY --from=builder /app/artefacts ./artefacts

# 安装依赖并立即清理冗余文件
RUN pip install --no-cache-dir /wheels/* && \
    rm -rf /wheels && \
    rm -rf /root/.cache/pip

# 启动服务
CMD python -m uvicorn app.main:app --host $HOST --port $PORT --workers 1

关键优化点 #

1. 统一依赖打包：在requirements.txt中明确写入torch==2.0.1+cpu，用一条pip wheel命令打包所有依赖并指定CPU索引，避免拆分命令导致的依赖解析错误。
2. 清理builder阶段冗余：打完wheel后立即卸载gcc并清理apt缓存，减小临时镜像体积，减少构建过程中的磁盘占用。
3. 预下载模型到builder阶段：将模型下载移到builder阶段，避免第二阶段重复执行登录和下载操作，减少临时文件残留。
4. 第二阶段彻底清理：安装完依赖后删除wheel目录和pip缓存，避免冗余文件占用镜像空间。

额外注意事项 #

• 确保requirements.txt中所有依赖都不依赖CUDA版本的PyTorch，必要时可添加--force-reinstall参数强制使用CPU版依赖。
• 如果huggingface token敏感，建议使用构建参数传递，避免硬编码在Dockerfile中。

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