使用Lakebase Postgres简化AI代理编排
本文介绍了Databricks如何利用Lakebase Postgres为AI代理构建一个可扩展、容错的任务队列，无需外部中间件。通过四个Postgres原生模式，实现了并发优先级感知的调度、基于租约的崩溃恢复、速率限制感知的节流以及幂等回调。同时，结合LISTEN/NOTIFY和SSE实现了实时操作仪表板。该架构已在CLA的审计解决方案中得到验证，将文档提取时间从数小时缩短至数分钟。
简介：传统的审计流程繁琐且耗时，需要详细审阅和提取文档信息。为了加速这一过程，CLA（CliftonLarsonAllen LLP）与Databricks团队合作，构建了一个基于AI的审计解决方案。该方案完全基于Databricks平台，利用Lakebase Postgres、Databricks Apps、Lakeflow Jobs、MLflow和Unity Catalog Volumes，将文档提取时间从数小时缩短至数分钟。本文重点介绍其中的编排层，它负责协调长时间运行的任务、管理重试、分配成本并提供实时可见性。
编排挑战：在规模化运行文档解析时，面临五个分布式系统问题：任务延迟不可预测、需感知API速率限制、工作量优先级排序、每任务成本归属以及实时进度可见性。传统做法是组合多个专用系统，但每个系统都带来额外的运维负担。Databricks的解决方案以Lakebase为核心，用一个统一平台满足所有要求。
解决方案架构：整个应用栈仅包含Databricks服务：Web应用（Databricks Apps）负责上传文档和提交请求；Lakebase作为Postgres数据库，托管编排器的关系状态；编排器（Databricks Apps）是一个长期运行的工作守护进程和操作仪表板；AI代理（Lakeflow Jobs）执行实际的文档解析。数据流简单清晰，无需外部消息代理或调度器。
任务队列实现：任务队列由Lakebase中的两个Postgres表支持。tasks表记录每个任务的逻辑单元，task_attempts表记录每次执行尝试。四个Postgres原生模式将其转变为健壮的队列：
- 并发优先级感知出队：使用“FOR UPDATE SKIP LOCKED”确保并发安全，并通过ORDER BY priority DESC, created_at实现优先级排序。
- 基于租约的崩溃恢复：出队时记录过期租约，定期清理器重新入队已过期的任务，确保崩溃后任务自动恢复。
- 速率限制感知节流：支持三种模式——并发上限、令牌预算或两者结合，在出队时动态检查，避免超出模型端点限制。
- 幂等Webhook回调：回调处理器设计为幂等，接受多种状态，避免重复处理。
实时操作仪表板：利用Postgres的LISTEN/NOTIFY机制和服务器发送事件（SSE），构建低延迟操作仪表板，实时显示任务状态、输入输出令牌数、LLM成本、计算成本和中位响应时间等指标，无需额外监控平台。
结论：通过Lakebase Postgres，Databricks提供了一种更简单、可扩展的编排模式，用于长时间运行的代理工作负载。该架构已在CLA的生产环境中得到验证，展示了其在实际应用中的有效性。