Lakebase PostgresでAIエージェントのオーケストレーションを簡素化

はじめに

従来の監査は、詳細な文書のレビューや情報抽出を必要とすることが多く、時間と手間がかかるプロセスです。このプロセスを迅速化するため、グローバルに展開する大手プロフェッショナルサービス企業であるCLA（CliftonLarsonAllen LLP）は、DatabricksのForward Deployed Engineeringチームと協力し、エージェント型監査ソリューションを構築して本番環境に導入しました。共同で開発した文書処理アプリケーションにより、品質を損なうことなく、抽出時間を数時間から数分へと短縮しました。このアプリケーションは完全にDatabricks上で構築されており、Lakebase Postgres、Databricks Apps、Lakeflow Jobs、MLflow、およびUnity Catalog Volumesを使用しています。このブログでは、そのシステムの主要コンポーネントの1つである、Lakebaseを活用したオーケストレーションレイヤーに焦点を当てます。

オーケストレーションレイヤーは、長時間実行されるタスクの調整、再試行の管理、コストの帰属、およびリアルタイムの可視性の提供を担当します。LakebaseとDatabricks Appsを使用することで、キューイング、オーケストレーション、オブザーバビリティのための個別のインフラストラクチャが不要になります。

また、Lakebaseはストレージとコンピューティングを分離することで、このアーキテクチャを大規模環境でも実用的なものにします。従来のPostgresデプロイとは異なり、ストレージの耐久性と独立性を維持しながら、需要に応じてコンピューティングを拡張できます。これらの機能が組み合わさることで、Lakebaseは、Databricks上で長時間実行されるエージェント型ワークロード向けの、よりシンプルでスケーラブルなオーケストレーションパターンの実用的な基盤となります。

エージェント型ワークロードにおけるオーケストレーションの課題

文書のパース（解析）は、非常に一般的で大量の処理を伴うエージェント型ワークロードです。あらゆる業界の企業が、大量の契約書、請求書、財務書類、その他の文書を構造化データに変換する必要があります。これを大規模に実行すると、分散システムにおける5つの明確な課題が浮き彫りになります。

• タスクごとの予測不可能なレイテンシ： 2ページの請求書は数秒で処理できる一方、200ページの契約書には数分かかる場合があり、個々のタスクの実行時間を予測することが困難になります。
• レート制限を考慮したスロットリング： LLMやビジョンモデルのエンドポイントは、一定期間内に処理できるリクエスト数やトークン数を制限しています。一度に数百のタスクを送信すると、これらの制限を超えてスロットリングが発生し、再試行が繰り返される原因になります。オーケレーターは、リアクティブな再試行だけに頼るのではなく、処理中のワークロード（同時タスク数、トークン予算、またはその両方）をプロアクティブに制限する必要があります。
• ワークロードの優先順位付け： 緊急の送信が、大量のバッチ処理の後回しになって遅延してはなりません。タスクごとの優先度を設定することで、優先度の高い処理（インタラクティブな送信、プレミアム層のリクエスト、オペレーターが開始した再処理など）が最初にディスパッチされるようになります。
• タスクごとのコスト帰属： 財務チームは、AIのトークン使用量やコンピューティング消費量に基づいて、特定のタスク、顧客、エージェントに支出を帰属させる必要があります。
• リアルタイムの進捗可視化： 数百の文書をアップロードするユーザーには、リアルタイムの進捗状況ビューが必要です。

多くの組織では、オーケストレーションとオブザーバビリティのために、複数の専用システムを組み合わせて使用しています。それぞれのシステムには、独自のインフラストラクチャ、認証、監視、運用の要件があり、それらを統合するための作業も発生します。長時間実行される独立したエージェント型タスクにとって、そのオーバーヘッドは実際のスケジューリングの複雑さに見合わないものになります。

私たちが開発したDatabricksネイティブのソリューションは、Lakebaseを基盤として、上記の要件をすべて満たしています。

ソリューションアーキテクチャ

アプリケーションスタック全体は、Databricksのサービスのみで構成されています。

