使用 pipelines 为 Web 服务器提供推理服务¶

from starlette.applications import Starlette
from starlette.responses import JSONResponse
from starlette.routing import Route
from transformers import pipeline
import asyncio

# 定义主页路由处理函数
async def homepage(request):
    # 获取请求体
    payload = await request.body()
    string = payload.decode("utf-8")
    
    # 创建一个队列用于存储推理结果
    response_q = asyncio.Queue()
    
    # 将请求字符串和队列放入模型队列中
    await request.app.model_queue.put((string, response_q))
    
    # 从队列中获取推理结果并返回
    output = await response_q.get()
    return JSONResponse(output)

# 定义服务器循环处理函数
async def server_loop(q):
    # 初始化推理 pipelines
    pipe = pipeline(model="google-bert/bert-base-uncased")
    
    while True:
        # 从队列中获取请求字符串和结果队列
        (string, response_q) = await q.get()
        
        # 执行推理并返回结果
        out = pipe(string)
        await response_q.put(out)

# 创建 Starlette 应用实例
app = Starlette(
    routes=[
        Route("/", homepage, methods=["POST"]),
    ],
)

# 定义启动事件
@app.on_event("startup")
async def startup_event():
    q = asyncio.Queue()  # 创建队列
    app.model_queue = q  # 将队列绑定到应用实例
    asyncio.create_task(server_loop(q))  # 创建推理循环任务

uvicorn server:app

curl -X POST -d "test [MASK]" http://localhost:8000/

(string, rq) = await q.get()
strings = []
queues = []

while True:
    try:
        # 尝试在1毫秒内获取队列中的请求
        (string, rq) = await asyncio.wait_for(q.get(), timeout=0.001)  # 1ms
    except asyncio.exceptions.TimeoutError:
        break
    strings.append(string)
    queues.append(rq)

# 批量推理
outs = pipe(strings, batch_size=len(strings))

# 将结果放入各自的队列中
for rq, out in zip(queues, outs):
    await rq.put(out)