Docker部署Flask+Hugging Face模型重复下载问题优化求助

问题概述 #

基于Flask+Gunicorn开发的应用，搭配850MB的Hugging Face模型，部署Docker容器时遇到以下问题：

• 容器启动时模型下载失败，陷入循环重试
• 收到API请求时，Gunicorn重启worker并重新触发模型下载
• 构建镜像时内置模型会使镜像体积超过2GB，不符合轻量化需求（无模型镜像仅50MB）

优化方案 #

1. 调整模型加载逻辑，避免重复下载 #

修改Flask应用代码，将模型初始化从应用创建阶段移出，改为容器启动后异步加载，并添加锁机制防止重复下载：

from flask import Flask, request, jsonify
from happytransformer import HappyTextToText, TTSettings
import threading
import os

model_happy_tt = None
args = None
model_loaded = False
load_lock = threading.Lock()

def load_model():
    global model_happy_tt, args, model_loaded
    with load_lock:
        if not model_loaded:
            # 指定模型缓存目录，便于后续持久化
            os.environ["TRANSFORMERS_CACHE"] = "/app/model_cache"
            os.makedirs("/app/model_cache", exist_ok=True)
            
            model_happy_tt = HappyTextToText("T5", "vennify/t5-base-grammar-correction")
            args = TTSettings(num_beams=5, min_length=1)
            # 预热模型
            model_happy_tt.generate_text("grammar: It's a jst run!!", args=args)
            model_loaded = True

def create_app():
    app = Flask(__name__)
    
    # 容器启动后异步加载模型
    threading.Thread(target=load_model, daemon=True).start()
    
    @app.route('/')
    def welcome():
        return 'Welcome to ProofReader Service'
    
    @app.route('/proofread', methods=["GET", "POST"])
    def proofread():
        global model_happy_tt, args
        # 若模型未加载，等待加载完成
        if not model_loaded:
            with load_lock:
                load_model()
        
        text = request.args.get('text')
        result = model_happy_tt.generate_text("grammar: " + text, args=args)
        return jsonify({'corrected_text': result.text})
    
    return app

app = create_app()

2. 配置模型缓存持久化 #

通过Docker卷将模型缓存目录挂载到外部存储，避免容器重启后重复下载：

修改Dockerfile

FROM python:3.10-slim
WORKDIR /app
COPY ./requirements.txt /app/requirements.txt
RUN pip install -r requirements.txt
COPY . .
# 创建模型缓存目录
RUN mkdir -p /app/model_cache
EXPOSE 5000
CMD ["gunicorn", "--bind", "0.0.0.0:5000", "wsgi:app", "--workers=1", "--timeout=120"]

启动容器时挂载卷

docker run -d -p 5000:5000 -v ./model_cache:/app/model_cache your-image-name

3. 优化Gunicorn启动参数 #

• 保持单worker模式（--workers=1），避免多worker重复触发模型下载
• 添加--timeout=120参数，延长启动超时时间，给模型下载预留足够时间

4. 可选：预下载模型到共享存储 #

若有稳定存储环境，可先在宿主机预下载模型，再挂载到容器：

# 宿主机预下载模型
mkdir -p ./model_cache
python -c "from happytransformer import HappyTextToText; HappyTextToText('T5', 'vennify/t5-base-grammar-correction')"

启动容器时挂载该目录即可直接使用已下载的模型。

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