• Webアプリケーション（Databricks Apps）： ユーザーがPDF（Unity Catalog Volumesに保存）をアップロードし、パースリクエストを送信するFastAPIベースのユーザーインターフェースです。リクエストはLakebaseのタスクテーブルに直接書き込まれます。
• Lakebase： 関連するテーブル全体でオーケレーターのリレーショナル状態をホストする、自動スケーリング対応のPostgresデータベースです。テーブルには、tasks（パース対象の文書、ステータス、リース情報、構造化された結果を保持）と、task_attempts（実行試行ごとに1行、Databricks Jobの実行ID、MLflowトレースID、試行ごとのコストメタデータを記録）があります。Lakebaseは、オーケレーターの状態に関する唯一の信頼できる情報源（Single Source of Truth）として機能します。
• オーケレーター（Databricks Apps）： 長時間実行されるワーカーデーモンとオペレーターダッシュボードです。デーモンはLakebaseからタスクをデキュー（取り出し）し、AI Agentsレイヤーにディスパッチして、結果を書き戻します。ダッシュボードは同じテーブルを読み込み、リアルタイムのステータスを表示します。
• AI Agents（Lakeflow Jobs）： Lakeflow Jobsがパース処理を実行します。各ジョブはUnity Catalog VolumesからPDFを読み込み、Intelligent Document Processingやビジョン/LLMの呼び出しを通じて処理を行い、パースされた出力をLakebaseに保存し、Webhookを介してオーケレーターに確認を通知します。MLflow Tracingは、モデルの呼び出し、トークン使用量、レイテンシ、コストメタデータなどの実行詳細をキャプチャします。

コンポーネント間のデータフローは以下の通りです。Web AppはPDFをUnity Catalog Volumesに書き込み、パースリクエストをLakebaseに書き込みます。OrchestratorはLakebaseからタスクをデキューし、Databricks JobsをAI Agentsレイヤーにディスパッチします。AI Agentsは文書を処理し、結果をLakebaseに書き戻し、ステータスの更新をOrchestratorにコールバックします。

Databricksに組み込まれたこれらの機能のおかげで、外部のメッセージブローカー（Kafka、Redis）、個別のスケジューラー（Airflow、Temporal）、または専用のキャッシュレイヤーに依存する必要はありませんでした。

タスクキューの実装

タスクキューは、Lakebase内の2つのPostgresテーブルによって支えられています。tasksテーブルは、論理的な作業単位ごとに1行を保持し、タスクの現在のステータス、リース情報、エージェントの割り当て、親の抽出、および最終結果を記録します。task_attemptsテーブルは、実行試行ごとに1行を保持し、Databricks Jobの実行ID、MLflowトレースID、および試行ごとのコストメタデータをキャプチャします。この親子関係により、再試行（単一のタスクに複数の試行が存在可能）がサポートされ、コスト帰属やデバッグのための試行レベルのオブザーバビリティが維持されます。

一対のPostgresテーブルだけでは、まだタスクキューとしては機能しません。4つのPostgresネイティブなパターンを適用することで、これらを、長時間実行されるエージェント型ワークロードに適した、堅牢で並行処理が可能、クラッシュ耐性があり、レート制限を考慮したキューへと変換します。

並行処理と優先度を考慮したデキュー

基本的なデキュークエリでは、WHERE status = 'enqueued' and LIMIT batch_sizeを使用して、次に対応可能なタスクを選択するかもしれません。このクエリはキューに入れられたタスクを正しく識別しますが、複数のワーカーが同時にデキューを行う場合には十分ではありません。行ロックがないと、ステータスが更新される前に複数のワーカーが同じタスクを選択してしまう可能性があります。

FOR UPDATE SKIP LOCKED を追加することで、デキュー（キューからの取り出し）の並行処理の安全性が確保されます。各ワーカーは選択した行をロックし、他のワーカーはその行をスキップして次の利用可能なタスクに進みます。さらに、ORDER BY priority DESC, created_at 句により、各優先度レベル内で FIFO 順序を維持しながら、優先度の高いタスクが最初に選択されるようになります。

完全に並行処理セーフで優先度が安定したステートメントは以下の通りです。

リースベースのロックによるクラッシュリカバリ

ワーカーは、VM のエビクション、メモリ不足（out-of-memory）状態、またはデプロイイベントなどにより、タスクの途中で終了することがあります。終了したタスクが処理中のまま残されると、無期限に保持されてしまいます。解決策は、デキュー時に有効期限付きのリースを記録することです。

定期的なスイーパーが、lease_expires_at を過ぎたタスクを再キューイングします。終了したワーカーが保持していたタスクは、外部の調整サービスを必要とせずに、数分以内に自動的に復旧されます。

レート制限を考慮したスロットリング

LLM およびビジョンモデルのエンドポイントでは、通常、1秒あたりのリクエスト数制限と1分あたりのトークン数（TPM）制限という2つの異なるクォータが適用されます。単一のスロットリング戦略が両方に適合することはほとんどありません。オーケストレーターは3つのモードをサポートしており、設定を通じてエージェントごとに選択できます。

並行処理数の制限。 MAX_CONCURRENT_TASKS パラメータは、オーケストレーターが並行してディスパッチするタスクの数を制限します。この制限は、デキュー時にタスクテーブル内の現在の PROCESSING 行数をカウントすることによって適用されます。

