2. 使用FastAPI构建高性能模型服务

任务目标

构建一个基于FastAPI的异步模型推理服务,使用一个简单的线性回归模型作为示例,支持高并发请求处理,并通过Pydantic模型实现输入数据的类型与范围验证。

环境依赖

pip install fastapi uvicorn pydantic numpy scikit-learn

完整实现代码

import asyncio
import numpy as np
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from sklearn.linear_model import LinearRegression
import joblib
from contextlib import asynccontextmanager

# ---------- 输入验证模型 ----------
class PredictRequest(BaseModel):
    features: list[float] = Field(..., min_length=2, max_length=2,
                                  description="两个特征值,范围[0,100]")
    
    @field_validator("features")
    @classmethod
    def validate_range(cls, v):
        for val in v:
            if val < 0 or val > 100:
                raise ValueError("特征值必须在0到100之间")
        return v

class PredictResponse(BaseModel):
    prediction: float

# ---------- 模型加载与生命周期 ----------
model = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global model
    # 训练一个简单的线性模型 y = 2*x1 + 3*x2 + 5
    X = np.random.rand(1000, 2) * 100
    y = 2 * X[:, 0] + 3 * X[:, 1] + 5 + np.random.randn(1000) * 0.1
    model = LinearRegression()
    model.fit(X, y)
    yield
    model = None

app = FastAPI(lifespan=lifespan)

# ---------- 异步推理接口 ----------
@app.post("/predict", response_model=PredictResponse)
async def predict(request: PredictRequest):
    """异步处理推理请求,支持并发"""
    # 模拟异步IO操作(如从数据库读取特征或写入日志)
    await asyncio.sleep(0.01)
    
    features = np.array(request.features).reshape(1, -1)
    result = model.predict(features)[0]
    return PredictResponse(prediction=float(result))

# ---------- 健康检查 ----------
@app.get("/health")
async def health():
    return {"status": "ok"}

运行方式

uvicorn main:app --host 0.0.0.0 --port 8000 --workers 4

验证方法

1. 输入验证测试:

# 合法请求
curl -X POST http://localhost:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"features": [1.5, 2.5]}'
# 预期输出:{"prediction": 13.0}(约等于2*1.5+3*2.5+5)

# 非法请求(特征数量错误)
curl -X POST http://localhost:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"features": [1.5]}'
# 预期输出:422 验证错误

# 非法请求(超出范围)
curl -X POST http://localhost:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"features": [150, 2.5]}'
# 预期输出:422 验证错误

2. 并发性能测试:

# 使用ab工具发送1000个并发请求
ab -n 1000 -c 100 -p request.json -T application/json http://localhost:8000/predict

# 观察指标:请求成功率应为100%,平均响应时间应小于50ms

3. 健康检查:

curl http://localhost:8000/health
# 预期输出:{"status":"ok"}

关键实现说明

  • 异步处理:async def + await asyncio.sleep()模拟真实推理中的IO等待,使服务在等待期间能处理其他请求,提升并发能力。
  • 输入验证:通过Pydantic的Field约束特征数量,通过field_validator检查数值范围,非法输入自动返回422错误。
  • 模型生命周期:使用lifespan上下文管理器在服务启动时训练模型,关闭时释放资源,避免每次请求重复加载。
  • 多进程部署:--workers 4启动4个进程,配合异步接口可充分利用多核CPU。

通过判定标准:上述三个验证步骤均能返回预期结果,且并发测试中无请求失败,即视为案例实现成功。