Django+Celery+Redis Docker部署：模型重复初始化致GPU内存耗尽求解

我在Docker容器中部署了包含Django、Redis和Celery的项目，使用Redis和Celery管理队列及其他进程。当前遇到的问题是：Django应用启动时仅初始化一次模型，但每个Celery Worker启动时都会尝试重建模型，由于GPU内存有限，导致应用因内存不足崩溃。

当前model.py中的模型初始化代码：

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "test-model",
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
    token = TOKEN
)
FastLanguageModel.for_inference(model)

我曾尝试在Django应用中初始化模型，但每个Celery Worker启动时仍会重新加载模型，导致GPU内存耗尽。

问题根源 #

Celery Worker是独立的Python进程，每个Worker启动时都会加载model.py模块，执行模块级别的模型初始化代码。这意味着每个Worker进程都会单独加载一份模型到GPU内存中，当Worker数量较多时，GPU内存会被快速耗尽。

解决方案：模型服务化（最优方案） #

将模型单独部署为一个独立的微服务，让所有Celery Worker通过网络请求调用该服务的推理接口。这样模型仅在服务进程中初始化一次，所有Worker共享同一模型实例，彻底解决重复加载的问题。

步骤1：创建模型推理服务 #

用FastAPI编写一个简单的模型服务，在服务启动时仅初始化一次模型：

# model_service.py
from fastapi import FastAPI
from pydantic import BaseModel
# 导入你的模型相关依赖
from your_module import FastLanguageModel, max_seq_length, dtype, load_in_4bit, TOKEN

app = FastAPI()

# 全局初始化模型（仅服务启动时执行一次）
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="test-model",
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit=load_in_4bit,
    token=TOKEN
)
FastLanguageModel.for_inference(model)

# 定义请求体格式
class InferenceRequest(BaseModel):
    input_text: str

@app.post("/predict")
async def predict(request: InferenceRequest):
    # 执行模型推理
    inputs = tokenizer(request.input_text, return_tensors="pt").to("cuda")
    outputs = model.generate(**inputs)
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return {"result": result}

步骤2：修改Celery任务代码 #

将Celery任务中的模型推理逻辑替换为调用模型服务的HTTP请求：

# celery_tasks.py
import requests
from celery import shared_task

@shared_task
def process_text_task(input_text):
    # 调用模型服务接口（替换为你的模型服务地址）
    response = requests.post(
        "http://model-service:8000/predict",
        json={"input_text": input_text}
    )
    response.raise_for_status()
    return response.json()["result"]

步骤3：Docker编排部署 #

在Docker Compose中添加模型服务容器，确保其能访问GPU资源，并与其他服务网络互通：

# docker-compose.yml
version: '3.8'
services:
  django-app:
    # 你的Django应用配置
    build: ./django-app
    ...

  celery-worker:
    # 你的Celery Worker配置
    build: ./celery-worker
    depends_on:
      - redis
      - model-service
    ...

  redis:
    image: redis:latest
    ...

  model-service:
    build: ./model-service
    ports:
      - "8000:8000"
    # 配置GPU资源（需提前安装nvidia-docker）
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    environment:
      - PYTHONUNBUFFERED=1

备选方案（低负载场景） #

如果你的业务并发量较低，可以限制Celery Worker的数量为1，确保仅加载一次模型：
启动Celery时指定Worker数量：

celery -A your_project worker --loglevel=info --concurrency=1

此方案会严重限制任务处理的并发能力，仅适合测试或低流量场景。

注意事项 #

• GPU资源配置：确保Docker环境支持nvidia runtime，模型服务容器能正常访问GPU。
• 接口超时设置：模型推理可能耗时较长，需在请求时设置合理的超时时间。
• 请求缓存：对于重复的输入请求，可以添加缓存逻辑（如用Redis缓存结果），减少模型服务的计算压力。

内容的提问来源于stack exchange，提问作者Mehmet Yıldırım