カウントが上限以上の場合は、新しいタスクはデキューされません。ローカルエグゼキュータのキューサイズではなく、データベースの行数に基づいてチェックを行うことで、ワーカーの再起動、リースの復旧、マルチレプリカデプロイをまたいでも制限を正確に維持できます。このモードは、各タスクのトークン使用量がほぼ均一で、1秒あたりのリクエスト制限によって制約されるエンドポイントに適しています。

トークン予算。 MAX_TPM パラメータは、実行中のタスク全体の予測トークンレートを制限します。オーケストレーターはタスクのトークン数を推定し、すべての PROCESSING タスクの予測トークンレートを合算します。新しいタスクは、その合計に新しいタスクの予測トークンを加えた値が予算内に収まる場合にのみデキューされます。

複合制限。 MAX_CONCURRENT_TASKS と MAX_TPM の両方が設定されている場合、オーケストレーターはより厳しい方の制約を適用します。このモードは、ある状況下では並行処理数に縛られ（多数の短く低コストなタスク）、別の状況下ではトークン数に縛られる（1つの非常に長いドキュメントが1分あたりのクォータを使い果たす）ようなワークロードを処理します。

3つのモードすべてにおいて、スロットリングの決定はデキュー時に FOR UPDATE SKIP LOCKED と同じトランザクション内で行われます。現在のクォータに収まらないタスクはキューに残り、次のデキューサイクルで再検討されます。個別のスケジューリング状態、インメモリの待機キュー、ワーカーレプリカ間の調整レイヤーは一切不要です。

べき等な Webhook コールバック

AI Agents レイヤーがタスクを完了すると、結果を含むコールバックをオーケストレーターに送信します。コールバックの配信は厳密に1回（exactly-once）とは限りません。Databricks が再試行したり、ネットワークが中断したり、プロキシが再配信したりする可能性があります。コールバックハンドラーはべき等（idempotent）になるよう設計されています。PROCESSING と ENQUEUED の両方の状態を受け入れ、すでに終了状態にあるタスクは何もしない処理（no-op）として扱います。同一のペイロードは同一の結果を生成するため、二重請求や重複処理のリスクが排除されます。

これら4つのパターンを組み合わせることで、並行処理下でも正確で、クラッシュに強く、負荷がかかってもレート制限を考慮し、再試行時にもべき等なタスクキューが実現します。稼働中のシステムに対するリアルタイムの可視性は、次のセクションで説明する別のメカニズムによって提供されます。

リアルタイムオペレーターダッシュボード

多数のドキュメントが処理中の場合、オペレーターはエージェントのパフォーマンス、タスクのステータス、ワークロードのコストを明確に把握する必要があります。オーケストレーターを継続的にポーリングしたり、別のメトリクスプラットフォームに依存したりする必要はありません。オーケストレーターは、この機能を、ワーカーデーモンを実行するのと同じ Databricks App によって表示される単一のダッシュボードに直接統合します。

ダッシュボードの機能

オペレーター向けのダッシュボードには、稼働中のシステムを特徴付ける一連の運用メトリクスが表示されます。すべてのメトリクスは、日付範囲、タスクステータス、およびエージェントによるフィルタリングをサポートしています。

• ステータス別のタスク総数。 各状態（enqueued、processing、completed、failed、cancelled）のタスク数。状態遷移が発生するとリアルタイムで更新されます。
• インプットおよびアウトプットトークン。 MLflow Traces から取得した、タスクごとおよび集計されたトークン数。
• LLM コスト。 MLflow Traces から取得したモデル出力の推定値（各モデル呼び出しから数秒以内に利用可能）。
• コンピュートコスト。 オーケストレーターのタスク実行に起因する Serverless Jobs のコンピュートコスト。system.billing.usage から取得されます。
• 応答時間の中央値。 完了したタスク全体で計算されます。再試行バックオフの外れ値や、飽和状態におけるキューイングテールのレイテンシによる歪みを避けるため、平均値の代わりに中央値が使用されます。
• 信頼度。 AI Agents レイヤーから返されるドキュメントごとの信頼度スコア。タスク結果とともに表示されます。

実装

tasks テーブルの状態変化により、Postgres の LISTEN/NOTIFY イベントがトリガーされます。バックエンドは単一の LISTEN 接続を維持し、接続されているダッシュボードクライアントに Server-Sent Events (SSE) を介してイベントを配信します。ブラウザは EventSource 接続を開き、意味のある状態変化が発生してから約1秒以内にライブアップデートを受信します。この実装には、Redis、WebSocket サーバー、メッセージバスは不要です。

