Dagster と NATS JetStream で後処理を分離する

gateway と消費側の分離

Go + NATS + Dagster AI 基盤では、gateway がイベントを送り、Dagster が独立に処理します。LLM アプリケーションの応答と、保存・データ処理を分ける構成です。

初期の NATS Consumer Worker + GraphQL API trigger bridge を、Dagster native sensor に置き換えました。トピックとジョブの対応はテーブルで管理し、PostgreSQL で重複を防ぎ、correlation_id で処理を追跡します。

Chat、Embedding、Retrieve、Compose は NATS → Dagster で計装しています。record → dataset の Asset ペアを定義し、Dagit UI にリネージを表示します。Go ブリッジも pipelineRoute で対応を管理します。

前提

以下の基盤で動かしています。全体の構成とあわせて説明します。

• メッセージング: NATS 2.11 (JetStream 有効)
• オーケストレーション: Dagster (Python) — daemon + sensor + asset materialization
• データストア: PostgreSQL 18 + pgvector (JIT 有効)
• イベント発行元: agent-gateway (Go/Gin) — fire-and-forget publish

なぜイベント駆動にしたか

初期の Knowledge Service は、チャット応答後に PostgreSQL へ同期書き込みしていました。

gateway をイベント発行だけにし、Dagster が処理する構成に変えました。記録では、この後処理による応答遅延はゼロになっています。後処理の追加も gateway から分離できます。

Consumer アーキテクチャの進化

Phase 1: NATS Consumer Worker + GraphQL bridge

最初は Python の独立プロセスが NATS を購読し、GraphQL API で Dagster ジョブを起動しました。

  agent-gateway → NATS pub → Consumer Worker → GraphQL mutation → Dagster job
  

  async def message_handler(msg):
    payload = json.loads(msg.data.decode())
    submit_dagster_run(
        job_name="knowledge_chat_persist",
        payload=payload,
        correlation_id=payload.get("correlation_id", "")
    )
  

GraphQL mutation:

  mutation LaunchRun($executionParams: LaunchRunExecutionParam!) {
  launchRun(executionParams: $executionParams) {
    __typename
    ... on LaunchRunSuccess {
      run { runId }
    }
    ... on RunConfigValidationInvalid {
      errors { message reason }
    }
  }
}
  

この構成の問題点:

• Consumer Worker が独立プロセスとして管理コストを増やす
• GraphQL API が中間の障害点になる
• Dagster の lineage tracking と consumer の状態が分離している

Phase 2: マルチジョブ fan-out + テーブル駆動

トピックごとの条件分岐が増えたため、対応をテーブルに移しました。

  TOPIC_JOB_MAP = {
    "pipeline.knowledge.chat.persist": [
        ("knowledge_chat_persist", build_audit_run_config),
        ("knowledge_chat_pair_materialize", build_asset_run_config),
    ],
    "pipeline.knowledge.embedding": [
        ("knowledge_embedding", build_event_run_config),
    ],
    "pipeline.knowledge.retrieve": [
        ("knowledge_retrieve", build_event_run_config),
    ],
    "pipeline.knowledge.compose": [
        ("knowledge_compose", build_event_run_config),
    ],
}
  

• トピック追加はテーブルへのエントリ追加のみ — コード複雑度は O(1)
• 各ジョブの設定ビルダ関数を独立管理
• fan-out（1 メッセージ → 複数ジョブ）を宣言的に表現

Phase 3: Dagster native sensor（現行）

その後、GraphQL bridge をやめ、Dagster native sensor に統一しました。

  agent-gateway → NATS JetStream pub (fire-and-forget)
                    ↓
Dagster daemon → sensor (pull subscribe, durable consumer)
                    ↓
                RunRequest → job 直接実行 → asset materialization
  

sensor が JetStream を pull subscribe し、RunRequest を発行します。別の Consumer Worker と GraphQL API が不要になりました。

現在稼働中の sensor 一覧:

トピックとジョブの対応をテーブルで管理する方法は、sensor でも引き継いでいます。

NATS メッセージスキーマ

