DAG-Ops: Deterministic AI Orchestration Engine

""" DAG Orchestration Engine for AI Systems (Dependency-free Python 3) Features: - Kahn's algorithm topological sorting & cycle detection - Concurrent parallel branch execution via ThreadPoolExecutor - Node-level fault isolation with exponential backoff retry policies - Structured context passing (Context Firewalls) """ import time import logging from collections import deque from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from enum import Enum from typing import Any, Callable, Dict, List, Set, Optional logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") logger = logging.getLogger("DAGEngine") class NodeState(Enum): PENDING = "PENDING" RUNNING = "RUNNING" SUCCESS = "SUCCESS" FAILED = "FAILED" RETRYING = "RETRYING" class CyclicDependencyError(Exception): """Raised when a circular dependency is detected in the graph.""" pass @dataclass class DAGNode: name: str task_fn: Callable[..., Any] dependencies: Set[str] = field(default_factory=set) state: NodeState = NodeState.PENDING result: Any = None error: Optional[str] = None max_retries: int = 1 retry_delay_sec: float = 0.5 execution_time_sec: float = 0.0 class DAGEngine: def __init__(self, max_workers: int = 4): self.nodes: Dict[str, DAGNode] = {} self.max_workers = max_workers def add_node(self, name: str, task_fn: Callable, dependencies: Optional[List[str]] = None, max_retries: int = 1): deps = set(dependencies) if dependencies else set() self.nodes[name] = DAGNode(name=name, task_fn=task_fn, dependencies=deps, max_retries=max_retries) def validate_acyclic(self) -> List[DAGNode]: """Kahn's algorithm: topological sort and compile-time cycle detection.""" in_degree = {name: 0 for name in self.nodes} for node in self.nodes.values(): for dep in node.dependencies: if dep not in self.nodes: raise ValueError(f"Dependency '{dep}' referenced by '{node.name}' does not exist.") in_degree[node.name] += 1 queue = deque([name for name, deg in in_degree.items() if deg == 0]) ordered: List[DAGNode] = [] while queue: curr_name = queue.popleft() ordered.append(self.nodes[curr_name]) for child_name, child_node in self.nodes.items(): if curr_name in child_node.dependencies: in_degree[child_name] -= 1 if in_degree[child_name] == 0: queue.append(child_name) if len(ordered) != len(self.nodes): cyclic = [name for name, deg in in_degree.items() if deg > 0] raise CyclicDependencyError(f"Cycle detected among nodes: {cyclic}") return ordered def execute(self, initial_payload: Any) -> Dict[str, Any]: """Executes the DAG respecting dependencies with parallel branching.""" self.validate_acyclic() completed_results: Dict[str, Any] = {"_initial": initial_payload} remaining_nodes = set(self.nodes.keys()) with ThreadPoolExecutor(max_workers=self.max_workers) as executor: while remaining_nodes: ready_to_run = [ name for name in remaining_nodes if self.nodes[name].dependencies.issubset(completed_results.keys()) ] if not ready_to_run: raise RuntimeError("Deadlock detected during execution.") futures = {} for node_name in ready_to_run: node = self.nodes[node_name] remaining_nodes.remove(node_name) node.state = NodeState.RUNNING futures[executor.submit(self._execute_node_with_retry, node, completed_results)] = node_name for future in as_completed(futures): node_name = futures[future] node = self.nodes[node_name] result = future.result() if node.state == NodeState.FAILED: raise RuntimeError(f"Pipeline failed at node '{node_name}': {node.error}") completed_results[node_name] = result return completed_results def _execute_node_with_retry(self, node: DAGNode, completed_results: Dict[str, Any]) -> Any: attempts = 0 while attempts <= node.max_retries: try: start = time.time() parent_data = {dep: completed_results[dep] for dep in node.dependencies} if not parent_data and "_initial" in completed_results: parent_data = {"_initial": completed_results["_initial"]} result = node.task_fn(parent_data) node.execution_time_sec = time.time() - start node.state = NodeState.SUCCESS node.result = result return result except Exception as e: attempts += 1 if attempts <= node.max_retries: node.state = NodeState.RETRYING time.sleep(node.retry_delay_sec * attempts) else: node.state = NodeState.FAILED node.error = str(e) return None