ポーリングは、デフォルトの10秒間隔で永続的なフォールバックとして保持されます。クラウドのイングレスプロキシを経由するストリーミング接続では、クライアント側のエラーイベントを発生させることなく、サイレントにバイトがドロップされることがあります。永続的なポーリングにより、このような場合でもダッシュボードを最新の状態に維持できます。UIインジケーターにより、live（SSEアクティブ）チャネルとpolling（SSE利用不可）チャネルが区別されます。

ダッシュボードのデータは、レイテンシ特性の異なる3つのソースにまたがっています。Postgres（即時）、MLflowのトレースAPI（1秒未満）、およびシステム課金テーブルに対するウェアハウスのクエリ（場合によっては数十秒）です。高速なクエリはすべての更新サイクルにデータを提供し、低速なクエリはユーザーのアクション時にのみ実行され、結果が利用可能になるまでローディング状態で楽観的に返されます。

アプリケーションごとのコスト割り当て

Databricksのシステム課金テーブルはアカウントスコープです。すべてのジョブ、すべてのモデル呼び出し、およびその他のすべてのアプリケーションが、同じsystem.billing.usage行に寄与します。スコープを設定しない場合、アプリケーションレベルの「OCRコスト」タイルは、ワークスペース内のすべてのモデル呼び出しにわたる使用量を集計してしまいます。

解決策は、オーケストレーターが送信したDatabricksジョブの実行（tasks.locked_byおよび task_attempts.run_id）を記録し、課金クエリをそのセットにフィルターすることです。単一のSQLウェアハウスで複数のアプリケーションをサポートでき、各ダッシュボードには独自の支出のみが表示されます。

同じクエリ構成は、オペレーター主導のフィルターとも自然に組み合わせることができます。コストの数値は、他のすべてのダッシュボードメトリクスと同様に、日付範囲、タスクのステータス、またはエージェントによってさらに絞り込むことができ、ダッシュボードを離れることなく、「過去7日間に失敗したタスクのコストはいくらだったか？」や「今月のエージェントXのタスクあたりの支出の中央値はいくらだったか？」といった疑問に対応できます。

これにより、コストの数値を簡単に監視、割り当て、レポートできるようになります。

オーケストレーションのバックボーンとしてのLakebase

Postgresをキューとして使用するパターンは、データエンジニアリングコミュニティで十分に確立されています。Lakebaseは、このパターンをDatabricks上の本番アーキテクチャとして実行可能にする、追加の運用特性を提供します。

• コンピューティングの自動スケーリング。 Lakebaseはワークロードに応じてPostgresのコンピューティングユニットをスケールアップおよびスケールダウンするため、オーケストレーターは24時間体制でピーク容量の料金を支払うことなく、データベースを活用できます。
• OAuthローテーション認証。 Lakebaseは、接続認証に有効期限の短いOAuthトークンを使用します。接続プールはトークンを自動的に更新するため、アプリケーション構成内の静的な資格情報が不要になり、ローテーションのランブックも不要になります。
• Unity Catalogの統合。 Lakebaseは、Databricksの他の機能とアイデンティティ、権限、ガバナンスを共有します。オーケストレーターのサービスプリンシパルは、tasksおよびresultsテーブルに対する明示的な権限を受け取るため、個別のIAM設定は不要です。
• ブランチ作成とスナップショット。 デバッグのために本番タスクテーブルを開発環境にクローンすることは、Lakebaseの標準的な操作であり、ネイティブにサポートされています。

これらの機能により、タスクキューイングのためにセルフホスト型のPostgresの代わりにマネージドメッセージブローカーを採用する動機となるような、運用上の負担が排除されます。

影響と結論

CLAでは、ここで説明するオーケストレーションパターンが本番のドキュメント処理ワークフローをサポートし、抽出時間を数時間から数分に短縮しています。このアーキテクチャは、Databricksネイティブのサービスを使用し、Lakebase Postgresを中心に据えて、外部システムを必要とすることなくキューイング、スケジューリング、およびオブザーバビリティを管理します。これにより、統合のオーバーヘッドを削減しながら、スケールするように構築された統合プラットフォームのメリットを最大限に活用できます。

本番環境において、このパターンは、耐久性のあるタスク管理、優先度制御、レート制限を考慮したスケジューリング、リアルタイムの可視化、およびタスクごとのコスト追跡を提供します。これらの機能が組み合わさることで、周囲のインフラストラクチャをシンプルに保ちながら、エージェント型ワークロードをオーケストレーションする実用的な方法が提供されます。

単一のプラットフォームでAIエージェントのオーケストレーションを簡素化する準備はできましたか？Databricks無料版をお試しいただき、最初のLakebase PostgresプロジェクトとDatabricks Appを作成してください。その後、MLflow Tracingの10分間デモに従って、エージェントワークフローにエンドツーエンドのオブザーバビリティを追加しましょう。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事