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