295 lines
12 KiB
Python
295 lines
12 KiB
Python
"""Scheduler service for external-task chain bindings (任务分析定时).
|
||
|
||
Mirrors ``ScheduledTaskService``'s proven skeleton: a 30-second poll loop
|
||
scans due bindings and spawns one asyncio task per due binding, refreshing
|
||
each binding's chain snapshot before executing, and skipping bindings whose
|
||
chain has been deleted / disabled (marked ``chain_missing``). A watchdog pass
|
||
fails runs stuck in pending/running past a grace window so one wedged
|
||
execution cannot block the schedule.
|
||
|
||
Overlap policy (数据交付第三期定稿口径)：调度不管上一轮跑没跑完——到期照常触发、
|
||
照常开新期，同任务允许多轮并行，不跳过不顺延。因此**没有**在跑守卫；去重仅靠
|
||
``create_run_once`` 的 (binding_id, scheduled_for) 幂等键。每次执行建 run 的同时开
|
||
一期数据交付刻度（``open_period``，run_id 幂等锚点），智能体提交的数据落在提交
|
||
时刻的当前最新一期。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
from datetime import UTC, datetime, timedelta
|
||
from typing import Any
|
||
from zoneinfo import ZoneInfo
|
||
|
||
from croniter import croniter
|
||
|
||
from deerflow.persistence.task_analysis.base import TaskAnalysisStore
|
||
from deerflow.runtime.task_analysis.executor import ChainBindingExecutor
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
STUCK_GRACE_SECONDS = 60 * 60 # 1h without finishing → watchdog marks failed
|
||
|
||
|
||
def next_run_at(cron_expr: str, timezone: str, base: datetime | None = None) -> datetime:
|
||
tz = ZoneInfo(timezone or "UTC")
|
||
base_local = (base or datetime.now(UTC)).astimezone(tz)
|
||
next_local = croniter(cron_expr, base_local).get_next(datetime)
|
||
if next_local.tzinfo is None:
|
||
next_local = next_local.replace(tzinfo=tz)
|
||
return next_local.astimezone(UTC)
|
||
|
||
|
||
def expand_chain_stages(chain: dict[str, Any]) -> list[list[str]]:
|
||
seats = chain.get("seats") or []
|
||
stages = chain.get("stages")
|
||
if stages:
|
||
return [list(stage) for stage in stages]
|
||
return [[seat["agent_id"]] for seat in seats]
|
||
|
||
|
||
def build_chain_snapshot(chain: dict[str, Any]) -> dict[str, Any]:
|
||
return {
|
||
"title": chain.get("title") or chain.get("id") or "",
|
||
"seats": chain.get("seats") or [],
|
||
"stages": expand_chain_stages(chain),
|
||
"stage_goals": chain.get("stage_goals") or [],
|
||
"coordinator_prompt": chain.get("coordinator_prompt"),
|
||
}
|
||
|
||
|
||
class TaskAnalysisService:
|
||
def __init__(
|
||
self,
|
||
app: Any,
|
||
store: TaskAnalysisStore,
|
||
*,
|
||
poll_interval_seconds: float = 30.0,
|
||
) -> None:
|
||
self.app = app
|
||
self.store = store
|
||
self.poll_interval_seconds = poll_interval_seconds
|
||
self._stopping = asyncio.Event()
|
||
self._loop_task: asyncio.Task | None = None
|
||
|
||
# ── lifecycle ───────────────────────────────────────────────────────────
|
||
|
||
async def start(self) -> None:
|
||
self._stopping = asyncio.Event()
|
||
self._loop_task = asyncio.create_task(self._poll_loop(), name="task-analysis-scheduler")
|
||
|
||
async def stop(self) -> None:
|
||
self._stopping.set()
|
||
if self._loop_task is not None:
|
||
try:
|
||
await asyncio.wait_for(self._loop_task, timeout=5)
|
||
except (TimeoutError, asyncio.CancelledError):
|
||
self._loop_task.cancel()
|
||
self._loop_task = None
|
||
|
||
async def _poll_loop(self) -> None:
|
||
while not self._stopping.is_set():
|
||
try:
|
||
await self._reconcile_stuck_runs()
|
||
except asyncio.CancelledError:
|
||
raise
|
||
except Exception:
|
||
logger.exception("task-analysis watchdog scan failed")
|
||
try:
|
||
due = await self.store.due_bindings(datetime.now(UTC), limit=20)
|
||
for binding in due:
|
||
# 到期即触发，不管上一轮是否结束：同任务允许多轮并行。
|
||
asyncio.create_task(self._execute_and_untrack(binding))
|
||
except asyncio.CancelledError:
|
||
raise
|
||
except Exception:
|
||
logger.exception("task-analysis scan failed")
|
||
try:
|
||
await asyncio.wait_for(self._stopping.wait(), timeout=self.poll_interval_seconds)
|
||
except TimeoutError:
|
||
pass
|
||
|
||
async def _execute_and_untrack(self, binding: dict[str, Any]) -> None:
|
||
try:
|
||
await self.execute_binding(binding["id"], manual=False)
|
||
except Exception:
|
||
logger.exception("task-analysis binding %s execution crashed", binding.get("id"))
|
||
|
||
async def _reconcile_stuck_runs(self) -> None:
|
||
cutoff = datetime.now(UTC) - timedelta(seconds=STUCK_GRACE_SECONDS)
|
||
for run in await self.store.stuck_runs(cutoff):
|
||
logger.warning("task-analysis run %s stuck since %s — failing", run["id"], run["started_at"])
|
||
await self.store.update_run(
|
||
run["id"],
|
||
{
|
||
"status": "failed",
|
||
"finished_at": datetime.now(UTC).isoformat(),
|
||
"error": "执行超时，由调度器看门狗标记失败",
|
||
},
|
||
)
|
||
|
||
# ── execution entry (scheduler + manual trigger share this) ─────────────
|
||
|
||
async def execute_binding(self, binding_id: str, *, manual: bool) -> dict[str, Any] | None:
|
||
"""Synchronous full path (scheduler loop / tests): prepare + carry out."""
|
||
prepared = await self.prepare_binding_run(binding_id, manual=manual)
|
||
if prepared is None:
|
||
return None
|
||
run, binding, task, snapshot = prepared
|
||
await self._carry_out(run, binding, task, snapshot)
|
||
return await self.store.get_run(run["id"])
|
||
|
||
async def launch_binding(self, binding_id: str, *, manual: bool) -> dict[str, Any] | None:
|
||
"""Async entry (manual trigger API): create the run, execute in background.
|
||
|
||
Returns the freshly created run row (status=pending) immediately so the
|
||
HTTP caller can start polling; None when skipped (chain/task missing).
|
||
"""
|
||
prepared = await self.prepare_binding_run(binding_id, manual=manual)
|
||
if prepared is None:
|
||
return None
|
||
run, binding, task, snapshot = prepared
|
||
asyncio.create_task(self._carry_out_and_swallow(run, binding, task, snapshot))
|
||
return run
|
||
|
||
async def prepare_binding_run(
|
||
self, binding_id: str, *, manual: bool
|
||
) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any], dict[str, Any]] | None:
|
||
"""Validate + snapshot + create the run row. Shared by sync/async paths."""
|
||
binding = await self.store.get_binding_by_id(binding_id)
|
||
if binding is None:
|
||
return None
|
||
task = await self.store.get_task(binding["external_task_id"])
|
||
if task is None:
|
||
await self.store.update_binding(binding_id, {"last_status": "task_missing"})
|
||
return None
|
||
|
||
chain = await self._load_chain(binding.get("chain_id") or "", binding.get("created_by") or "")
|
||
if chain is None or chain.get("enabled") is False:
|
||
# 口径：链被删/停用 → 跳过执行并标记，绑定保留。
|
||
await self.store.update_binding(
|
||
binding_id,
|
||
{
|
||
"last_status": "chain_missing",
|
||
"next_run_at": None
|
||
if (manual or self._is_once(binding))
|
||
else self._advance(binding).isoformat(),
|
||
},
|
||
)
|
||
return None
|
||
|
||
snapshot = build_chain_snapshot(chain)
|
||
stages = snapshot["stages"]
|
||
seat_names = {seat["agent_id"]: seat.get("name") or seat["agent_id"] for seat in snapshot["seats"]}
|
||
scheduled_for = datetime.now(UTC).replace(microsecond=0) if manual else _scheduled_for(binding)
|
||
run = await self.store.create_run_once(
|
||
{
|
||
"binding_id": binding_id,
|
||
"external_task_id": binding["external_task_id"],
|
||
"trigger": "manual" if manual else "schedule",
|
||
"scheduled_for": scheduled_for,
|
||
"status": "pending",
|
||
"stages_snapshot": stages,
|
||
"stage_goals_snapshot": snapshot["stage_goals"],
|
||
"seat_names_snapshot": seat_names,
|
||
"seat_results": {},
|
||
}
|
||
)
|
||
await self.store.update_binding(
|
||
binding_id,
|
||
{
|
||
"chain_snapshot": snapshot,
|
||
# once 绑定永不排期（next_run_at 恒空，调度器自然不再触发）。
|
||
**(
|
||
{}
|
||
if (manual or self._is_once(binding))
|
||
else {"next_run_at": self._advance(binding).isoformat()}
|
||
),
|
||
},
|
||
)
|
||
# 数据交付已去掉时间轴（期）维度：调度轮次不再开新期，席位提交统一
|
||
# 合并进任务唯一的当前数据（首次提交时由 submit 兜底开期）。
|
||
return run, binding, task, snapshot
|
||
|
||
async def _carry_out_and_swallow(
|
||
self,
|
||
run: dict[str, Any],
|
||
binding: dict[str, Any],
|
||
task: dict[str, Any],
|
||
snapshot: dict[str, Any],
|
||
) -> None:
|
||
try:
|
||
await self._carry_out(run, binding, task, snapshot)
|
||
except Exception:
|
||
logger.exception("task-analysis background run %s crashed", run.get("id"))
|
||
|
||
async def _carry_out(
|
||
self,
|
||
run: dict[str, Any],
|
||
binding: dict[str, Any],
|
||
task: dict[str, Any],
|
||
snapshot: dict[str, Any],
|
||
) -> None:
|
||
executor = ChainBindingExecutor(self.app, self.store)
|
||
user_id = binding.get("created_by") or ""
|
||
ctx_token = None
|
||
if user_id:
|
||
from deerflow.runtime.task_analysis.executor import _RuntimeUser
|
||
from deerflow.runtime.user_context import reset_current_user, set_current_user
|
||
|
||
ctx_token = set_current_user(_RuntimeUser(id=user_id))
|
||
try:
|
||
await executor.execute(run, {**binding, "chain_snapshot": snapshot}, task)
|
||
finally:
|
||
if ctx_token is not None:
|
||
reset_current_user(ctx_token)
|
||
final = await self.store.get_run(run["id"])
|
||
status = (final or {}).get("status") or "unknown"
|
||
await self.store.update_binding(
|
||
binding["id"],
|
||
{
|
||
"last_run_at": datetime.now(UTC).isoformat(),
|
||
"last_status": status,
|
||
},
|
||
)
|
||
|
||
# ── helpers ─────────────────────────────────────────────────────────────
|
||
|
||
def _advance(self, binding: dict[str, Any]) -> datetime:
|
||
return next_run_at(
|
||
binding.get("cron_expr") or "0 9 * * *",
|
||
binding.get("timezone") or "Asia/Shanghai",
|
||
)
|
||
|
||
@staticmethod
|
||
def _is_once(binding: dict[str, Any]) -> bool:
|
||
return (binding.get("schedule_type") or "cron") == "once"
|
||
|
||
async def _load_chain(self, chain_id: str, user_id: str) -> dict[str, Any] | None:
|
||
chain_store = getattr(self.app.state, "roundtable_chain_store", None)
|
||
if chain_store is None or not chain_id:
|
||
return None
|
||
# created_by 为空（旧数据/未登录创建）时回退 default 用户，再试公共链可见性。
|
||
effective_user = user_id or "default"
|
||
try:
|
||
chain = await chain_store.get_chain(chain_id, effective_user)
|
||
if chain is None and effective_user != "default":
|
||
chain = await chain_store.get_chain(chain_id, "default")
|
||
return chain
|
||
except Exception:
|
||
logger.exception("task-analysis load chain %s failed", chain_id)
|
||
return None
|
||
|
||
|
||
def _scheduled_for(binding: dict[str, Any]) -> datetime:
|
||
value = binding.get("next_run_at")
|
||
if isinstance(value, datetime):
|
||
return value
|
||
if isinstance(value, str) and value:
|
||
try:
|
||
return datetime.fromisoformat(value)
|
||
except ValueError:
|
||
pass
|
||
return datetime.now(UTC).replace(microsecond=0)
|