解决Celery任务中加载joblib保存的wav2vec2模型挂起问题

可能的原因 #

• 模型序列化兼容性问题：joblib对包含PyTorch张量的复杂模型对象支持有限，Celery Worker加载时可能触发隐式死锁或资源加载异常，导致任务无响应。
• Celery进程模型冲突：默认的prefork进程模式下，Worker的进程隔离机制可能与模型的CUDA上下文、多线程推理设置产生冲突，引发挂起。
• Pipeline懒加载特性影响：transformers的Pipeline默认懒加载模型资源，在Celery任务中首次加载时，资源竞争可能导致任务停滞。

解决方法 #

1. 放弃joblib序列化，直接在Worker中初始化模型（推荐） #

避免序列化复杂模型对象，改为在Worker启动时或任务内直接加载模型：

from celery import Celery
from transformers import pipeline

app = Celery('asr_tasks', broker='你的消息队列地址')

# 全局模型实例，Worker启动时预加载一次
asr_pipeline = None

@app.on_worker_init.connect
def init_model(**kwargs):
    global asr_pipeline
    asr_pipeline = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-base-960h")

@app.task
def transcribe_audio(audio_path):
    return asr_pipeline(audio_path)

2. 调整Celery Worker运行参数 #

修改Worker的进程模式，避免资源冲突：

• 使用单进程模式启动Worker：

celery -A 你的应用模块 worker --pool=solo --loglevel=info

• 强制模型使用CPU推理并关闭多线程，减少资源竞争：

# 初始化Pipeline时指定参数
pipeline("automatic-speech-recognition", model="facebook/wav2vec2-base-960h", device=-1, model_kwargs={"num_workers": 1})

3. 修正joblib序列化逻辑（不推荐，但可尝试） #

序列化前将模型移至CPU，避免CUDA张量的序列化问题：

# 保存模型时
asr_pipeline = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-base-960h")
asr_pipeline.model = asr_pipeline.model.to("cpu")
joblib.dump(asr_pipeline, "asr_model.joblib")

# Celery任务中加载
import joblib
@app.task
def transcribe_audio(audio_path):
    asr_pipeline = joblib.load("asr_model.joblib")
    return asr_pipeline(audio_path)

调试技巧 #

• 开启Celery Debug日志，查看Worker的资源加载细节：

celery -A 你的应用模块 worker --loglevel=debug

• 在任务中添加打印语句，定位挂起发生在模型加载阶段还是推理阶段。

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