Docker构建/运行Flask应用时，transformers模型分片下载卡顿

问题背景 #

本地运行的Flask应用使用transformers加载BLIP2模型正常，但Docker镜像构建或运行时，进程卡在模型分片下载步骤：

Downloading shards:   0%|          | 0/2 [01:07<?, ?it/s]

尝试在Dockerfile中添加构建阶段预加载模型的命令后，构建过程同样卡在下载环节。

相关代码 #

Flask主程序（blip-server.py）

import os
from flask import Flask, request, jsonify
from utils import load_blip_model, generate_caption

# Load the BLIP model
processor, model = load_blip_model()

app = Flask(__name__)

@app.route("/predict", methods=["POST"])
def predict():
    if "image" not in request.files:
        return jsonify({"error": "No image part in the request"}), 400

    file_storage = request.files["image"]
    image_path = "temp_image.jpg"
    file_storage.save(image_path)

    generated_text = generate_caption(model, processor, image_path)
    os.remove(image_path)

    return jsonify({"caption": generated_text})

if __name__ == "__main__":
    app.run(debug=False, host="0.0.0.0", port=5000)

模型加载工具（utils.py）

from transformers import AutoProcessor, Blip2ForConditionalGeneration, BitsAndBytesConfig

def load_blip_model():
    processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
    quantization_config = BitsAndBytesConfig(load_in_8bit_fp32_cpu_offload=True, device_map="auto")
    model = Blip2ForConditionalGeneration.from_pretrained(
        "Salesforce/blip2-opt-2.7b",
        quantization_config=quantization_config
    )
    return processor, model

原Dockerfile

FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu20.04

RUN apt-get update && \
    DEBIAN_FRONTEND=noninteractive apt-get install -y \
    wget build-essential libreadline-gplv2-dev libncursesw5-dev \
    libssl-dev libsqlite3-dev tk-dev libgdbm-dev libc6-dev libbz2-dev \
    libffi-dev zlib1g-dev liblzma-dev libexpat1-dev libglib2.0-0 \
    libsm6 libxext6 libxrender-dev libgl1-mesa-glx patch git tree \
    && rm -rf /var/lib/apt/lists/*

# Install Python
ENV PYTHON_VERSION 3.10.14
RUN cd /usr/src && \
    wget https://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VERSION.tar.xz && \
    tar xvf Python-$PYTHON_VERSION.tar.xz && \
    cd Python-$PYTHON_VERSION && \
    ./configure --enable-optimizations --enable-shared && \
    make -j $(nproc) && make altinstall && ldconfig

# Install pip
RUN wget https://bootstrap.pypa.io/get-pip.py && \
    python3.10 get-pip.py && rm get-pip.py

WORKDIR /usr/src/app
COPY requirements.txt .
RUN pip3.10 install --no-cache-dir -r requirements.txt
COPY . .

EXPOSE 5000
ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
CMD ["python3.10", "blip-server.py"]

原因分析 #

1. 容器网络限制：Docker默认网络环境下，访问Hugging Face模型仓库速度较慢，大体积模型分片容易超时卡住。
2. 构建阶段无GPU资源：预加载模型时使用了8bit量化配置，但Docker构建阶段没有GPU，导致模型加载逻辑异常，间接卡住下载进程。
3. 模型分片体积大：BLIP2-opt-2.7b包含2个大体积分片，网络条件差时下载极易中断或停滞。

解决方案 #

方案1：本地提前下载模型，构建时复制进容器 #

• 本地运行一次load_blip_model()，模型会自动下载到~/.cache/huggingface/hub目录。
• 修改Dockerfile，添加COPY命令将本地模型缓存复制到容器对应路径：

  # 在COPY . .之前添加
  COPY ~/.cache/huggingface/hub /root/.cache/huggingface/hub
  # 或自定义缓存路径
  ENV TRANSFORMERS_CACHE=/usr/src/app/model_cache
  COPY ./model_cache $TRANSFORMERS_CACHE
  

• 修改utils.py开启离线加载模式：

  def load_blip_model():
      processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b", local_files_only=True)
      quantization_config = BitsAndBytesConfig(load_in_8bit_fp32_cpu_offload=True, device_map="auto")
      model = Blip2ForConditionalGeneration.from_pretrained(
          "Salesforce/blip2-opt-2.7b",
          quantization_config=quantization_config,
          local_files_only=True
      )
      return processor, model
  

方案2：优化Docker构建网络 #

• 构建镜像时使用宿主机网络提升下载速度：

  docker build --network host -t blip:v1 .
  

• 在Dockerfile中添加Hugging Face镜像源加速下载：

  # 安装pip后添加
  ENV HF_ENDPOINT=https://hf-mirror.com
  

方案3：调整构建阶段预加载逻辑 #

• 构建阶段预加载时暂时去掉量化配置（无GPU环境不支持），运行时再加载量化模型：

  # 安装依赖后添加
  RUN python3.10 -c "from transformers import AutoProcessor, Blip2ForConditionalGeneration; AutoProcessor.from_pretrained('Salesforce/blip2-opt-2.7b'); Blip2ForConditionalGeneration.from_pretrained('Salesforce/blip2-opt-2.7b')"
  

  此方式会在构建阶段下载完整模型，运行时直接加载缓存文件，配合原有量化配置使用。

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