356 lines
16 KiB
Python
356 lines
16 KiB
Python
"""圆桌后台作业 dispatcher（Phase 3）：跨 worker 领取 + 崩溃接管 + outbox 补偿。
|
||
|
||
把「谁来跑这个作业」从「发起请求的那个 worker 的进程内存」收敛到 **DB 租约**：
|
||
|
||
- **领取（claim）**：后台循环定期 ``list_claimable``（``queued`` 新作业/续跑，或 ``running``
|
||
但租约过期的僵尸），对每个候选用 ``claim_job`` 的**原子条件 UPDATE** 抢占 —— 多 worker
|
||
并发扫到同一候选时只有一个 ``rowcount==1``，其余拿到 None 自动跳过。一致性来自 DB，
|
||
不是进程内锁。
|
||
- **心跳**：领取成功后交给 ``RoundtableJobExecutor`` 跑，执行器内部定期 ``renew_lease``；
|
||
worker 崩溃则心跳停止，租约最长一个 TTL 后过期，下一轮循环在其它 worker 重新领取
|
||
（``attempt+1``）—— 崩溃可接管。
|
||
- **启动扫描**：本 worker 启动时立即跑一轮 ``dispatch_once``，尽快接走历史遗留的
|
||
queued / 租约过期作业（不用等第一个轮询周期）。
|
||
- **outbox 补偿器**：``draft_persist_pending=1`` 的作业（后台研讨结果因乐观锁冲突没能写回
|
||
草稿）在这里被重新投影到草稿 step2.runs，成功后清标记 —— 「研讨结果不丢」由 DB 标记驱动。
|
||
|
||
所有 worker 都跑同一个 dispatcher；谁先 ``claim_job`` 成功谁执行，天然幂等、互斥。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
from contextlib import suppress
|
||
from datetime import UTC, datetime, timedelta
|
||
from typing import Any
|
||
|
||
from app.gateway.roundtable_job_executor import (
|
||
_LEASE_TTL_SECONDS,
|
||
WORKER_ID,
|
||
params_from_snapshot,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# dispatcher 轮询间隔（秒）。路由 start/resume 成功后会 ``nudge()`` 立即唤醒，故正常路径
|
||
# 几乎无感延迟；该间隔只是「没有 nudge 时的兜底」，兼作崩溃接管的发现周期。
|
||
_DISPATCH_POLL_INTERVAL = 3.0
|
||
# outbox 补偿器轮询间隔（秒）。pending 标记是低频异常路径，慢一点无妨。
|
||
_COMPENSATE_INTERVAL = 10.0
|
||
# 单轮最多领取的作业数（防止一次扫太多把事件循环占死）。
|
||
_CLAIM_BATCH = 20
|
||
|
||
|
||
class RoundtableJobDispatcher:
|
||
"""跨 worker 的圆桌后台作业调度器（每个 worker 进程一个实例，lifespan 起停）。"""
|
||
|
||
def __init__(
|
||
self,
|
||
store,
|
||
executor,
|
||
*,
|
||
draft_store=None,
|
||
task_draft_store=None,
|
||
worker_id: str | None = None,
|
||
) -> None:
|
||
self._store = store
|
||
self._executor = executor
|
||
# outbox 补偿器投影草稿用（个人草稿 / task 草稿两个存储）。
|
||
self._draft_store = draft_store
|
||
self._task_draft_store = task_draft_store
|
||
self._worker_id = worker_id or WORKER_ID
|
||
self._wake = asyncio.Event()
|
||
self._task: asyncio.Task | None = None
|
||
self._stopping = False
|
||
self._last_compensate = datetime.now(UTC)
|
||
|
||
# ── 生命周期 ─────────────────────────────────────────────────────────────
|
||
|
||
def start(self) -> None:
|
||
"""启动后台循环（lifespan startup 调用）。幂等。"""
|
||
if self._task is not None and not self._task.done():
|
||
return
|
||
self._stopping = False
|
||
self._task = asyncio.create_task(self._loop(), name="roundtable-job-dispatcher")
|
||
logger.info("[roundtable-jobs] dispatcher started (worker=%s)", self._worker_id)
|
||
|
||
async def stop(self) -> None:
|
||
"""停止后台循环（lifespan shutdown 调用）。"""
|
||
self._stopping = True
|
||
self._wake.set()
|
||
if self._task is not None:
|
||
self._task.cancel()
|
||
with suppress(asyncio.CancelledError, Exception):
|
||
await self._task
|
||
self._task = None
|
||
|
||
def nudge(self) -> None:
|
||
"""立即唤醒一轮领取（路由 start/resume 成功后调用，降低启动延迟）。"""
|
||
self._wake.set()
|
||
|
||
# ── 主循环 ───────────────────────────────────────────────────────────────
|
||
|
||
async def _loop(self) -> None:
|
||
# 启动即扫一轮：尽快接管历史遗留 queued / 租约过期（僵尸）作业。
|
||
await self._safe_dispatch()
|
||
while not self._stopping:
|
||
try:
|
||
await asyncio.wait_for(self._wake.wait(), timeout=_DISPATCH_POLL_INTERVAL)
|
||
except TimeoutError:
|
||
pass
|
||
self._wake.clear()
|
||
if self._stopping:
|
||
break
|
||
await self._safe_dispatch()
|
||
# outbox 补偿器按自己的节奏跑（低频）。
|
||
if (datetime.now(UTC) - self._last_compensate).total_seconds() >= _COMPENSATE_INTERVAL:
|
||
self._last_compensate = datetime.now(UTC)
|
||
await self._safe_compensate()
|
||
|
||
async def _safe_dispatch(self) -> None:
|
||
try:
|
||
await self.dispatch_once()
|
||
except Exception: # noqa: BLE001 — 单轮失败不能杀死循环，下轮重试
|
||
logger.exception("[roundtable-jobs] dispatch round failed")
|
||
|
||
async def _safe_compensate(self) -> None:
|
||
try:
|
||
await self.compensate_once()
|
||
except Exception: # noqa: BLE001
|
||
logger.exception("[roundtable-jobs] compensate round failed")
|
||
|
||
# ── 领取 ─────────────────────────────────────────────────────────────────
|
||
|
||
async def dispatch_once(self) -> int:
|
||
"""扫一轮可领取作业并逐个条件抢占。返回本轮成功启动的作业数。"""
|
||
now = datetime.now(UTC)
|
||
await self._reap_cancel_requested(now)
|
||
try:
|
||
candidate_ids = await self._store.list_claimable(now=now, limit=_CLAIM_BATCH)
|
||
except Exception: # noqa: BLE001
|
||
logger.warning("[roundtable-jobs] list_claimable failed", exc_info=True)
|
||
return 0
|
||
if not candidate_ids:
|
||
return 0
|
||
|
||
started = 0
|
||
lease_until = now + timedelta(seconds=_LEASE_TTL_SECONDS)
|
||
for job_id in candidate_ids:
|
||
try:
|
||
row = await self._store.claim_job(
|
||
job_id, lease_owner=self._worker_id, lease_until=lease_until
|
||
)
|
||
except Exception: # noqa: BLE001
|
||
logger.warning("[roundtable-jobs] claim_job failed for %s", job_id, exc_info=True)
|
||
continue
|
||
if row is None:
|
||
# 别的 worker 先领了 / 状态已变 —— 条件 UPDATE rowcount=0，跳过即可。
|
||
continue
|
||
params = params_from_snapshot(
|
||
row.get("input_snapshot"), job_id=job_id, lease_owner=self._worker_id
|
||
)
|
||
if params is None:
|
||
# 快照缺失/损坏：无法重建执行参数 —— 标记 error 并释放活跃键，避免僵死占位。
|
||
logger.error(
|
||
"[roundtable-jobs] job %s claimed but snapshot missing/corrupt; marking error",
|
||
job_id,
|
||
)
|
||
with suppress(Exception):
|
||
await self._store.update_progress(
|
||
job_id,
|
||
lease_owner=self._worker_id,
|
||
status="error",
|
||
error="input_snapshot 缺失或损坏，无法在后台重建研讨",
|
||
)
|
||
continue
|
||
self._executor.start_job(params)
|
||
started += 1
|
||
logger.info(
|
||
"[roundtable-jobs] worker=%s claimed job %s (attempt=%s, resume=%s)",
|
||
self._worker_id, job_id, row.get("attempt"), params.is_resume,
|
||
)
|
||
return started
|
||
|
||
async def _reap_cancel_requested(self, now: datetime) -> int:
|
||
"""收口持租 worker 崩溃后遗留的 ``cancel_requested`` 作业。
|
||
|
||
有效租约仍在时，必须让持租 worker 的取消看门狗收口；无租约或租约过期时，已经
|
||
不存在可信的持租者，此处通过 ``finalize_cancel`` 的条件 UPDATE 安全地回收。
|
||
"""
|
||
try:
|
||
candidate_ids = await self._store.list_cancel_reapable(
|
||
now=now, limit=_CLAIM_BATCH
|
||
)
|
||
except Exception: # noqa: BLE001
|
||
logger.warning("[roundtable-jobs] list_cancel_reapable failed", exc_info=True)
|
||
return 0
|
||
|
||
finalized = 0
|
||
for job_id in candidate_ids:
|
||
try:
|
||
row = await self._store.finalize_cancel(job_id)
|
||
except Exception: # noqa: BLE001
|
||
logger.warning(
|
||
"[roundtable-jobs] finalize orphaned cancellation failed for %s",
|
||
job_id,
|
||
exc_info=True,
|
||
)
|
||
continue
|
||
if row is not None:
|
||
finalized += 1
|
||
logger.info(
|
||
"[roundtable-jobs] reaped orphaned cancellation for job %s", job_id
|
||
)
|
||
return finalized
|
||
|
||
# ── outbox 补偿器 ────────────────────────────────────────────────────────
|
||
|
||
def _draft_store_for(self, task_id: str | None):
|
||
"""task 作业 → task-draft 存储；普通作业 → 个人草稿存储（与执行器同规则）。"""
|
||
return self._task_draft_store if task_id else self._draft_store
|
||
|
||
async def compensate_once(self) -> int:
|
||
"""把 ``draft_persist_pending=1`` 的作业结果投影回草稿，成功后清标记。
|
||
|
||
投影逻辑与执行器 ``_write_run_to_draft`` 的终态写一致：读最新草稿 → 按
|
||
``run_id=job-<jobId>`` upsert 一个 run（用作业行里的 dialogues/step3/status）→
|
||
带 ``expected_version`` 乐观锁写回。写成功（或草稿已含该 run 的终态）即清标记。
|
||
返回本轮成功投影的作业数。
|
||
"""
|
||
try:
|
||
pending = await self._store.list_draft_persist_pending(limit=_CLAIM_BATCH)
|
||
except Exception: # noqa: BLE001
|
||
logger.warning("[roundtable-jobs] list_draft_persist_pending failed", exc_info=True)
|
||
return 0
|
||
if not pending:
|
||
return 0
|
||
|
||
projected = 0
|
||
for job in pending:
|
||
job_id = job.get("id")
|
||
draft_id = job.get("draft_id")
|
||
draft_store = self._draft_store_for(job.get("task_id"))
|
||
if not draft_id or draft_store is None:
|
||
# 无草稿可投影（如纯 task 作业但 task_store 未配）——直接清标记，避免死循环。
|
||
with suppress(Exception):
|
||
await self._store.clear_draft_persist_pending(job_id)
|
||
continue
|
||
try:
|
||
ok = await self._project_job_to_draft(draft_store, job)
|
||
except Exception: # noqa: BLE001
|
||
logger.warning(
|
||
"[roundtable-jobs] compensate project failed for job %s", job_id, exc_info=True
|
||
)
|
||
continue
|
||
if ok:
|
||
with suppress(Exception):
|
||
await self._store.clear_draft_persist_pending(job_id)
|
||
projected += 1
|
||
logger.info("[roundtable-jobs] compensated draft write for job %s", job_id)
|
||
return projected
|
||
|
||
async def _project_job_to_draft(self, draft_store: Any, job: dict[str, Any]) -> bool:
|
||
"""把单个作业的最新状态投影成草稿 step2.runs 里的一个 run。成功返回 True。
|
||
|
||
与执行器不同：这里没有 ``params``（作业可能跑在别的 worker / 已结束），只能用作业行
|
||
里已落库的 ``dialogues``/``step3``/``status`` 重建 run —— 这些正是「研讨结果」本体，
|
||
足够让前端加载草稿时看到完整的后台研讨记录与报告。
|
||
"""
|
||
from deerflow.persistence.roundtable_drafts.sql import DraftConcurrentWriteError
|
||
|
||
job_id = job.get("id")
|
||
draft_id = job.get("draft_id")
|
||
user_id = job.get("user_id")
|
||
run_id = f"job-{job_id}"
|
||
status = job.get("status") or "running"
|
||
dialogues = job.get("dialogues") or []
|
||
step3 = job.get("step3")
|
||
|
||
# task 草稿不分权：get/update 不传 user_id；个人草稿按 user。
|
||
is_task = bool(job.get("task_id"))
|
||
|
||
for _attempt in range(3):
|
||
try:
|
||
draft = (
|
||
await draft_store.get_draft(draft_id)
|
||
if is_task
|
||
else await draft_store.get_draft(draft_id, user_id)
|
||
)
|
||
except Exception: # noqa: BLE001
|
||
draft = None
|
||
if draft is None:
|
||
# 草稿已被删除 —— 无处投影，视作完成（清标记）。
|
||
return True
|
||
expected_version = draft.get("version")
|
||
|
||
step2 = dict(draft.get("step2") or {})
|
||
runs = list(step2.get("runs") or [])
|
||
idx = next((i for i, r in enumerate(runs) if r.get("id") == run_id), -1)
|
||
created = runs[idx].get("createdAt") if idx >= 0 else (job.get("created_at") or datetime.now(UTC).isoformat())
|
||
prev_step3 = runs[idx].get("step3") if idx >= 0 else None
|
||
effective_step3 = step3 if step3 is not None else prev_step3
|
||
|
||
# 从 dialogues 里提取席位信息（尽量还原 selectedAgents）。
|
||
agents_seen: list[dict[str, str]] = []
|
||
seen_ids: set[str] = set()
|
||
for d in dialogues:
|
||
aid = str(d.get("agentId") or "")
|
||
if aid and aid not in seen_ids and d.get("role") == "seat":
|
||
seen_ids.add(aid)
|
||
agents_seen.append({"agent_id": aid, "name": str(d.get("name") or aid)})
|
||
|
||
run = {
|
||
"id": run_id,
|
||
"title": "后台研讨",
|
||
"createdAt": created,
|
||
"source": job.get("orchestration_mode") or "recommend",
|
||
"selectedAgents": agents_seen,
|
||
"threadIds": job.get("thread_ids"),
|
||
"coordinatorName": job.get("coordinator_name"),
|
||
"step2RoundtableDialogues": dialogues,
|
||
"lastLeaderContent": next(
|
||
(d.get("content") for d in reversed(dialogues) if d.get("role") == "leader"), ""
|
||
) or "",
|
||
"hasConsensus": status == "done",
|
||
"consensusPercentage": int(job.get("consensus_percentage") or (100 if status == "done" else 0)),
|
||
"budgetLimit": 0,
|
||
"orchestrationMode": job.get("orchestration_mode") or "recommend",
|
||
"chain": job.get("chain"),
|
||
"orchestrationPlan": job.get("orchestration_plan"),
|
||
"step3": effective_step3,
|
||
"jobId": job_id,
|
||
"jobStatus": status,
|
||
}
|
||
runs = [run if i == idx else r for i, r in enumerate(runs)] if idx >= 0 else [*runs, run]
|
||
step2["runs"] = runs
|
||
step2["activeRunId"] = run_id
|
||
has_report = effective_step3 is not None
|
||
update_kwargs: dict[str, Any] = {
|
||
"step2": step2,
|
||
"furthest_step": 3 if (status == "done" and has_report) else 2,
|
||
}
|
||
if has_report:
|
||
update_kwargs["step3"] = effective_step3
|
||
|
||
try:
|
||
if is_task:
|
||
await draft_store.update_draft(
|
||
draft_id, expected_version=expected_version, **update_kwargs
|
||
)
|
||
else:
|
||
await draft_store.update_draft(
|
||
draft_id, user_id, expected_version=expected_version, **update_kwargs
|
||
)
|
||
return True
|
||
except DraftConcurrentWriteError:
|
||
# 并发写冲突 —— 重读最新草稿再合并重试（与执行器同款乐观锁重试）。
|
||
continue
|
||
except Exception: # noqa: BLE001
|
||
logger.warning(
|
||
"[roundtable-jobs] compensate update_draft failed for %s", draft_id, exc_info=True
|
||
)
|
||
return False
|
||
# 乐观锁重试耗尽：保留 pending 标记，等下一轮补偿器再试。
|
||
return False
|