通过 Bifrost 对长时间运行的扩散任务进行异步推理
通过 Bifrost 对长时间运行的扩散任务进行异步推理
TL;DR: 通过 Bifrost 进行异步推理,让长时间运行的扩散任务可以通过 x-bf-async 头提交并轮询结果,使 SDXL 批次能够跨越 60 秒的代理超时——而此前正是这个超时导致我们的产品照片流水线频繁崩溃。
在 Photoroom 的流水线中,一个大型产品变体批次需要 70 到 110 秒才能通过 SDXL 渲染完成,而 AWS ALB 默认会关闭任何空闲超过 60 秒的连接。当我们增大批次规模以降低每张图像的 GPU 成本时,同步调用在扩散步骤完成前就开始返回 504 错误。客户端对 504 进行重试,导致同一渲染任务被重复排队,在高峰时段 GPU 负载几乎翻倍。我们将生成流量迁移到了 Maxim AI 的开源 AI 网关 Bifrost 后面,并将慢速任务切换为异步推理,这样 HTTP 连接就不再需要为整个渲染过程保持打开状态。
AI 网关中的异步推理意味着什么
AI 网关中的异步推理,允许客户端提交一个生成任务,获取一个任务 ID,然后轮询结果,而不是为了整个计算过程保持一个 HTTP 连接保持打开。Bifrost 通过 x-bf-async: true 请求头和提交时返回的 x-bf-async-id 来暴露此功能,因此一个 100 秒的扩散调用就可以与客户端和网关之间的任何代理或负载均衡器的空闲限制解耦。
这里的细微差别在于,GPU 工作本身并不会变快。改变的是连接模型。同步请求将一个 100 秒渲染的成功与否,绑定到一个跨越两个网络跳点、在 100 秒内保持健康的 TCP 连接上。异步打破了这种耦合:提交调用在毫秒级返回,而轮询调用既短小又幂等。
使用 x-bf-async 提交和轮询任务
提交请求看起来像是一个通过 OpenAI 兼容端点的正常调用,只是多了一个额外的头。Bifrost 作为 直接替代方案 运行,因此我们现有的图像客户端只在头部层发生了变化,请求体没有改动。
# 提交一个长时间运行的生成任务
curl -X POST http://localhost:8080/v1/images/generations \
-H "Content-Type: application/json" \
-H "x-bf-async: true" \
-d '{\n "model": "openai/gpt-image-1",\n "prompt": "studio product shot, white seamless background",\n "n": 8\n }'
# 响应返回:x-bf-async-id: job_8f2c...
# 使用返回的任务 ID 轮询结果
curl http://localhost:8080/v1/images/generations \
-H "x-bf-async-id: job_8f2c..."
准确说明我们测量的内容:提交调用在模型开始解码之前就返回了,因此客户端线程在远小于一秒的时间内就释放了。我们确定的轮询间隔是两秒,这保持了队列工作器的轻量化,同时又不会在完成时增加明显的尾部延迟。我们完全废弃了旧的 504 重试逻辑,因为已经没有需要保持长时间连接而可能失败的情况了。
标记和观测进行中的任务
一旦任务以分离方式运行,就需要一种方法来归属每个任务,否则直到客户投诉时,一个慢速渲染才会被发现。Bifrost 将前缀为 x-bf-dim-* 的自定义维度头转发到日志、追踪和 Prometheus 中,因此我们用创建任务的团队和实验来标记每个提交。
-H "x-bf-dim-team: catalog-enrichment" \
-H "x-bf-dim-experiment: sdxl-batch-v3" \
这些标签进入可观测性层,Bifrost 以每次请求不超过 0.1ms 的开销异步写入。我们现在为每个实验绘制完成时间图,而不是一个聚合图,正因如此我们才发现某个提示模板的速度是批次中其他模板的三倍。对于团队间的成本归属,我们将维度标签与有范围的虚拟密钥配对,这样每个业务部门针对同一个提供商池都有自己的预算。
路由在这里也很重要。该网关统一了 20+ 提供商 到一个端点后面,无论任务落在自托管的 SDXL 部署上还是托管的图像模型上,相同的异步机制都适用,因此我们可以在不重写客户端的情况下进行批次故障转移。
权衡与局限性
对于快速路径,异步是错误的选择。一个在 900ms 内渲染完成的交互式缩略图,从提交-轮询中得不到任何好处;你增加了一次往返和一个轮询循环,而任务本可以在原始连接内完成。我们仅将预期渲染时间约 30 秒以上的批次通过 x-bf-async 路由。
Bifrost 一方的实际限制在于运维层面。生产部署需要 Postgres 作为网关的后端,你需要自行托管整套系统,这意味着你实际运行和维护的是真正的基础设施,而非一个托管端点。基准测试数据表现强劲:Bifrost 在单个 t3.xlarge 实例上可维持 5,000 RPS,成功率 100%,额外开销约 11µs,但这些数据描述的是你自行运维的节点。相比 LiteLLM 这类更成熟的代理,其生态系统也较年轻,因此某些集成路径可参考的社区示例较少。对我们团队而言,这个权衡显然是值得的,因为替代方案是针对每条路由调整负载均衡器超时设置,但最终仍会在尾部丢任务。
总结
异步推理并未让我们的扩散模型跑得更快;它通过消除对单个长连接依赖,让长时间渲染变得可维持。x-bf-async 的提交并轮询模型,加上用于归因的维度标签,将一类间歇性 504 错误转化为可度量的队列,我们能够对其进行推理。如果你运行的图像或视频生成任务经常超出代理超时,我建议优先尝试这种模式。
如果你想针对自己的负载测试异步推理及网关其余功能,请预约演示:https://getmaxim.ai/bifrost/book-a-demo
延伸阅读
- Bifrost 可观测性文档 —— 了解异步写入路径与指标接收器
- Bifrost 基准测试 —— 了解额外开销与吞吐量数据
- SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis
- AWS Application Load Balancer 连接空闲超时