Flask跨请求复用GPU内模型避免重复初始化的解决方案

报错核心原因是：CUDA上下文与进程强绑定，fork生成的子进程无法继承父进程已初始化的CUDA状态，任何跨进程传递已加载CUDA模型对象的操作都会触发这个错误，不管是用标准库BaseManager还是其他多进程通信工具都走不通。

以下方案按实现复杂度从低到高排序，都可以实现模型单次加载、跨请求复用，无需每次请求重新初始化：

方案1：单Flask进程全局加载（90%场景适用，实现成本最低） #

完全不需要引入多进程管理器，直接在Flask服务启动阶段全局加载一次模型，所有请求复用同一个全局模型对象即可。

实现代码 #

from flask import Flask, request
from transformers import pipeline
import torch

app = Flask(__name__)
generator = None

@app.before_first_request
def init_model():
    global generator
    # 模型仅在服务首次启动时加载一次，常驻GPU显存
    generator = pipeline(
        'text-generation',
        model=YOUR_MODEL_NAME,
        device=1
    )

@app.route("/generation", methods=["POST"])
def text_gen():
    params = request.get_json()
    # 直接调用已加载的全局模型，无重复初始化开销
    output = generator(
        params["prompt"],
        max_new_tokens=params.get("max_new_tokens", 100)
    )
    return {"code": 0, "data": output}

if __name__ == "__main__":
    # 关键配置：单进程+多线程模式，避免多worker重复加载模型
    app.run(
        host="0.0.0.0",
        port=5000,
        threaded=True,
        processes=1
    )

注意事项 #

• 必须将processes参数设为1，否则Flask会启动多个worker进程，每个进程都会独立加载一份模型到GPU，不仅会重复消耗60秒加载时间，还会直接占满显存。
• 开启threaded=True是安全的，Transformers库的pipeline内部自带推理锁，不会出现多线程同时调用CUDA的冲突问题，常规QPS场景完全够用。

方案2：独立模型Worker架构（高QPS场景适用） #

如果单进程Flask的性能扛不住并发量，就把模型加载逻辑和Web请求逻辑拆成两个独立部分：

1. 单独启动一个模型Worker进程，进程启动时仅加载一次模型到GPU，常驻运行，专门处理推理任务
2. Flask Web服务只负责处理HTTP请求，通过进程队列把推理任务发给模型Worker，拿到结果后返回给用户

核心实现逻辑 #

import uuid
import torch
from flask import Flask, request
from transformers import pipeline
from torch.multiprocessing import Process, Queue, set_start_method

# 模型常驻进程逻辑
def model_worker(task_queue: Queue, result_queue: Queue, model_name: str):
    # 关键：模型必须在Worker进程内部初始化，不能从父进程传
    local_generator = pipeline('text-generation', model=model_name, device=1)
    while True:
        task_id, prompt, gen_params = task_queue.get()
        if prompt == "__EXIT__":
            break
        res = local_generator(prompt, **gen_params)
        result_queue.put((task_id, res))

def create_app(task_q: Queue, result_q: Queue):
    app = Flask(__name__)
    app.config["TASK_Q"] = task_q
    app.config["RESULT_Q"] = result_q

    @app.route("/generation", methods=["POST"])
    def text_gen():
        params = request.get_json()
        task_id = str(uuid.uuid4())
        # 把推理任务丢给模型进程
        task_q.put((task_id, params["prompt"], params.get("gen_params", {})))
        # 轮询拿对应任务的结果
        while True:
            tid, res = result_q.get()
            if tid == task_id:
                return {"code":0, "data":res}
    return app

if __name__ == "__main__":
    # 必须用spawn启动方式，避免fork导致的CUDA上下文问题
    set_start_method("spawn", force=True)
    task_q = Queue()
    result_q = Queue()
    # 启动模型Worker，仅加载一次模型
    worker = Process(
        target=model_worker,
        args=(task_q, result_q, YOUR_MODEL_NAME)
    )
    worker.start()
    # 启动Flask服务
    app = create_app(task_q, result_q)
    app.run(host="0.0.0.0", port=5000, threaded=True)

注意事项 #

• 不要在父进程初始化CUDA模型再传给子进程，所有CUDA相关的初始化必须放在Worker进程的执行逻辑里。
• 如果服务规模更大，可以把模型Worker换成独立部署的推理服务（比如单独用FastAPI写一个推理接口），Flask通过内网HTTP调用推理服务，本质逻辑和队列方案一致，都是保证模型只在一个独立进程内初始化一次，不会重复加载。

避坑总结 #

• 不用尝试用BaseManager或者其他多进程共享工具传递已初始化的CUDA模型，CUDA本身的进程隔离机制决定了这种操作不可能成功。
• torch.multiprocessing没有提供BaseManager是正常的，就算自行封装了Manager能力，也解决不了CUDA上下文不能跨进程传递的问题。
• 不要用Flask自带的多进程模式部署加载了大模型的服务，多Worker会导致多份模型重复加载到GPU，浪费显存和加载时间。

内容的提问来源于stack exchange，提问作者red-devil