如何在Celery多进程中共享PyTorch模型以避免重复加载？

一、单GPU场景：Worker启动时全局加载模型 #

核心思路是在Celery Worker启动阶段一次性加载模型到GPU，避免每个任务重复执行6秒的加载流程。利用Celery的worker_init信号实现初始化逻辑，配合forkserver进程池保证PyTorch CUDA上下文的安全性。

代码实现 #

import torch
from celery import Celery

# 初始化Celery实例
celery = Celery('model_tasks', broker='你的Broker地址（如Redis/RabbitMQ）')

# 全局变量存储预加载的模型与设备
model = None
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

@celery.signals.worker_init.connect
def init_worker(sender, **kwargs):
    """Worker启动时加载模型到GPU"""
    global model
    if model is None:
        # 加载模型（可替换为torch.jit.load提升加载速度）
        model = torch.load('你的模型文件路径.pt')
        model = model.to(device)
        model.eval()  # 切换至评估模式，关闭训练相关层（如Dropout）

@celery.task
def predict(data):
    """异步预测任务，直接使用预加载的模型"""
    global model, device
    # 数据预处理（根据你的业务逻辑调整）
    input_tensor = torch.tensor(data).to(device)
    
    # 推理（禁用梯度计算减少内存开销）
    with torch.no_grad():
        output = model(input_tensor)
    
    # 结果后处理并返回
    return output.cpu().numpy().tolist()

Worker启动命令 #

celery -A model_tasks worker --concurrency=1 --pool=forkserver --loglevel=info

• --concurrency=1：保证单个Worker串行处理任务，避免GPU内存冲突
• --pool=forkserver：符合PyTorch多进程最佳实践，避免CUDA上下文损坏问题

二、未来多GPU扩展方案 #

当新增GPU后，可为每个GPU分配独立的Worker进程，每个进程绑定专属GPU并预加载模型，实现多GPU并行处理（每个GPU仍串行处理任务）。

代码调整 #

import torch
from celery import Celery
import os

celery = Celery('model_tasks', broker='你的Broker地址')

model = None
device = None

@celery.signals.worker_init.connect
def init_worker(sender, **kwargs):
    """根据环境变量绑定GPU并加载模型"""
    global model, device
    # 通过环境变量指定当前Worker绑定的GPU索引
    gpu_idx = int(os.environ.get('CUDA_VISIBLE_DEVICES', 0))
    device = torch.device(f'cuda:{gpu_idx}')
    torch.cuda.set_device(device)
    
    model = torch.load('你的模型文件路径.pt')
    model = model.to(device)
    model.eval()

@celery.task
def predict(data):
    global model, device
    input_tensor = torch.tensor(data).to(device)
    with torch.no_grad():
        output = model(input_tensor)
    return output.cpu().numpy().tolist()

多GPU Worker启动命令 #

# 启动绑定GPU 0的Worker
CUDA_VISIBLE_DEVICES=0 celery -A model_tasks worker --concurrency=1 --pool=forkserver --loglevel=info --name=gpu0 --queues=gpu0

# 启动绑定GPU 1的Worker
CUDA_VISIBLE_DEVICES=1 celery -A model_tasks worker --concurrency=1 --pool=forkserver --loglevel=info --name=gpu1 --queues=gpu1

• 配合Celery队列路由，Flask提交任务时可指定对应GPU的队列，实现任务定向分配：

  # Flask中提交任务到GPU 0的队列
  predict.apply_async(args=[data], queue='gpu0')
  

三、forkserver配置验证与注意事项 #

1. 为什么选择forkserver：PyTorch的CUDA上下文在默认fork模式下易出现损坏，forkserver会先启动独立服务器进程，所有Worker从该进程fork，保证CUDA上下文的独立性与正确性，你的选择符合最佳实践。
2. 模型加载优化：
   • 用torch.jit.save将模型脚本化，torch.jit.load的加载速度远快于普通torch.load，且跨环境兼容性更好。
   • 若模型支持，可转换为半精度（model.half()），减少GPU内存占用并提升加载速度。
3. 内存验证：启动Worker后用nvidia-smi查看GPU内存占用，确认模型已预加载，任务执行时内存无大幅波动（说明无重复加载）。

四、额外效率优化建议 #

1. 批量任务处理：若业务允许，将多个小任务合并为批量任务，减少GPU数据传输开销（需注意GPU内存上限）。
2. 模型量化：使用PyTorch量化工具（torch.quantization）将模型转换为INT8，降低内存占用并提升推理速度。
3. 模型缓存：若模型文件过大，可将模型加载至内存后利用forkserver的进程继承机制共享，无需额外实现共享内存逻辑。

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