from __future__ import annotations import asyncio import logging from dataclasses import dataclass, field from typing import Any, Dict, Set from velai.dataflow.enums import ConnectMultiplicity, DataPortState, NodeState, PortConnectionState from velai.dataflow.graph import DataGraph from velai.dataflow.nodes import NodeInstance from velai.nodes.base_node import BaseNode, BaseNodeData from velai.ui.node_progress import ProgressStateUpdater from velai.ui.vueflow_canvas import VueFlowCanvas logger = logging.getLogger(__name__) @dataclass class ExecutionGraph: root_node: NodeInstance nodes: Dict[str, NodeInstance] = field(default_factory=dict) # node_id -> set of upstream node_ids that must finish before this node starts dependencies: Dict[str, Set[str]] = field(default_factory=dict) nodes_to_execute: Set[str] = field(default_factory=set) class _ProgressManager: def __init__(self, canvas: VueFlowCanvas, updater: ProgressStateUpdater) -> None: self._canvas = canvas self._updater = updater self._stop_events: dict[str, asyncio.Event] = {} self._tasks: dict[str, asyncio.Task[None]] = {} def start(self, node: BaseNode[BaseNodeData]) -> None: node_id = node.node_id if node_id in self._tasks: return stop_event = asyncio.Event() self._stop_events[node_id] = stop_event self._tasks[node_id] = asyncio.create_task(self._poll(node, stop_event)) async def stop(self, node_id: str) -> None: stop_event = self._stop_events.pop(node_id, None) if stop_event is not None: stop_event.set() task = self._tasks.pop(node_id, None) if task is None: return task.cancel() try: await task except asyncio.CancelledError: pass async def stop_all(self) -> None: for event in self._stop_events.values(): event.set() self._stop_events.clear() tasks = list(self._tasks.values()) self._tasks.clear() for task in tasks: task.cancel() for task in tasks: try: await task except asyncio.CancelledError: pass async def _poll(self, node: BaseNode[BaseNodeData], stop_event: asyncio.Event) -> None: last_value: Any = object() last_message: Any = object() while not stop_event.is_set(): await asyncio.sleep(0.1) current_value = node.data.progress_value current_message = node.data.progress_message if current_value != last_value or current_message != last_message: await self._updater.update(node) last_value = current_value last_message = current_message @dataclass class GraphRuntime: graph: DataGraph canvas: VueFlowCanvas _node_tasks: Dict[str, asyncio.Task[None]] = field(default_factory=dict, init=False) _progress: _ProgressManager = field(init=False) _progress_updater: ProgressStateUpdater = field(init=False) def __post_init__(self) -> None: self._progress_updater = ProgressStateUpdater(self.canvas) self._progress = _ProgressManager(self.canvas, self._progress_updater) def _resolve_node(self, node: NodeInstance | str) -> NodeInstance | None: if isinstance(node, str): return self.graph.nodes.get(node) return node def _create_execution_graph(self, root: NodeInstance) -> ExecutionGraph: """Create an execution graph for the given root node. Also checks for circular dependencies. """ exe_graph = ExecutionGraph(root_node=root) visited: Dict[str, int] = {} # node_id -> status (0=visiting, 1=visited) def visit(node: NodeInstance): node_id = node.node_id if node_id in visited: if visited[node_id] == 0: raise ValueError(f"Circular dependency detected at node {node_id}") return visited[node_id] = 0 exe_graph.nodes[node_id] = node exe_graph.dependencies[node_id] = set() for conn in self.graph.upstream_of(node): upstream_node = conn.start_node exe_graph.dependencies[node_id].add(upstream_node.node_id) visit(upstream_node) visited[node_id] = 1 visit(root) return exe_graph def _compute_execution_plan(self, exe_graph: ExecutionGraph) -> None: # 1. Identify initially dirty nodes relevant to this execution dirty_nodes = set() for node in exe_graph.nodes.values(): is_relevant_dirty = False for output in node.all_outputs(): if output.state == DataPortState.DIRTY: # Check if this output is connected to any node in exe_graph for conn in self.graph.downstream_of(node): if conn.start_port is output and conn.end_node.node_id in exe_graph.nodes: is_relevant_dirty = True break if is_relevant_dirty: break if is_relevant_dirty: dirty_nodes.add(node.node_id) # 2. Propagate downstream # Build local downstream map downstream = {nid: set() for nid in exe_graph.nodes} for nid, deps in exe_graph.dependencies.items(): for dep in deps: if dep in downstream: downstream[dep].add(nid) queue = list(dirty_nodes) exe_graph.nodes_to_execute = set(dirty_nodes) while queue: nid = queue.pop(0) for child_id in downstream.get(nid, []): if child_id not in exe_graph.nodes_to_execute: exe_graph.nodes_to_execute.add(child_id) queue.append(child_id) async def execute_node(self, node: NodeInstance | str) -> None: """Execute a node and its dependencies in parallel where possible.""" node_obj = self._resolve_node(node) if node_obj is None: return logger.info("Runtime: Preparing execution for %s...", node_obj.node_id) # If already running or queued, wait for it to finish first waited = False while node_obj.node_id in self._node_tasks or node_obj.state in (NodeState.QUEUED, NodeState.PROCESSING): if node_obj.node_id in self._node_tasks: await self._node_tasks[node_obj.node_id] else: await asyncio.sleep(0.1) waited = True try: # 1. Create Execution Graph exe_graph = self._create_execution_graph(node_obj) # 2. Compute Execution Plan self._compute_execution_plan(exe_graph) # 3. Check if we need to run if not exe_graph.nodes_to_execute: if waited: # If we waited and it's clean, assume previous run satisfied request return else: # Force run if not waited (user clicked execute on clean node) exe_graph.nodes_to_execute.add(node_obj.node_id) # 4. Reset outputs for nodes that will run (specifically the root if forced) if node_obj.node_id in exe_graph.nodes_to_execute: node_obj.reset_outputs() # 5. Mark all nodes that will be executed as QUEUED for node_id in exe_graph.nodes_to_execute: exe_node = exe_graph.nodes[node_id] if exe_node.state not in (NodeState.PROCESSING, NodeState.QUEUED): exe_node.state = NodeState.QUEUED if isinstance(exe_node, BaseNode): await exe_node.on_queue_for_execution() await self._progress_updater.update(exe_node) # 6. Run Execution Graph await self._run_execution_graph(exe_graph) except Exception as e: logger.exception("Runtime execution failed: %s", e) # Cleanup any nodes that are still in QUEUED state due to this failure for exe_node in exe_graph.nodes.values(): if exe_node.state == NodeState.QUEUED and exe_node.node_id not in self._node_tasks: exe_node.state = NodeState.IDLE await self._progress_updater.update(exe_node) async def _run_execution_graph(self, exe_graph: ExecutionGraph) -> None: """Execute nodes in the execution graph topologically and in parallel.""" # Use a local cache of tasks to avoid re-executing the same node in this run # but also coordinate with the global _node_tasks. local_tasks: Dict[str, asyncio.Task[None]] = {} async def run_node_recursive(node_id: str): if node_id in local_tasks: await local_tasks[node_id] return node = exe_graph.nodes[node_id] # 1. Ensure all dependencies are executed first # We run them in parallel dep_tasks = [asyncio.create_task(run_node_recursive(dep_id)) for dep_id in exe_graph.dependencies[node_id]] if dep_tasks: await asyncio.gather(*dep_tasks) # 2. Execute this node (coordinated globally) # We wrap this in a local task so others in this run can wait for it local_tasks[node_id] = asyncio.create_task(self._ensure_node_executed(node, exe_graph)) await local_tasks[node_id] await run_node_recursive(exe_graph.root_node.node_id) async def _ensure_node_executed(self, node: NodeInstance, exe_graph: ExecutionGraph) -> None: node_id = node.node_id # If already being executed globally, just wait for it if node_id in self._node_tasks: logger.debug("Node %s is already executing, waiting...", node_id) await self._node_tasks[node_id] return # Check if node actually needs execution according to plan if node_id not in exe_graph.nodes_to_execute: logger.debug("Node %s not in execution plan, skipping.", node_id) return # Double check: if node became clean (executed by another task) and is not the forced root, skip it if node_id != exe_graph.root_node.node_id and not node.has_dirty_outputs(): logger.debug("Node %s became clean, skipping.", node_id) return # Create a new global task for this node task = asyncio.create_task(self._execute_single_node(node)) self._node_tasks[node_id] = task try: await task finally: self._node_tasks.pop(node_id, None) async def _execute_single_node(self, node: NodeInstance) -> None: """The actual execution of a single node.""" node_id = node.node_id logger.info("Runtime: Executing node %s", node_id) # 1. Update status to QUEUED node.state = NodeState.QUEUED if isinstance(node, BaseNode): await node.on_queue_for_execution() await self._progress_updater.update(node) # 2. Pull values from upstream self._pull_upstream_values(node) # 3. Update status to PROCESSING node.state = NodeState.PROCESSING if isinstance(node, BaseNode): self._progress.start(node) await self._progress_updater.update(node) try: # 4. Run process await node.process() # 5. Push values to downstream (mark them dirty) self._push_downstream_values(node) except Exception as e: logger.exception("Node %s failed: %s", node_id, e) # In case of failure, we might want to set error message if it's a BaseNode if isinstance(node, BaseNode): node.data.error_message = str(e) raise finally: # 6. Update status to IDLE node.state = NodeState.IDLE # 7. UI Cleanup await self._progress.stop(node_id) await self._progress_updater.update(node) await self.sync_node_to_ui(node) def _pull_upstream_values(self, node: NodeInstance) -> None: incoming_conns = list(self.graph.upstream_of(node)) for inp in node.all_inputs(): # Find all connections feeding this port feeds = [c for c in incoming_conns if c.end_port is inp] if not feeds: # No connections to this port inp.value = None if inp.schema.multiplicity == ConnectMultiplicity.SINGLE else [] elif inp.schema.multiplicity == ConnectMultiplicity.MULTIPLE: values = [] for c in feeds: if c.start_port.value is not None: values.append(c.start_port.value) inp.value = values else: # Single connection inp.value = feeds[-1].start_port.value inp.state = DataPortState.CLEAN inp.connection_state = PortConnectionState.CONNECTED if feeds else PortConnectionState.DISCONNECTED def _push_downstream_values(self, node: NodeInstance) -> None: for outp in node.all_outputs(): outp.state = DataPortState.CLEAN for c in self.graph.downstream_of(node): # For now, we always mark downstream nodes as dirty when an upstream node finishes. # A more optimized version would check if the output value actually changed. c.end_port.state = DataPortState.DIRTY c.end_node.mark_dirty() async def sync_node_to_ui(self, node: NodeInstance) -> None: """Push relevant node state back to the Vue nodes.""" if not isinstance(node, BaseNode): return values = node.get_state() if values: await self.canvas.update_node_values(node.node_id, values)