Transformers Pipeline Web 推理服务:用 Starlette + asyncio 队列构建单线程模型推理服务器

本文以 Transformers 官方文档中的"Web 服务器推理"(Web server inference)指南为主体,完整讲解如何把 pipeline 作为推理引擎嵌入 Web 服务器:从完整的 Starlette/uvicorn 服务器代码,到单模型加载、请求队列、动态批处理伪代码,再到错误处理、熔断、主线程阻塞等生产化要点。读完后,你能够独立搭建一个"Web 层轻量收发 + 单线程模型计算"的最小推理服务,并理解其背后的并发设计取舍。

一、先明确前提:不存在放之四海皆准的推理引擎

构建一个推理引擎是一项复杂的工作,"最佳"方案取决于你的问题领域:你用 CPU 还是 GPU?你要的是最低延迟、最高吞吐、支持多模型,还是把某个特定模型调到极致?有很多方法可以解决这些问题,因此文档给出的只是一个良好的默认起点,而不一定是针对你的场景的最优解。

这是原文档开篇给出的重要立场:下面所有内容都是一个可运行的"默认起点",而非性能标杆。理解了这一点,才不会把示例代码直接当生产系统用。

二、核心矛盾:Web 服务器的多路复用 vs Pipeline 的串行计算

要理解整套设计,只需抓住两点事实:

1. Web 服务器本质上是"等待请求、到来即处理"的系统,因此可以像在数据集上迭代那样使用 Pipeline 的迭代器能力,逐个处理到达的请求;
2. Web 服务器通常多路复用(多线程、异步等)以并发处理请求,而 Pipeline(及其底层模型)并不适合并行:它们占用大量 RAM,最合适的用法是在运行时把全部可用资源交给它,尤其是在计算密集型任务中。

文档给出的解法是:让 Web 服务器承担"接收请求、发送响应"的轻量负载,只用一个线程/任务去承担真正执行 Pipeline 的重负载。示例基于 starlette 实现;如果换用其他框架,需要调整代码才能达到同样效果。

三、创建服务器:完整可运行的 server.py

示例基于轻量级框架 Starlette 和 ASGI 服务器 uvicorn 构建。先安装依赖:

pip install starlette uvicorn

然后创建 server.py,完整代码如下(与原文档一致,此处补充了关键注释):

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):
    # 1. 读取并解码 HTTP 请求体
    payload = await request.body()
    string = payload.decode("utf-8")
    # 2. 为当前请求创建一个"响应队列",与输入一起放入模型队列
    response_q = asyncio.Queue()
    await request.app.model_queue.put((string, response_q))
    # 3. 挂起当前协程,等待模型线程回填结果
    output = await response_q.get()
    return JSONResponse(output)


async def server_loop(q):
    # 关键:模型只在这里加载一次
    pipe = pipeline(model="google-bert/bert-base-uncased")
    while True:
        (string, response_q) = await q.get()
        out = pipe(string)
        await response_q.put(out)


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))

代码中的几个设计要点:

• pipeline(model="google-bert/bert-base-uncased") 在 server_loop 中只执行一次。FillMaskPipeline 类定义在 fill_mask.py,当直接传入 model 而不指定 task 时,pipeline 工厂函数会根据模型配置自动推断任务类型——对 google-bert/bert-base-uncased 这类掩码语言模型,BERT 的默认任务正是 fill-mask(见 fill_mask.py 中的官方示例),因此无需显式写 task="fill-mask";
• 每个请求自带一个 response_q:(string, response_q) 这个元组进入全局队列后,Web 侧协程在 await response_q.get() 处挂起,模型侧计算完成后把结果放回该私有队列,实现了"一请求一通道"的回传,无需额外维护请求到响应的映射表;
• app.model_queue 挂在 app 对象上,使请求处理函数与启动钩子之间共享同一个队列,避免全局变量。

运行与查询

启动服务器:

uvicorn server:app

用 POST 请求查询:

curl -X POST -d "test [MASK]" http://localhost:8000/
#[{"score":0.7742936015129089,"token":1012,"token_str":".","sequence":"test."},...]

返回的是 fill-mask 任务的标准输出:一组按分数排序的候选填充词,每项包含 score、token、token_str、sequence 四个字段。默认最多返回 5 个预测(top_k 默认值为 5,见 fill_mask.py 的参数说明)。

四、源码级原理:为什么这套结构成立

4.1 pipeline 工厂与 Pipeline 基类

从源码结构看,pipeline 工厂函数位于 pipelines/init.py,函数签名支持 model、config、tokenizer、revision、use_fast、device、device_map、dtype(默认 "auto")、model_kwargs、pipeline_class 等参数,并针对 fill-mask、text-generation、zero-shot-classification 等全部任务提供了独立的类型重载(见 第 387 行起的 overload 定义)。这意味着示例中的调用返回的是一个强类型的 FillMaskPipeline 实例,其内部完成模型加载、tokenizer 装配与推理封装。

4.2 模型只需加载一次:内存是核心约束

原文档强调的"加载模型一次"直接对应 RAM 约束:Pipeline 基类(定义于 base.py)持有模型与 tokenizer 的引用,若在 Web 服务器每来一个请求就新建一个 pipeline,就会产生多份模型副本,白白消耗显存/内存。本例中模型对象只存在于 server_loop 这一个任务内,Web 层(Starlette 的各请求协程)完全不接触模型对象,这就是"轻量收发 + 单线程重计算"架构的内存收益来源。

