140_异步推理：队列管理框架 - 使用Celery处理高并发请求的独特设计

┌─────────────────┐     ┌─────────────────┐     ┌─────────────────┐
│   API Gateway   │────▶│   Message Queue │────▶│  Worker Nodes   │
│                 │     │                 │     │                 │
│   处理请求      │◀────│   存储任务      │◀────│   执行推理      │
└─────────────────┘     └─────────────────┘     └─────────────────┘
        ▲                      ▲                      │
        │                      │                      ▼
        │                      │             ┌─────────────────┐
        │                      └─────────────▶│   结果存储      │
        │                                    │  (Redis/MongoDB)│
        └────────────────────────────────────┘                 ┘

from celery import Celery
import time
import uuid

# Celery实例初始化
celery_app = Celery('llm_inference', 
                   broker='redis://localhost:6379/0',
                   backend='redis://localhost:6379/1')

# 配置任务序列化和压缩
celery_app.conf.update(
    task_serializer='pickle',  # 支持复杂对象序列化
    result_serializer='pickle',
    accept_content=['pickle', 'json'],
    result_compression='gzip',  # 压缩结果减少网络传输
    task_compression='gzip',    # 压缩任务减少队列占用
    result_expires=3600,        # 结果过期时间（秒）
    timezone='UTC',
    enable_utc=True,
    worker_prefetch_multiplier=1  # 预取任务数量，LLM场景建议设为1
)

@celery_app.task(bind=True, autoretry_for=(Exception,), retry_backoff=2, retry_kwargs={'max_retries': 3})
def llm_inference_task(self, model_id, prompt, parameters=None, user_id=None):
    """LLM推理任务定义
    
    Args:
        model_id: 模型标识符
        prompt: 提示文本
        parameters: 推理参数（温度、最大长度等）
        user_id: 用户标识，用于跟踪和统计
    
    Returns:
        dict: 包含生成文本、推理时间等信息的结果
    """
    # 生成唯一的推理ID
    inference_id = str(uuid.uuid4())
    
    # 记录任务开始时间
    start_time = time.time()
    
    try:
        # 这里将实现实际的LLM推理逻辑
        # 例如加载模型、执行推理、后处理等
        result = {
            'inference_id': inference_id,
            'model_id': model_id,
            'generated_text': "This is a sample generated text.",
            'input_tokens': len(prompt.split()),
            'output_tokens': 100,
            'execution_time': time.time() - start_time,
            'timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
        }
        
        return result
        
    except Exception as e:
        # 记录错误日志
        error_info = f"Inference error: {str(e)}"
        self.update_state(
            state='FAILURE',
            meta={'error': error_info, 'inference_id': inference_id}
        )
        raise e

celery -A llm_worker worker --loglevel=info --concurrency=2 --max-tasks-per-child=10

# 使用CUDA_VISIBLE_DEVICES环境变量控制GPU可见性
# 或在代码中使用以下方式
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# 模型缓存装饰器
def model_cache(max_models=3):
    cache = {}
    
    def decorator(func):
        def wrapper(model_id, *args, **kwargs):
            if model_id not in cache:
                # 如果缓存已满，删除最早的模型
                if len(cache) >= max_models:
                    oldest_key = next(iter(cache))
                    del cache[oldest_key]
                # 加载新模型
                model = load_model(model_id)
                cache[model_id] = model
            return func(cache[model_id], *args, **kwargs)
        return wrapper
    return decorator

# Celery配置 - 任务路由
celery_app.conf.update(
    task_routes={
        # 高优先级任务路由到专用队列
        'llm_inference.high_priority_task': {
            'queue': 'high_priority',
            'routing_key': 'high.priority'
        },
        # 长文本生成任务路由到高性能GPU节点
        'llm_inference.long_generation_task': {
            'queue': 'gpu_high_memory',
            'routing_key': 'gpu.high_memory'
        },
        # 默认路由配置
        'llm_inference.*': {
            'queue': 'default',
            'routing_key': 'default'
        }
    },
    task_queue_max_priority=10  # 支持优先级队列，0-10，数字越大优先级越高
)

# 定义不同优先级的任务
@celery_app.task(bind=True, queue='high_priority', priority=9)
def high_priority_task(self, *args, **kwargs):
    # 高优先级任务实现
    pass

