532 lines
21 KiB
Python
532 lines
21 KiB
Python
"""Persisted DAG scheduler for workflow runs.
|
||
|
||
Scheduling model
|
||
----------------
|
||
Every edge carries a state: ``pending`` → ``active`` (its source chose this
|
||
port) or ``pruned`` (the source chose a different port). A node is *ready* when
|
||
none of its non-back incoming edges is ``pending`` and at least one is
|
||
``active``; a node whose incoming edges all resolved to ``pruned`` is *skipped*
|
||
and prunes its own outgoing edges in turn. That single rule gives us fan-out
|
||
parallelism, condition branching, and merge joins without a topological
|
||
pre-pass, and it terminates because each superstep resolves at least one edge.
|
||
|
||
Loops are the one cycle we allow: the back edge (``loop`` node, port
|
||
``continue``) is excluded from readiness gating, and activating it clears the
|
||
body's results so the next iteration re-runs them. Every re-entry is capped by
|
||
``settings.max_loop_iterations``.
|
||
|
||
Durability: node results are written to the run store as they complete, so a
|
||
resumed or re-leased run replays only what is missing. ``human_input`` raises
|
||
:class:`WorkflowPaused`, which the caller turns into ``awaiting_input``.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from datetime import UTC, datetime
|
||
from typing import Any, Protocol
|
||
|
||
from deerflow.persistence.workflow_runs.base import WorkflowRunStore
|
||
from deerflow.workflows.errors import WorkflowError
|
||
from deerflow.workflows.ports import ACTIVE_PORTS, LOOP_CONTINUE_PORT
|
||
from deerflow.workflows.runtime.context import RunContext, WorkflowPaused
|
||
from deerflow.workflows.schemas import NodeResult, WorkflowEdge, WorkflowGraph, WorkflowNode
|
||
|
||
|
||
class ExecutorRegistry(Protocol):
|
||
"""Structural view of ``deerflow.workflows.nodes.NodeRegistry``.
|
||
|
||
Declared here rather than imported so the scheduler has no dependency on
|
||
the node package (which depends on the runtime context).
|
||
"""
|
||
|
||
def get(self, node_type: str) -> Any: ...
|
||
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
PENDING = "pending"
|
||
ACTIVE = "active"
|
||
PRUNED = "pruned"
|
||
|
||
|
||
@dataclass
|
||
class EngineResult:
|
||
output: dict[str, Any] = field(default_factory=dict)
|
||
node_results: dict[str, dict[str, Any]] = field(default_factory=dict)
|
||
loop_state: dict[str, int] = field(default_factory=dict)
|
||
steps: int = 0
|
||
|
||
|
||
def _is_back_edge(edge: WorkflowEdge, nodes: dict[str, WorkflowNode]) -> bool:
|
||
source = nodes.get(edge.source)
|
||
return bool(source and source.type == "loop" and (edge.source_port or "") == LOOP_CONTINUE_PORT)
|
||
|
||
|
||
def _node_started_event_data(node: WorkflowNode, node_attempt: int) -> dict[str, Any]:
|
||
"""Safe presentation metadata for workflow clients.
|
||
|
||
This stays intentionally small: UI clients need stable agent identity and
|
||
an optional explicit parallel group, but never prompt templates, tools,
|
||
credentials, or the full node config. A missing group is valid; clients
|
||
may still use their backwards-compatible name-based presentation.
|
||
"""
|
||
|
||
data: dict[str, Any] = {
|
||
"nodeType": node.type,
|
||
"nodeKind": node.type,
|
||
"name": node.name,
|
||
"attempt": node_attempt,
|
||
}
|
||
group_id = node.config.get("parallelGroupId") or node.config.get("parallel_group_id")
|
||
if isinstance(group_id, str) and group_id.strip():
|
||
data["parallelGroupId"] = group_id.strip()
|
||
round_index = node.config.get("roundIndex")
|
||
if round_index is None:
|
||
round_index = node.config.get("round_index")
|
||
if isinstance(round_index, int) and round_index >= 0:
|
||
data["roundIndex"] = round_index
|
||
|
||
if node.type == "agent":
|
||
agent_id = node.config.get("agentId") or node.config.get("agent_id")
|
||
if isinstance(agent_id, str) and agent_id.strip():
|
||
agent: dict[str, str] = {
|
||
"agentId": agent_id.strip(),
|
||
"name": str(
|
||
node.config.get("agentName")
|
||
or node.config.get("agent_name")
|
||
or node.name
|
||
or agent_id
|
||
),
|
||
}
|
||
avatar_url = node.config.get("agentAvatarUrl") or node.config.get("agent_avatar_url")
|
||
if isinstance(avatar_url, str) and avatar_url.strip():
|
||
agent["avatarUrl"] = avatar_url.strip()
|
||
data["agent"] = agent
|
||
return data
|
||
|
||
|
||
class WorkflowEngine:
|
||
def __init__(self, registry: ExecutorRegistry, *, run_store: WorkflowRunStore | None = None) -> None:
|
||
self._registry = registry
|
||
self._store = run_store
|
||
|
||
async def run(
|
||
self,
|
||
ctx: RunContext,
|
||
*,
|
||
preloaded: dict[str, dict[str, Any]] | None = None,
|
||
loop_state: dict[str, int] | None = None,
|
||
attempt: int = 1,
|
||
) -> EngineResult:
|
||
graph: WorkflowGraph = ctx.graph
|
||
nodes = graph.node_map()
|
||
outgoing: dict[str, list[WorkflowEdge]] = {node_id: [] for node_id in nodes}
|
||
incoming: dict[str, list[WorkflowEdge]] = {node_id: [] for node_id in nodes}
|
||
for edge in graph.edges:
|
||
outgoing[edge.source].append(edge)
|
||
incoming[edge.target].append(edge)
|
||
|
||
edge_state: dict[str, str] = {edge.id: PENDING for edge in graph.edges}
|
||
done: set[str] = set()
|
||
skipped: set[str] = set()
|
||
|
||
ctx.loop_state.update(loop_state or {})
|
||
if preloaded:
|
||
ctx.node_results.update(preloaded)
|
||
done |= set(preloaded)
|
||
# Re-derive edge states from the ports each replayed node selected.
|
||
for node_id in list(done):
|
||
self._apply_ports(
|
||
node_id,
|
||
ctx.node_results.get(node_id, {}).get("metadata") or {},
|
||
outgoing,
|
||
edge_state,
|
||
)
|
||
|
||
semaphore = asyncio.Semaphore(max(1, graph.settings.max_parallelism))
|
||
steps = 0
|
||
max_steps = max(1, graph.settings.max_steps)
|
||
|
||
while True:
|
||
ctx.cancel.raise_if_cancelled()
|
||
self._propagate_skips(nodes, incoming, outgoing, edge_state, done, skipped)
|
||
ready = [node_id for node_id in nodes if node_id not in done and node_id not in skipped and self._readiness(node_id, incoming, edge_state, nodes) == "ready"]
|
||
if not ready:
|
||
break
|
||
|
||
steps += len(ready)
|
||
if steps > max_steps:
|
||
raise WorkflowError(
|
||
"WORKFLOW_MAX_STEPS",
|
||
f"运行步数超过上限 {max_steps}",
|
||
details={"maxSteps": max_steps},
|
||
)
|
||
|
||
results = await asyncio.gather(
|
||
*(self._run_one(nodes[node_id], ctx, semaphore, attempt) for node_id in ready),
|
||
return_exceptions=True,
|
||
)
|
||
|
||
paused: WorkflowPaused | None = None
|
||
failure: BaseException | None = None
|
||
for node_id, outcome in zip(ready, results, strict=True):
|
||
if isinstance(outcome, WorkflowPaused):
|
||
paused = paused or outcome
|
||
continue
|
||
if isinstance(outcome, BaseException):
|
||
failure = failure or outcome
|
||
continue
|
||
ctx.node_results[node_id] = outcome
|
||
done.add(node_id)
|
||
metadata = outcome.get("metadata") or {}
|
||
self._apply_ports(node_id, metadata, outgoing, edge_state)
|
||
if nodes[node_id].type == "loop":
|
||
await self._maybe_reset_loop_body(nodes[node_id], ctx, nodes, outgoing, edge_state, done, skipped)
|
||
if failure is not None:
|
||
raise failure
|
||
if paused is not None:
|
||
raise paused
|
||
|
||
return EngineResult(
|
||
output=self._collect_output(nodes, ctx),
|
||
node_results=dict(ctx.node_results),
|
||
loop_state=dict(ctx.loop_state),
|
||
steps=steps,
|
||
)
|
||
|
||
# ── scheduling helpers ──────────────────────────────────────────────
|
||
def _readiness(
|
||
self,
|
||
node_id: str,
|
||
incoming: dict[str, list[WorkflowEdge]],
|
||
edge_state: dict[str, str],
|
||
nodes: dict[str, WorkflowNode],
|
||
) -> str:
|
||
edges = incoming.get(node_id) or []
|
||
if not edges:
|
||
return "ready"
|
||
back = [e for e in edges if _is_back_edge(e, nodes)]
|
||
if any(edge_state[e.id] == ACTIVE for e in back):
|
||
return "ready"
|
||
normal = [e for e in edges if not _is_back_edge(e, nodes)]
|
||
if not normal:
|
||
return "blocked"
|
||
if any(edge_state[e.id] == PENDING for e in normal):
|
||
return "blocked"
|
||
if any(edge_state[e.id] == ACTIVE for e in normal):
|
||
return "ready"
|
||
return "skip"
|
||
|
||
def _apply_ports(
|
||
self,
|
||
node_id: str,
|
||
metadata: dict[str, Any],
|
||
outgoing: dict[str, list[WorkflowEdge]],
|
||
edge_state: dict[str, str],
|
||
) -> None:
|
||
active_ports = metadata.get(ACTIVE_PORTS)
|
||
allowed = {str(p) for p in active_ports} if isinstance(active_ports, list) else None
|
||
for edge in outgoing.get(node_id) or []:
|
||
port = edge.source_port or ""
|
||
if allowed is None or not port or port in allowed:
|
||
edge_state[edge.id] = ACTIVE
|
||
else:
|
||
edge_state[edge.id] = PRUNED
|
||
|
||
def _propagate_skips(
|
||
self,
|
||
nodes: dict[str, WorkflowNode],
|
||
incoming: dict[str, list[WorkflowEdge]],
|
||
outgoing: dict[str, list[WorkflowEdge]],
|
||
edge_state: dict[str, str],
|
||
done: set[str],
|
||
skipped: set[str],
|
||
) -> None:
|
||
changed = True
|
||
while changed:
|
||
changed = False
|
||
for node_id in nodes:
|
||
if node_id in done or node_id in skipped:
|
||
continue
|
||
if self._readiness(node_id, incoming, edge_state, nodes) != "skip":
|
||
continue
|
||
skipped.add(node_id)
|
||
for edge in outgoing.get(node_id) or []:
|
||
edge_state[edge.id] = PRUNED
|
||
changed = True
|
||
|
||
async def _maybe_reset_loop_body(
|
||
self,
|
||
loop_node: WorkflowNode,
|
||
ctx: RunContext,
|
||
nodes: dict[str, WorkflowNode],
|
||
outgoing: dict[str, list[WorkflowEdge]],
|
||
edge_state: dict[str, str],
|
||
done: set[str],
|
||
skipped: set[str],
|
||
) -> None:
|
||
back_edges = [e for e in outgoing.get(loop_node.id) or [] if _is_back_edge(e, nodes) and edge_state[e.id] == ACTIVE]
|
||
if not back_edges:
|
||
return
|
||
body_entry = str(loop_node.config.get("bodyEntry") or loop_node.config.get("body_entry") or "")
|
||
if not body_entry or body_entry not in nodes:
|
||
raise WorkflowError(
|
||
"WORKFLOW_LOOP_FAILED",
|
||
"循环节点的 bodyEntry 不存在",
|
||
node_id=loop_node.id,
|
||
details={"bodyEntry": body_entry},
|
||
)
|
||
body = self._loop_body(body_entry, loop_node.id, outgoing, nodes)
|
||
for node_id in body:
|
||
done.discard(node_id)
|
||
skipped.discard(node_id)
|
||
ctx.node_results.pop(node_id, None)
|
||
for edge in outgoing.get(node_id) or []:
|
||
if not _is_back_edge(edge, nodes):
|
||
edge_state[edge.id] = PENDING
|
||
# The back edge stays ACTIVE so the body entry is reachable again; the
|
||
# loop node prunes it itself once it selects the ``done`` port.
|
||
if self._store is not None and body:
|
||
try:
|
||
await self._store.invalidate_node_runs(ctx.run_id, sorted(body))
|
||
except Exception: # noqa: BLE001 - bookkeeping must not kill the run
|
||
logger.warning("failed to invalidate node runs for loop reset", exc_info=True)
|
||
await ctx.emit_event(
|
||
"node.progress",
|
||
data={
|
||
"phase": "loop.iteration",
|
||
"iteration": ctx.loop_state.get(loop_node.id, 0),
|
||
"bodyEntry": body_entry,
|
||
},
|
||
node_id=loop_node.id,
|
||
)
|
||
|
||
def _loop_body(
|
||
self,
|
||
body_entry: str,
|
||
loop_id: str,
|
||
outgoing: dict[str, list[WorkflowEdge]],
|
||
nodes: dict[str, WorkflowNode],
|
||
) -> set[str]:
|
||
"""Nodes reachable from ``body_entry`` up to (and including) the loop node."""
|
||
body: set[str] = set()
|
||
stack = [body_entry]
|
||
while stack:
|
||
current = stack.pop()
|
||
if current in body:
|
||
continue
|
||
body.add(current)
|
||
if current == loop_id:
|
||
continue
|
||
for edge in outgoing.get(current) or []:
|
||
if not _is_back_edge(edge, nodes):
|
||
stack.append(edge.target)
|
||
return body
|
||
|
||
# ── node execution ──────────────────────────────────────────────────
|
||
async def _run_one(
|
||
self,
|
||
node: WorkflowNode,
|
||
ctx: RunContext,
|
||
semaphore: asyncio.Semaphore,
|
||
attempt: int,
|
||
) -> dict[str, Any]:
|
||
async with semaphore:
|
||
ctx.cancel.raise_if_cancelled(node.id)
|
||
retry_cfg = node.config.get("retry") if isinstance(node.config.get("retry"), dict) else {}
|
||
max_attempts = max(1, min(5, int(retry_cfg.get("maxAttempts") or 1)))
|
||
backoff = max(0.0, float(retry_cfg.get("backoffSeconds") or 0.5))
|
||
on_error = str(node.config.get("onError") or "fail")
|
||
timeout = self._node_timeout(node, ctx)
|
||
|
||
last_error: WorkflowError | None = None
|
||
for node_attempt in range(1, max_attempts + 1):
|
||
node_run = await self._begin_node_run(ctx, node, attempt, node_attempt)
|
||
started = time.monotonic()
|
||
await ctx.emit_event(
|
||
"node.started",
|
||
data=_node_started_event_data(node, node_attempt),
|
||
node_id=node.id,
|
||
)
|
||
try:
|
||
executor = self._registry.get(node.type)
|
||
result: NodeResult = await asyncio.wait_for(executor.execute(node, ctx), timeout=timeout)
|
||
except WorkflowPaused:
|
||
await self._finish_node_run(ctx, node_run, status="queued", output=None)
|
||
raise
|
||
except TimeoutError as exc:
|
||
last_error = WorkflowError(
|
||
"WORKFLOW_TIMEOUT",
|
||
f"节点执行超时（{timeout}s）",
|
||
retryable=True,
|
||
node_id=node.id,
|
||
details={"timeoutSeconds": timeout},
|
||
)
|
||
await self._fail_node_run(ctx, node_run, last_error, started)
|
||
if node_attempt < max_attempts:
|
||
await asyncio.sleep(backoff * node_attempt)
|
||
continue
|
||
if on_error == "continue":
|
||
return self._error_result(last_error)
|
||
raise last_error from exc
|
||
except WorkflowError as exc:
|
||
last_error = exc
|
||
exc.node_id = exc.node_id or node.id
|
||
await self._fail_node_run(ctx, node_run, exc, started)
|
||
if exc.retryable and node_attempt < max_attempts:
|
||
await asyncio.sleep(backoff * node_attempt)
|
||
continue
|
||
if on_error == "continue":
|
||
return self._error_result(exc)
|
||
raise
|
||
except asyncio.CancelledError:
|
||
raise
|
||
except Exception as exc: # noqa: BLE001 - unexpected node bug
|
||
logger.exception("workflow node %s crashed", node.id)
|
||
last_error = WorkflowError(
|
||
"WORKFLOW_INTERNAL",
|
||
"节点执行发生内部错误",
|
||
node_id=node.id,
|
||
details={"reason": type(exc).__name__},
|
||
)
|
||
await self._fail_node_run(ctx, node_run, last_error, started)
|
||
if on_error == "continue":
|
||
return self._error_result(last_error)
|
||
raise last_error from exc
|
||
|
||
payload = result.model_dump(by_alias=False)
|
||
duration = int((time.monotonic() - started) * 1000)
|
||
await self._finish_node_run(ctx, node_run, status="completed", output=payload, duration_ms=duration)
|
||
await ctx.emit_event(
|
||
"node.completed",
|
||
data={
|
||
"nodeType": node.type,
|
||
"durationMs": duration,
|
||
"data": payload.get("data"),
|
||
"warnings": payload.get("warnings") or [],
|
||
},
|
||
node_id=node.id,
|
||
)
|
||
return payload
|
||
|
||
raise last_error or WorkflowError("WORKFLOW_INTERNAL", "节点执行失败", node_id=node.id)
|
||
|
||
def _node_timeout(self, node: WorkflowNode, ctx: RunContext) -> int:
|
||
configured = node.config.get("timeoutSeconds") or node.config.get("timeout_seconds")
|
||
graph_limit = ctx.graph.settings.node_timeout_seconds
|
||
try:
|
||
value = int(configured) if configured else graph_limit
|
||
except (TypeError, ValueError):
|
||
value = graph_limit
|
||
return max(1, min(value, graph_limit))
|
||
|
||
def _error_result(self, error: WorkflowError) -> dict[str, Any]:
|
||
return {
|
||
"data": {},
|
||
"messages": [],
|
||
"artifacts": [],
|
||
"metadata": {"failed": True, "error": error.to_body().model_dump(by_alias=True)},
|
||
"warnings": [error.message],
|
||
}
|
||
|
||
async def _begin_node_run(self, ctx: RunContext, node: WorkflowNode, run_attempt: int, node_attempt: int) -> dict[str, Any] | None:
|
||
if self._store is None:
|
||
return None
|
||
try:
|
||
row = await self._store.record_node_run(
|
||
{
|
||
"run_id": ctx.run_id,
|
||
"node_id": node.id,
|
||
"node_type": node.type,
|
||
"status": "running",
|
||
"attempt": node_attempt,
|
||
"iteration": int(ctx.loop_state.get(node.id, 0)),
|
||
"started_at": datetime.now(UTC),
|
||
"is_current": True,
|
||
}
|
||
)
|
||
except Exception: # noqa: BLE001
|
||
logger.warning("failed to record node run start for %s", node.id, exc_info=True)
|
||
return None
|
||
ctx.node_run_ids[node.id] = str(row.get("id") or "")
|
||
return row
|
||
|
||
async def _finish_node_run(
|
||
self,
|
||
ctx: RunContext,
|
||
node_run: dict[str, Any] | None,
|
||
*,
|
||
status: str,
|
||
output: dict[str, Any] | None,
|
||
duration_ms: int | None = None,
|
||
) -> None:
|
||
if self._store is None or node_run is None:
|
||
return
|
||
try:
|
||
await self._store.record_node_run(
|
||
{
|
||
"run_id": ctx.run_id,
|
||
"node_id": node_run["node_id"],
|
||
"node_type": node_run["node_type"],
|
||
"attempt": node_run["attempt"],
|
||
"iteration": node_run["iteration"],
|
||
"status": status,
|
||
"output": output,
|
||
"duration_ms": duration_ms,
|
||
"finished_at": datetime.now(UTC),
|
||
"is_current": True,
|
||
}
|
||
)
|
||
except Exception: # noqa: BLE001
|
||
logger.warning("failed to record node run finish", exc_info=True)
|
||
|
||
async def _fail_node_run(
|
||
self,
|
||
ctx: RunContext,
|
||
node_run: dict[str, Any] | None,
|
||
error: WorkflowError,
|
||
started: float,
|
||
) -> None:
|
||
duration = int((time.monotonic() - started) * 1000)
|
||
body = error.to_body().model_dump(by_alias=True)
|
||
if self._store is not None and node_run is not None:
|
||
try:
|
||
await self._store.record_node_run(
|
||
{
|
||
"run_id": ctx.run_id,
|
||
"node_id": node_run["node_id"],
|
||
"node_type": node_run["node_type"],
|
||
"attempt": node_run["attempt"],
|
||
"iteration": node_run["iteration"],
|
||
"status": "failed",
|
||
"error": body,
|
||
"duration_ms": duration,
|
||
"finished_at": datetime.now(UTC),
|
||
"is_current": True,
|
||
}
|
||
)
|
||
except Exception: # noqa: BLE001
|
||
logger.warning("failed to record node run failure", exc_info=True)
|
||
await ctx.emit_event(
|
||
"node.failed",
|
||
data={"error": body, "durationMs": duration},
|
||
node_id=error.node_id,
|
||
)
|
||
|
||
def _collect_output(self, nodes: dict[str, WorkflowNode], ctx: RunContext) -> dict[str, Any]:
|
||
output: dict[str, Any] = {}
|
||
for node_id, node in nodes.items():
|
||
if node.type != "output":
|
||
continue
|
||
result = ctx.node_results.get(node_id)
|
||
if isinstance(result, dict) and isinstance(result.get("data"), dict):
|
||
output.update(result["data"])
|
||
return output
|
||
|
||
|
||
__all__ = ["EngineResult", "WorkflowEngine"]
|