butler/dag_scheduler.py
Rama main · réplica de lectura del repositorio original.
"""DAG topology helpers extracted from ``TaskOrchestrator.execute_graph`` (ENG-4)."""
from __future__ import annotations
from collections.abc import Callable
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from butler.task_orchestrator import AgentResult, TaskNode
def topological_sort(nodes: list["TaskNode"]) -> list[str]:
"""Kahn's algorithm for topological ordering."""
graph: dict[str, list[str]] = {n.id: [] for n in nodes}
in_degree: dict[str, int] = {n.id: 0 for n in nodes}
for n in nodes:
for dep in n.depends_on:
if dep in graph:
graph[dep].append(n.id)
in_degree[n.id] += 1
queue = [nid for nid, deg in in_degree.items() if deg == 0]
result: list[str] = []
while queue:
nid = queue.pop(0)
result.append(nid)
for neighbor in graph.get(nid, []):
in_degree[neighbor] -= 1
if in_degree[neighbor] == 0:
queue.append(neighbor)
if len(result) != len(nodes):
raise ValueError("Task graph contains a cycle")
return result
def group_into_layers(order: list[str], node_map: dict[str, "TaskNode"]) -> list[list[str]]:
"""Group topologically sorted nodes into parallel layers."""
depth: dict[str, int] = {}
for nid in order:
node = node_map[nid]
if not node.depends_on:
depth[nid] = 0
else:
depth[nid] = max(depth.get(d, 0) for d in node.depends_on) + 1
max_depth = max(depth.values()) if depth else 0
layers: list[list[str]] = [[] for _ in range(max_depth + 1)]
for nid in order:
layers[depth[nid]].append(nid)
return [layer for layer in layers if layer]
def first_cancelled_dependency(node: "TaskNode", cancelled: set[str]) -> str:
for dep_id in node.depends_on:
if dep_id in cancelled:
return str(dep_id)
return ""
def direct_dependents(node_map: dict[str, "TaskNode"], parent_id: str) -> list["TaskNode"]:
return [node for node in node_map.values() if parent_id in node.depends_on]
def cancel_direct_dependents(
node_map: dict[str, "TaskNode"],
parent_id: str,
cancelled: set[str],
) -> None:
for child in direct_dependents(node_map, parent_id):
cancelled.add(child.id)
def first_failed_dependency(
node: "TaskNode",
completed: dict[str, "AgentResult"],
node_map: dict[str, "TaskNode"],
) -> str:
from butler.core.workflow_flags import workflow_optional_enabled
for dep_id in node.depends_on:
dep_result = completed.get(dep_id)
if dep_result is not None and not dep_result.success:
dep_node = node_map.get(dep_id)
if (
workflow_optional_enabled()
and dep_node is not None
and dep_node.optional
):
continue
return str(dep_id)
return ""
def graph_all_required_ok(
completed: dict[str, "AgentResult"],
node_map: dict[str, "TaskNode"],
) -> bool:
from butler.core.workflow_flags import workflow_optional_enabled
for nid, result in completed.items():
node = node_map.get(nid)
if result.success:
continue
if workflow_optional_enabled() and node is not None and node.optional:
continue
return False
return True
def format_dependency_context(
dep_id: str,
dep_result: "AgentResult",
*,
handoff_only: bool,
) -> str:
if handoff_only and dep_result.report is not None:
from butler.core.handoff import render_handoff_block
rep = dep_result.report
return str(render_handoff_block(
from_role=dep_id,
to_role="",
current_state=(rep.summary or dep_result.response or "")[:400],
deliverable=rep.headline or dep_id,
prior_report=rep,
))
text = (dep_result.response or dep_result.error or "").strip()
if not text:
return ""
return f"[{dep_id} 结果]: {text[:1500]}"
def inject_dependency_context(
node: "TaskNode",
completed: dict[str, "AgentResult"],
) -> None:
if not node.depends_on:
return
dep_contexts: list[str] = []
for dep_id in node.depends_on:
dep_result = completed.get(dep_id)
if dep_result is None:
continue
block = format_dependency_context(
dep_id,
dep_result,
handoff_only=node.handoff_only,
)
if block:
dep_contexts.append(block)
if dep_contexts:
node.config.context = (
(node.config.context or "") + "\n\n" + "\n".join(dep_contexts)
)
def inject_supervisor_note(node: "TaskNode") -> None:
if not node.supervisor_note.strip():
return
sup = f"## Supervisor 指令\n{node.supervisor_note.strip()[:600]}"
node.config.context = ((node.config.context or "") + "\n\n" + sup).strip()
def prepare_layer_node(
node: "TaskNode",
*,
completed: dict[str, "AgentResult"],
cancelled: set[str],
node_map: dict[str, "TaskNode"],
on_approval: Callable[["TaskNode"], bool] | None,
) -> "AgentResult | None":
"""Return a skip result, or ``None`` if the node is ready to execute."""
from butler.task_orchestrator import AgentResult
cancelled_dep = first_cancelled_dependency(node, cancelled)
if cancelled_dep:
return AgentResult(
success=False,
error=f"Skipped due to cancelled dependency: {cancelled_dep}",
)
failed_dep = first_failed_dependency(node, completed, node_map)
if failed_dep:
return AgentResult(
success=False,
error=f"Skipped due to failed dependency: {failed_dep}",
)
inject_dependency_context(node, completed)
inject_supervisor_note(node)
if node.requires_approval:
if on_approval is None or not on_approval(node):
return AgentResult(success=False, error="workflow_step_approval_pending")
return None
def split_layer_batches(
layer_tasks: list[tuple[str, "TaskNode"]],
*,
serial: bool,
max_parallel: int | None,
) -> list[list[tuple[str, "TaskNode"]]]:
if serial:
return [[item] for item in layer_tasks]
if max_parallel and max_parallel > 0:
return [
layer_tasks[i : i + max_parallel]
for i in range(0, len(layer_tasks), max_parallel)
]
return [layer_tasks]
def apply_node_router(
nid: str,
node: "TaskNode",
result: "AgentResult",
*,
node_map: dict[str, "TaskNode"],
cancelled: set[str],
) -> list[str]:
"""Apply optional router; return human-readable error strings."""
if not node.router or not result.success:
return []
errors: list[str] = []
from butler.dag_scheduler_ops import run_node_router_safe
next_id, router_error = run_node_router_safe(node.router, result)
if router_error:
errors.append(f"Router for {nid!r} failed: {router_error}")
cancel_direct_dependents(node_map, nid, cancelled)
return errors
if not next_id:
return errors
if next_id not in node_map:
errors.append(f"Router for {nid!r} returned unknown target {next_id!r}")
cancel_direct_dependents(node_map, nid, cancelled)
return errors
deps = direct_dependents(node_map, nid)
if next_id not in {dependent.id for dependent in deps}:
errors.append(
f"Router for {nid!r} returned non-direct dependent {next_id!r}"
)
cancel_direct_dependents(node_map, nid, cancelled)
return errors
for dependent in deps:
if dependent.id != next_id:
cancelled.add(dependent.id)
return errors
def finalize_unexecuted_nodes(
order: list[str],
*,
completed: dict[str, "AgentResult"],
cancelled: set[str],
node_map: dict[str, "TaskNode"],
errors: list[str],
) -> None:
from butler.task_orchestrator import AgentResult
for node_id in order:
if node_id in completed or node_id in cancelled:
continue
cancelled_dep = first_cancelled_dependency(node_map[node_id], cancelled)
if cancelled_dep:
completed[node_id] = AgentResult(
success=False,
error=f"Skipped due to cancelled dependency: {cancelled_dep}",
)
else:
errors.append(f"Node {node_id!r} was not executed")
__all__ = [
"apply_node_router",
"cancel_direct_dependents",
"direct_dependents",
"finalize_unexecuted_nodes",
"first_cancelled_dependency",
"first_failed_dependency",
"format_dependency_context",
"graph_all_required_ok",
"group_into_layers",
"inject_dependency_context",
"inject_supervisor_note",
"prepare_layer_node",
"split_layer_batches",
"topological_sort",
]