@celery_app.task(bind=True, queue='gpu_high_memory', priority=5)
def long_generation_task(self, *args, **kwargs):
    # 长文本生成任务实现
    pass

# 任务前置处理 - 记录任务开始
@celery_app.task_prerun.connect
def task_prerun_handler(task_id, task, args, kwargs, **extra):
    # 记录任务开始信息到日志或监控系统
    print(f"Task {task_id} ({task.name}) started")
    # 可以在这里更新任务状态到监控系统

# 任务成功处理 - 记录结果
@celery_app.task_success.connect
def task_success_handler(sender=None, result=None, **kwargs):
    # 记录成功结果，更新统计信息
    print(f"Task {sender.request.id} succeeded, result length: {len(str(result))}")

# 任务失败处理 - 错误记录与通知
@celery_app.task_failure.connect
def task_failure_handler(task_id, exception, traceback, sender=None, **kwargs):
    # 记录错误信息
    error_info = {
        'task_id': task_id,
        'exception': str(exception),
        'task_name': sender.name if sender else 'unknown'
    }
    print(f"Task failed: {error_info}")
    # 可以在这里发送告警通知
    # send_alert(f"LLM inference task failed", error_info)

@celery_app.task(bind=True, autoretry_for=(TemporaryError,), 
                retry_backoff=True, retry_backoff_max=60, 
                retry_jitter=True, retry_kwargs={'max_retries': 5})

from fastapi import FastAPI, BackgroundTasks, HTTPException, status
from pydantic import BaseModel
import asyncio

app = FastAPI()

class InferenceRequest(BaseModel):
    model_id: str
    prompt: str
    parameters: dict = {}

@app.post("/api/v1/inference/async")
async def async_inference(request: InferenceRequest):
    """异步推理API，返回任务ID"""
    # 提交任务到Celery队列
    task = llm_inference_task.delay(
        model_id=request.model_id,
        prompt=request.prompt,
        parameters=request.parameters
    )
    
    return {
        "task_id": task.id,
        "status": "pending",
        "message": "Inference task submitted successfully",
        "eta": "Check status at /api/v1/inference/status/{task.id}"
    }

@app.get("/api/v1/inference/status/{task_id}")
async def get_task_status(task_id: str):
    """获取任务状态和结果"""
    task = llm_inference_task.AsyncResult(task_id)
    
    if task.state == 'PENDING':
        return {
            "task_id": task_id,
            "status": "pending",
            "progress": 0,
            "message": "Task is waiting to be processed"
        }
    elif task.state == 'PROGRESS':
        return {
            "task_id": task_id,
            "status": "processing",
            "progress": task.info.get('progress', 0),
            "message": "Task is being processed"
        }
    elif task.state == 'SUCCESS':
        return {
            "task_id": task_id,
            "status": "completed",
            "progress": 100,
            "result": task.result,
            "message": "Inference completed successfully"
        }
    elif task.state == 'FAILURE':
        return {
            "task_id": task_id,
            "status": "failed",
            "progress": 0,
            "error": str(task.result),
            "message": "Inference failed"
        }
    else:
        return {
            "task_id": task_id,
            "status": task.state,
            "message": f"Task is in state: {task.state}"
        }

@celery_app.task(bind=True)
def llm_inference_with_progress(self, model_id, prompt, parameters=None):
    """带进度反馈的LLM推理任务"""
    # 初始化进度
    self.update_state(state='PROGRESS', meta={'progress': 0})
    
    # 步骤1: 加载模型 (20%)
    time.sleep(1)  # 模拟模型加载
    self.update_state(state='PROGRESS', meta={'progress': 20})
    
    # 步骤2: 处理输入 (30%)
    time.sleep(0.5)  # 模拟输入处理
    self.update_state(state='PROGRESS', meta={'progress': 30})
    
    # 步骤3: 执行推理 (80%)
    time.sleep(2)  # 模拟推理过程
    self.update_state(state='PROGRESS', meta={'progress': 80})
    
    # 步骤4: 后处理结果 (90%)
    time.sleep(0.5)  # 模拟结果后处理
    self.update_state(state='PROGRESS', meta={'progress': 90})
    
    # 步骤5: 存储结果 (100%)
    time.sleep(0.3)  # 模拟结果存储
    
    return {
        'generated_text': "Final generated text with progress tracking",
        'execution_time': 4.3,
        'status': 'completed'
    }