使用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的生產環境中得到驗證，展示了其在實際應用中的有效性。