2. 配置模型加载与预热策略

任务说明

在LLM推理服务中，模型加载与权重初始化通常耗时数秒至数十秒。若在首个请求到达时才执行加载，该请求将承受极高的首延迟（TTFT）。本案例目标：实现服务启动时自动完成模型加载与预热，使首个业务请求到达时模型已就绪，从而消除冷启动延迟。

实现方案

基于FastAPI构建最小推理服务，通过lifespan上下文管理器在应用启动阶段执行模型加载与预热推理。预热方式为使用固定占位输入执行一次前向传播，触发CUDA kernel加载与显存分配。服务暴露/health端点供就绪检查，暴露/generate端点处理实际推理请求。

环境依赖

pip install fastapi uvicorn torch transformers

完整代码

import asyncio
import time
from contextlib import asynccontextmanager

import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_NAME = "sshleifer/tiny-"  # 轻量模型，便于演示
PREWARM_TEXT = "The quick brown fox"  # 预热输入

class GenerateRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 20

class ModelHolder:
    """持有模型与tokenizer，避免全局变量重复加载"""

    def __init__(self):
        self.tokenizer = None
        self.model = None
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

holder = ModelHolder()

def load_and_prewarm():
    """加载模型并执行一次预热推理"""
    t0 = time.perf_counter()
    holder.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    holder.model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(holder.device)
    holder.model.eval()
    load_time = time.perf_counter() - t0

    # 预热：触发CUDA kernel加载与显存分配
    t1 = time.perf_counter()
    inputs = holder.tokenizer(PREWARM_TEXT, return_tensors="pt").to(holder.device)
    with torch.no_grad():
        _ = holder.model.generate(**inputs, max_new_tokens=5)
    warm_time = time.perf_counter() - t1

    print(f"[startup] 模型加载耗时 {load_time:.2f}s, 预热耗时 {warm_time:.2f}s")

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动阶段：同步加载在事件循环外执行，避免阻塞
    await asyncio.to_thread(load_and_prewarm)
    yield
    # 关闭阶段释放显存
    del holder.model
    if holder.device == "cuda":
        torch.cuda.empty_cache()

app = FastAPI(lifespan=lifespan)

@app.get("/health")
async def health():
    if holder.model is None:
        raise HTTPException(status_code=503, detail="模型尚未就绪")
    return {"status": "ready", "device": holder.device}

@app.post("/generate")
async def generate(req: GenerateRequest):
    if holder.model is None:
        raise HTTPException(status_code=503, detail="模型尚未就绪")
    inputs = holder.tokenizer(req.prompt, return_tensors="pt").to(holder.device)
    with torch.no_grad():
        outputs = holder.model.generate(**inputs, max_new_tokens=req.max_new_tokens)
    text = holder.tokenizer.decode(outputs[0], skip_special_tokens=True)
    return {"generated_text": text}

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)

运行与验证

  1. 启动服务：python server.py，观察控制台输出[startup] 模型加载耗时...，说明加载与预热已在启动阶段完成。
  2. 就绪检查：curl http://localhost:8000/health，返回{"status":"ready"}即表示模型已就绪。
  3. 首请求延迟验证：服务启动完成后立即执行curl -X POST http://localhost:8000/generate -H "Content-Type: application/json" -d '{"prompt":"Hello world"}'，记录响应时间。由于模型已预热，首个请求无需等待加载，TTFT显著低于未预热场景。
  4. 判定通过条件：启动日志中出现加载与预热耗时输出；/health返回ready；首个/generate请求成功返回文本且无超时。

关键说明

  • 使用asyncio.to_thread将阻塞的模型加载放入线程池，避免阻塞事件循环导致启动超时。
  • 预热推理使用max_new_tokens=5的短生成，足以触发kernel初始化与显存分配，无需完整长文本。
  • 生产环境可将MODEL_NAME替换为实际模型路径，并增加device_map="auto"以支持多卡分片。
  • 若使用vLLM等推理框架，预热逻辑可替换为engine.add_request发送一次虚拟请求，原理与本例一致。
← 上一章 📖 返回目录 下一章 →