Flask多进程应用中HuggingFace模型无返回结果的解决方法

当Flask以多进程模式（processes=4）运行时，主进程初始化的Encoder实例（包含预加载的SentenceTransformer模型）会被所有子进程继承。而基于PyTorch的SentenceTransformer模型在多进程环境下，尤其是通过fork创建子进程时，会出现CUDA上下文冲突（GPU场景）或内存资源竞争（CPU场景），导致模型推理卡住无响应。单进程模式下没有进程间资源共享问题，因此能正常运行。

核心思路是避免在主进程预加载模型，让每个子进程独立初始化模型，具体有两种实现方式：

方式一：模型懒加载（推荐） #

修改Encoder类，在第一次调用encode方法时才初始化模型，确保每个子进程都拥有独立的模型实例：

from flask import Flask, request
import json
from sentence_transformers import SentenceTransformer

app = Flask(__name__)

class Encoder():
    def __init__(self):
        self.model = None  # 先不初始化模型

    def encode(self, text):
        # 第一次调用时初始化模型
        if self.model is None:
            self.model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
        # 将numpy数组转为列表，方便JSON序列化
        return self.model.encode(text).tolist()

encoder = Encoder()

@app.route("/get_encoded_vector", methods=['POST'])
def get_encoded_vector():
    data = json.loads(request.data)
    text = data['text']
    embeddings = encoder.encode(text)
    return json.dumps({'embeddings': embeddings})

if __name__ == '__main__':
    app.run(port=5000, threaded=False, processes=4)

方式二：子进程启动时初始化模型 #

利用Flask的请求钩子，在每个子进程处理第一个请求前初始化模型（适配Flask 2.0+版本，替代已弃用的before_first_request）：

from flask import Flask, request
import json
from sentence_transformers import SentenceTransformer

app = Flask(__name__)
encoder = None  # 全局变量存储模型实例

@app.before_request
def init_model():
    global encoder
    # 每个子进程仅初始化一次模型
    if encoder is None:
        encoder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

@app.route("/get_encoded_vector", methods=['POST'])
def get_encoded_vector():
    data = json.loads(request.data)
    text = data['text']
    embeddings = encoder.encode(text).tolist()
    return json.dumps({'embeddings': embeddings})

if __name__ == '__main__':
    app.run(port=5000, threaded=False, processes=4)

• JSON序列化时，将numpy数组转为tolist()而非str()，避免前端解析字符串格式的数组，提升数据交互效率。
• 如果使用GPU，可在模型初始化前添加import torch; torch.cuda.set_device(0)（指定GPU编号），避免多进程抢占GPU资源。

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