チャット永続化メッセージ

  {
  "pipeline_id": "knowledge.chat.persist",
  "correlation_id": "unique-trace-id",
  "idempotency_key": "idempotency-unique-id",
  "model": "llm-model-name",
  "prompt": "user-provided-prompt-text",
  "response": "llm-generated-response-text",
  "messages": [{"role": "user", "content": "..."}],
  "usage": {"prompt_tokens": 100, "completion_tokens": 50},
  "requested_at": 1697845200000,
  "payload_json": "additional structured data"
}
  

設計上のポイント

JetStream ストリーム設定

PIPELINE ストリームは Dagster sensor が消費する。TELEMETRY ストリームは Vector (Rust) が subscribe して Prometheus / Loki に流す。

PostgreSQL スキーマ設計

knowledge_chat_audit テーブル

チャット永続化の監査ログ。

  CREATE TABLE knowledge_chat_audit (
  id BIGSERIAL PRIMARY KEY,
  pipeline_id TEXT NOT NULL,
  correlation_id TEXT,
  idempotency_key TEXT,
  model TEXT,
  prompt TEXT,
  response TEXT,
  payload_json JSONB NOT NULL DEFAULT '{}',
  created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
  

knowledge_chat_pairs テーブル

knowledge_chat_audit はリクエストごとの監査ログです。knowledge_chat_pairs は、学習用にプロンプトと応答をペアで保存します。

  CREATE TABLE knowledge_chat_pairs (
  id BIGSERIAL PRIMARY KEY,
  pair_id TEXT UNIQUE NOT NULL,
  correlation_id TEXT,
  pipeline_id TEXT,
  llm_model TEXT,
  prompt TEXT,
  response TEXT,
  prompt_tokens INT,
  completion_tokens INT,
  total_tokens INT,
  payload_json JSONB NOT NULL DEFAULT '{}'::jsonb,
  created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
  

設計判断:

• pair_id の UNIQUE 制約 — deterministic な ID 生成により、同一ペアの重複投入を ON CONFLICT DO NOTHING で防ぐ
• トークン数カラム — fine-tuning 用データセットとしての利用を想定し、トークン統計をペイロードから抽出して正規化カラムに保持
• audit テーブルとの分離 — 監査ログとペアデータセットは用途が異なるため、同一 NATS メッセージから2ジョブをファンアウトして並列に永続化する

knowledge_events テーブル（統一イベント）

embedding、retrieve、compose が共有するイベントテーブルです。

  CREATE TABLE knowledge_events (
    id             BIGSERIAL PRIMARY KEY,
    event_id       TEXT UNIQUE,          -- SHA256(event_type:correlation_id)[:32]
    event_type     TEXT NOT NULL,        -- 'embedding' | 'retrieve' | 'compose'
    correlation_id TEXT NOT NULL,
    pipeline_id    TEXT NOT NULL,
    llm_model      TEXT,
    payload_json   JSONB DEFAULT '{}'::jsonb,
    created_at     TIMESTAMPTZ DEFAULT NOW()
);

CREATE INDEX idx_knowledge_events_type ON knowledge_events (event_type);
CREATE INDEX idx_knowledge_events_corr ON knowledge_events (correlation_id);
CREATE INDEX idx_knowledge_events_date ON knowledge_events (created_at::date);
  

設計判断:

• event_id の UNIQUE 制約 — 冪等性の最終防衛線。同一イベントが複数回処理されても重複挿入されない
• event_type 判別カラム — テーブルを分散させず、1 テーブルに統合することで分析クエリを簡素化
• インデックス — event_type と correlation_id で高速フィルタリング、created_at::date で日次集計

冪等性とエラーハンドリング

冪等設計

すべての副作用は、UPSERT（ON CONFLICT）かUNIQUE制約で重複を吸収します。

エラーハンドリング

• NATS メッセージの ACK は Dagster ジョブ起動成功後に行う（Phase 1/2 では GraphQL 呼び出し成功後）
• 失敗したメッセージは JetStream の再配送キューに保持される
• Dagster 側のジョブ失敗は Dagit UI と slog ログの両方で追跡可能
• correlation_id がすべてのログに含まれるため、gateway → NATS → Dagster → PostgreSQL の全経路を 1 クエリで追跡できる

順序保証

NATS は embedding → retrieve → compose の順序を保証しません。Dagster Asset の依存関係で順序を制御し、correlation_id で事後に関連イベントを結びます。

Dagster Software-Defined Asset の設計

各操作を、抽出する record と保存する dataset の Asset ペアにしました。

Chat Pair アセット

chat_pair_record と chat_pair_dataset の2つのアセットで構成します。

• chat_pair_record: NATS ペイロードからの純粋な抽出。deterministic な pair_id を生成し、トークン数やモデルメタデータを付与する。入力の変換のみで副作用なし
• chat_pair_dataset: chat_pair_record に依存する永続化レイヤー。knowledge_chat_pairs テーブルに INSERT し、ON CONFLICT DO NOTHING で冪等性を保証

Asset の依存関係を、Dagit UI がリネージグラフとして表示します。

Knowledge Event アセット

embedding、retrieve、compose にも、同じ record → dataset の構成を使います。

record は event_id、event_type、correlation_id、pipeline_id、llm_model を抽出します。dataset は knowledge_events に保存します。種類は event_type で区別し、ペイロードの違いは JSONB に保持して、テーブルの増加を避けます。

ファンアウトパターン

NATS の pipeline.knowledge.chat.persist トピックに到着した1つのメッセージは、2つの Dagster ジョブを並列にトリガーする:

1. knowledge_chat_persist → knowledge_chat_audit テーブル（監査ログ）
2. knowledge_chat_pair_materialize → chat_pair_record → chat_pair_dataset → knowledge_chat_pairs テーブル（ペアデータセット）

embedding、retrieve、compose は各 1 トピック・1 ジョブです。fan-out が必要なら、対応テーブルに追加します。

Go 側パイプラインブリッジ

cmd/server/main.go のブリッジもテーブルで対応を管理します。

  type pipelineRoute struct {
    sourceTopic string
    targetJob   string
}
  

startPipelineToDagsterBridge() が pipelineRoute を受け取り、startBridgeForRoute() の goroutine が NATS 購読と Dagster 起動を担当します。

新しい Knowledge 操作を追加する際のステップ:

1. Go 側で NATS トピックに publish するコードを追加
2. Python 側でアセットとジョブを定義
3. Go 側の route テーブルと Python 側の TOPIC_JOB_MAP（または sensor 定義）にエントリを追加

newPipelineService() は NATS に接続できないとき、mock repository を使います。開発時は gateway 単体でも動かせます。

Knowledge モジュールの計装状況

注意事項

• Phase 1/2 の GraphQL trigger bridge コードは参考実装として残っているが、現行の sensor 方式とは併用しない
• JetStream の max-age (72h) を超えた未消費メッセージは失われる。Dagster daemon が 72 時間以上停止した場合は手動リカバリが必要
• knowledge_events.event_id の生成ロジック（SHA256(event_type:correlation_id)[:32]）は、同一 correlation_id で同一 event_type のイベントが複数回発生するケースに注意が必要
• PostgreSQL の init スクリプト（020-create-chat-pairs.sql、030-create-knowledge-events.sql）は初回起動時のみ実行される。既存ボリュームでは手動で DDL を流す必要がある
• Produce、Plan、Execute は未実装のまま。将来フェーズで計装対象に追加予定

検証

• Dagster UI で asset materialization 状態を確認: chat_pair, flow_lineage, tool_call が正常に Materialized
• NATS モニタリング (:8222) で stream 状態と consumer lag を監視
• PostgreSQL で SELECT * FROM knowledge_events WHERE correlation_id = '...' による全経路トレーシングを確認
• 冪等性テスト: 同一メッセージを複数回 publish し、PostgreSQL に重複レコードが生成されないことを検証

関連ファイル

この計装で変更・追加されたファイル:

次のアクション

記録時点では Phase 1（リアルタイム基盤）が稼働しています。以下は拡張案です。

• Dagster Asset Materialization を活用したチャットペアの自動キュレーション
• pg → Parquet → Iceberg バッチ変換（Phase 2、devstack 済み）
• fine-tuning 用データセットの自動生成パイプライン

全体ロードマップはGo 基盤の設計全記録を参照。