4.3 __call__ 的批处理能力是队列机制的伏笔

Pipeline 基类的调用入口签名为 __call__(self, inputs, *args, num_workers=None, batch_size=None, **kwargs)(base.py),原生支持传入批量输入并通过 batch_size 控制分批推理。示例中 pipe(string) 传单条字符串,而下文动态批处理示例中的 pipe(strings, batch_size=len(strings)) 也是同一入口——这正是"队列机制允许做动态批处理这类高级操作"的底层依据。

五、请求队列的进阶用法:动态批处理

⚠️ 下面的代码是刻意写成便于阅读的伪代码。在未确认它符合你的系统资源情况之前,不要直接运行!

队列机制天然支持"攒一批再推理"的动态批处理:先累积若干请求,再用 Pipeline 一次性推理。

(string, rq) = await q.get()
strings = []
queues = []
while True:
    try:
        (string, rq) = await asyncio.wait_for(q.get(), timeout=0.001) # 1ms
    except asyncio.exceptions.TimeoutError:
        break
    strings.append(string)
    queues.append(rq)
strings
outs = pipe(strings, batch_size=len(strings))
for rq, out in zip(queues, outs):
    await rq.put(out)

原文档同时明确指出了这段示例代码的三处已知缺陷,这是理解"何时该用动态批处理"的关键:

1. 没有批大小上限。无上限的批次通常不是好主意——批次越大,单批推理越久,内存压力越大;
2. 超时在每次出队时都会被重置。这意味着你可能等待远超 1ms 才真正开始处理请求,并且第一个推理请求也会被推迟同样的时间。更合理的做法是设置一个固定的 1ms 截止期,而不是每次出队都重新计时;
3. 即使队列为空也固定等待 1ms。由于你大概率希望队列一有东西就开始推理,这种"永远等 1ms"的做法未必最优;但如果动态批处理对你的场景至关重要,这种取舍又可能是合理的。原文档再次强调:不存在单一的最佳方案。

六、生产化之前,原文档列出的四个重点考量

以下四节完整继承自原文档"你可能需要考虑的事情"部分,并补充了与源码结构的对应关系。

6.1 错误检查(Error checking)

把模型公开对外服务时,有大量事情可能出错:内存溢出、磁盘空间不足、模型加载失败、查询本身有误、查询合法但因模型配置错误而运行失败……

一般做法是:让服务器把错误输出给用户,因此加入大量 try...except 来呈现这些错误是合理的。但要意识到:根据你的安全上下文,向用户暴露全部错误本身可能构成安全风险(例如泄露堆栈、内部路径或模型配置细节)。

6.2 熔断(Circuit breaking)

Web 服务器在过载时"切断电路"通常表现更好:即过载时返回正确的错误,而不是让用户无限等待——与其长时间等待,不如返回 503;与其等很久,不如在超长时间后返回 504。

在本例代码中实现这一点相对简单,因为只有一个队列:观察队列长度是开始返回错误、避免 Web 服务器在过载下崩溃的基本手段。可以据此设定阈值,例如队列长度超过 N 时直接对后续请求返回 503。

6.3 主线程阻塞(Block the main thread)

PyTorch 目前不支持异步感知(async-aware)的运算:一旦开始计算,就会阻塞主线程。也就是说,让 PyTorch 跑在独立的线程/进程中会更好。本例没有这么做,因为代码会复杂得多——主要原因在线程/进程与 asyncio 队列并不天然兼容;但从最终效果看,它会起到相同的作用。

这一点在**单次推理耗时较长(> 1 秒)**时尤为关键:因为这意味着推理期间到来的每个查询都至少要等 1 秒才能收到错误响应(对应 6.2 中熔断生效的下界延迟)。

6.4 动态批处理(Dynamic batching)

作为一般性结论:动态批处理并不必然优于每次只传一个元素(批推理的更多细节可参见 pipelines 文档中的批处理章节)。但在正确的场景下它可能非常有效。文档给出的两个边界案例值得记住:

• 在 API 服务中,默认不做动态批处理(那样很容易出现明显变慢的情况);
• 但对于 BLOOM 这类超大型模型的推理,动态批处理是必要的,才能给所有用户提供良好的体验。

七、适用前提与限制小结

要点	说明
依赖	starlette、uvicorn;示例默认加载 google-bert/bert-base-uncased(需可访问模型仓库)
架构	Web 层只做请求收发;模型在唯一任务 server_loop 中加载一次并串行处理
回传机制	每请求一个私有 asyncio.Queue,通过 app.model_queue 共享全局队列
默认任务	未显式指定 task 时由模型配置推断,BERT 系列默认 fill-mask(fill_mask.py)
动态批处理	示例为可读性优先的伪代码:无批次上限、超时会重置、空队列仍等待 1ms,生产前必须修正
并发限制	PyTorch 非 async-aware,长推理(>1s)会阻塞主线程,建议将模型计算移到独立线程/进程
过载策略	应基于队列长度实现 503/504 熔断,而非让请求无限排队

这套"轻量 Web 层 + 单队列 + 单模型工作循环"的结构,是 Transformers 官方推荐的 Web 推理默认起点:它用最小的并发复杂度换取了内存安全与行为可预测,并在此之上留出了动态批处理、熔断等进阶优化的挂载点。