chore: import upstream snapshot with attribution
Validate YAML Workflows / Validate YAML Configuration Files (push) Waiting to run
Validate YAML Workflows / Validate YAML Configuration Files (push) Waiting to run
This commit is contained in:
wehub-resource-sync
Executable
Executable
+614
@@ -0,0 +1,614 @@
|
||||
"""Cycle executor that runs workflow graphs containing loops."""
|
||||
|
||||
import copy
|
||||
import threading
|
||||
from typing import Dict, List, Callable, Any, Set, Optional
|
||||
|
||||
from entity.configs import Node, EdgeLink
|
||||
from utils.log_manager import LogManager
|
||||
from workflow.cycle_manager import CycleManager
|
||||
from workflow.executor.parallel_executor import ParallelExecutor
|
||||
from workflow.topology_builder import GraphTopologyBuilder
|
||||
|
||||
|
||||
class CycleExecutor:
|
||||
"""Execute workflow graphs that contain cycles.
|
||||
|
||||
Features:
|
||||
- Scheduling is based on "super nodes"
|
||||
- Parallel execution inside cycles
|
||||
- Automatic detection of exit conditions
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
log_manager: LogManager,
|
||||
nodes: Dict[str, Node],
|
||||
cycle_execution_order: List[Dict[str, Any]],
|
||||
cycle_manager: CycleManager,
|
||||
execute_node_func: Callable[[Node], None],
|
||||
):
|
||||
"""Initialize the cycle executor.
|
||||
|
||||
Args:
|
||||
log_manager: Logger instance
|
||||
nodes: Mapping of node ids to nodes
|
||||
cycle_execution_order: Super-node execution order with cycles
|
||||
cycle_manager: Cycle manager coordinating iterations
|
||||
execute_node_func: Callable that executes a single node
|
||||
"""
|
||||
self.log_manager = log_manager
|
||||
self.nodes = nodes
|
||||
self.cycle_execution_order = cycle_execution_order
|
||||
self.cycle_manager = cycle_manager
|
||||
self.execute_node_func = execute_node_func
|
||||
self.parallel_executor = ParallelExecutor(log_manager, nodes)
|
||||
|
||||
def execute(self) -> None:
|
||||
"""Run the workflow that contains cycles."""
|
||||
self.log_manager.debug("Executing graph with cycles using super-node scheduler")
|
||||
|
||||
for layer_idx, layer_items in enumerate(self.cycle_execution_order):
|
||||
self.log_manager.debug(f"Executing super-node layer {layer_idx} with {len(layer_items)} items")
|
||||
self._execute_super_layer(layer_items)
|
||||
|
||||
def _execute_super_layer(self, layer_items: List[Dict[str, Any]]) -> None:
|
||||
"""Execute a single super-node layer."""
|
||||
self._execute_super_layer_parallel(layer_items)
|
||||
|
||||
def _execute_super_layer_parallel(self, layer_items: List[Dict[str, Any]]) -> None:
|
||||
"""Execute a super-node layer in parallel."""
|
||||
def item_desc_func(item: Dict[str, Any]) -> str:
|
||||
if item["type"] == "cycle":
|
||||
return f"cycle {item['cycle_id']}"
|
||||
elif item["type"] == "node":
|
||||
# New format
|
||||
return f"node {item['node_id']}"
|
||||
else:
|
||||
# Old format: "layer"
|
||||
return f"node {item['nodes'][0]}"
|
||||
|
||||
self.parallel_executor.execute_items_parallel(
|
||||
layer_items,
|
||||
self._execute_super_item,
|
||||
item_desc_func
|
||||
)
|
||||
|
||||
def _execute_super_item(self, item: Dict[str, Any]) -> None:
|
||||
"""Execute a single super-node item (node or cycle)."""
|
||||
if item["type"] == "layer":
|
||||
# Old format: {"type": "layer", "nodes": [node_id]}
|
||||
self._execute_single_node(item["nodes"][0])
|
||||
elif item["type"] == "node":
|
||||
# New format from GraphTopologyBuilder: {"type": "node", "node_id": "..."}
|
||||
self._execute_single_node(item["node_id"])
|
||||
elif item["type"] == "cycle":
|
||||
self._execute_cycle(item)
|
||||
|
||||
def _execute_single_node(self, node_id: str) -> None:
|
||||
"""Execute a non-cycle node."""
|
||||
self.log_manager.debug(f"Executing non-cycle node: {node_id}")
|
||||
|
||||
node = self.nodes[node_id]
|
||||
if node.is_triggered():
|
||||
self.execute_node_func(node)
|
||||
else:
|
||||
self.log_manager.warning(f"Node {node_id} is not triggered, skipping execution")
|
||||
|
||||
def _execute_cycle(self, cycle_info: Dict[str, Any]) -> None:
|
||||
"""Execute a cycle using the multi-iteration logic."""
|
||||
cycle_id = cycle_info["cycle_id"]
|
||||
nodes = cycle_info["nodes"]
|
||||
|
||||
self.log_manager.debug(f"Executing cycle {cycle_id} with nodes: {nodes}")
|
||||
|
||||
# Step 2: Validate cycle entry uniqueness
|
||||
try:
|
||||
initial_node_id = self._validate_cycle_entry(cycle_id, nodes)
|
||||
except ValueError as e:
|
||||
self.log_manager.error(str(e))
|
||||
raise
|
||||
|
||||
if initial_node_id is None:
|
||||
self.log_manager.debug(
|
||||
f"Cycle {cycle_id} has no triggered entry node in this pass; skipping execution"
|
||||
)
|
||||
return
|
||||
|
||||
# Store initial node in cycle_manager
|
||||
self.cycle_manager.cycles[cycle_id].initial_node = initial_node_id
|
||||
self.log_manager.debug(f"Cycle {cycle_id} initial node: {initial_node_id}")
|
||||
|
||||
# Activate cycle
|
||||
self.cycle_manager.activate_cycle(cycle_id)
|
||||
|
||||
# Step 4: Execute cycle with iterations
|
||||
self._execute_cycle_with_iterations(
|
||||
cycle_id,
|
||||
nodes,
|
||||
initial_node_id,
|
||||
max_iterations=self.cycle_manager.cycles[cycle_id].get_max_iterations()
|
||||
)
|
||||
|
||||
# Cleanup
|
||||
self.cycle_manager.deactivate_cycle(cycle_id)
|
||||
self.log_manager.debug(f"Cycle {cycle_id} completed")
|
||||
|
||||
# ==================== New Methods for Refactored Cycle Execution ====================
|
||||
|
||||
def _validate_cycle_entry(self, cycle_id: str, nodes: List[str]) -> str | None:
|
||||
"""
|
||||
Validate that exactly one node in the cycle is triggered by external edges.
|
||||
|
||||
Args:
|
||||
cycle_id: The cycle ID
|
||||
nodes: List of node IDs in the cycle
|
||||
|
||||
Returns:
|
||||
The ID of the unique initial node
|
||||
|
||||
Raises:
|
||||
ValueError: If no node or multiple nodes are triggered
|
||||
"""
|
||||
triggered_nodes: List[str] = []
|
||||
|
||||
for node_id in nodes:
|
||||
node = self.nodes[node_id]
|
||||
# Check if any external predecessor (node outside the cycle) triggers this node
|
||||
for predecessor in node.predecessors:
|
||||
if predecessor.id not in nodes: # External node
|
||||
edge = predecessor.find_outgoing_edge(node_id)
|
||||
if edge and edge.trigger and edge.triggered:
|
||||
triggered_nodes.append(node_id)
|
||||
break
|
||||
|
||||
cycle_info = self.cycle_manager.cycles.get(cycle_id)
|
||||
configured_entry = cycle_info.configured_entry_node if cycle_info else None
|
||||
|
||||
if len(triggered_nodes) == 0:
|
||||
if configured_entry:
|
||||
return configured_entry
|
||||
return None
|
||||
elif len(triggered_nodes) > 1:
|
||||
raise ValueError(
|
||||
f"Cycle {cycle_id} has multiple triggered entry nodes: {triggered_nodes}. "
|
||||
"Only one entry node must be triggered when entering a cycle."
|
||||
)
|
||||
entry_node = triggered_nodes[0]
|
||||
if configured_entry and entry_node != configured_entry:
|
||||
raise ValueError(
|
||||
f"Cycle {cycle_id} entry mismatch: configured '{configured_entry}' "
|
||||
f"but triggered '{entry_node}'",
|
||||
)
|
||||
|
||||
return entry_node
|
||||
|
||||
def _execute_cycle_with_iterations(
|
||||
self,
|
||||
cycle_id: str,
|
||||
cycle_nodes: List[str],
|
||||
initial_node_id: str,
|
||||
max_iterations: int,
|
||||
) -> Set[str]:
|
||||
"""
|
||||
Execute a cycle with multiple iterations.
|
||||
|
||||
Args:
|
||||
cycle_id: Cycle ID
|
||||
cycle_nodes: List of all nodes in the cycle
|
||||
initial_node_id: Initial node ID
|
||||
max_iterations: Maximum number of iterations
|
||||
|
||||
Returns:
|
||||
A tuple of two sets:
|
||||
- exit_nodes: nodes triggered outside the *current* cycle scope
|
||||
- external_nodes: subset of exit_nodes that are also outside the
|
||||
provided parent_cycle_nodes scope
|
||||
"""
|
||||
iteration = 0
|
||||
|
||||
while iteration < max_iterations:
|
||||
self.log_manager.debug(
|
||||
f"Cycle {cycle_id} iteration {iteration + 1}/{max_iterations}"
|
||||
)
|
||||
|
||||
# Step 1: Detect nested cycles in the scoped subgraph
|
||||
inner_cycles = self._detect_cycles_in_scope(cycle_nodes, initial_node_id)
|
||||
|
||||
# Build topological layers (whether there are nested cycles or not)
|
||||
execution_layers = self._build_topological_layers_in_scope(
|
||||
cycle_nodes, initial_node_id, inner_cycles,
|
||||
is_first_iteration=(iteration == 0)
|
||||
)
|
||||
|
||||
# Execute the topological layers
|
||||
external_nodes = self._execute_scope_layers(
|
||||
execution_layers,
|
||||
cycle_id,
|
||||
cycle_nodes,
|
||||
initial_node_id=initial_node_id,
|
||||
is_first_iteration=(iteration == 0)
|
||||
)
|
||||
|
||||
if external_nodes:
|
||||
self.log_manager.debug(
|
||||
f"Cycle {cycle_id} exited - external nodes triggered: {sorted(external_nodes)}"
|
||||
)
|
||||
return external_nodes
|
||||
|
||||
# Step 4: Check if initial node is retriggered
|
||||
if not self._is_initial_node_retriggered(initial_node_id, cycle_nodes):
|
||||
self.log_manager.debug(
|
||||
f"Cycle {cycle_id} completed - initial node not retriggered"
|
||||
)
|
||||
break
|
||||
|
||||
iteration += 1
|
||||
|
||||
if iteration >= max_iterations:
|
||||
self.log_manager.warning(
|
||||
f"Cycle {cycle_id} reached max iterations ({max_iterations})"
|
||||
)
|
||||
return set()
|
||||
def _detect_cycles_in_scope(
|
||||
self,
|
||||
scope_nodes: List[str],
|
||||
initial_node_id: str
|
||||
) -> List[Set[str]]:
|
||||
"""
|
||||
Detect nested cycles within the scoped subgraph.
|
||||
|
||||
Constructs a subgraph containing only:
|
||||
1. Nodes in scope_nodes
|
||||
2. Edges where both source and target are in scope_nodes
|
||||
3. Initial node's incoming edges are REMOVED (to break the outer cycle)
|
||||
|
||||
Args:
|
||||
scope_nodes: List of node IDs in the current scope
|
||||
initial_node_id: Initial node ID (whose incoming edges are removed)
|
||||
|
||||
Returns:
|
||||
List of detected nested cycles (excluding the current cycle itself)
|
||||
"""
|
||||
# Build scoped nodes with initial node's incoming edges removed
|
||||
scoped_nodes = self._build_scoped_nodes(scope_nodes, clear_entry_node=initial_node_id)
|
||||
|
||||
# Use GraphTopologyBuilder to detect cycles
|
||||
all_cycles = GraphTopologyBuilder.detect_cycles(scoped_nodes)
|
||||
|
||||
# Filter out single-node "cycles" (unless they have self-loops)
|
||||
nested_cycles = [
|
||||
cycle for cycle in all_cycles
|
||||
if len(cycle) > 1
|
||||
]
|
||||
|
||||
return nested_cycles
|
||||
|
||||
def _build_scoped_nodes(
|
||||
self,
|
||||
scope_nodes: List[str],
|
||||
clear_entry_node: Optional[str] = None
|
||||
) -> Dict[str, Node]:
|
||||
"""
|
||||
Build a scoped subgraph containing only nodes and edges within the scope.
|
||||
|
||||
Args:
|
||||
scope_nodes: List of node IDs in the scope
|
||||
clear_entry_node: If specified, this node's incoming edges will be removed
|
||||
(used to break the outer cycle when detecting nested cycles)
|
||||
|
||||
Returns:
|
||||
Dictionary of scoped nodes
|
||||
"""
|
||||
scoped_nodes = {}
|
||||
scope_nodes_set = set(scope_nodes)
|
||||
|
||||
for node_id in scope_nodes:
|
||||
original_node = self.nodes[node_id]
|
||||
# Shallow copy the node
|
||||
scoped_node = copy.copy(original_node)
|
||||
|
||||
# Filter outgoing edges: only keep edges where target is in scope AND trigger=true
|
||||
# Special case: if target is clear_entry_node, remove this edge
|
||||
scoped_edges = [
|
||||
edge_link for edge_link in original_node.iter_outgoing_edges()
|
||||
if edge_link.target.id in scope_nodes_set
|
||||
and edge_link.trigger
|
||||
and edge_link.target.id != clear_entry_node # Remove edges to entry node
|
||||
]
|
||||
scoped_node._outgoing_edges = scoped_edges
|
||||
|
||||
# Filter predecessors: only keep predecessors in scope AND with trigger=true edge
|
||||
# Special case: if this node is clear_entry_node, clear all predecessors
|
||||
if node_id == clear_entry_node:
|
||||
scoped_node.predecessors = []
|
||||
else:
|
||||
scoped_predecessors = []
|
||||
for pred in original_node.predecessors:
|
||||
if pred.id in scope_nodes_set:
|
||||
# Check if the edge from pred to node has trigger=true
|
||||
edge = pred.find_outgoing_edge(node_id)
|
||||
if edge and edge.trigger:
|
||||
scoped_predecessors.append(pred)
|
||||
scoped_node.predecessors = scoped_predecessors
|
||||
|
||||
# Filter successors: only keep successors in scope AND with trigger=true edge
|
||||
# Special case: remove clear_entry_node from successors
|
||||
scoped_successors = [
|
||||
succ for succ in original_node.successors
|
||||
if succ.id in scope_nodes_set
|
||||
and succ.id != clear_entry_node # Remove entry node from successors
|
||||
and any(
|
||||
edge_link.target.id == succ.id and edge_link.trigger
|
||||
for edge_link in original_node.iter_outgoing_edges()
|
||||
)
|
||||
]
|
||||
scoped_node.successors = scoped_successors
|
||||
|
||||
scoped_nodes[node_id] = scoped_node
|
||||
|
||||
return scoped_nodes
|
||||
|
||||
def _build_topological_layers_in_scope(
|
||||
self,
|
||||
scope_nodes: List[str],
|
||||
initial_node_id: str,
|
||||
inner_cycles: List[Set[str]],
|
||||
is_first_iteration: bool = False
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Build topological execution order for the scoped subgraph.
|
||||
|
||||
Args:
|
||||
scope_nodes: List of node IDs in the scope
|
||||
initial_node_id: Initial node ID
|
||||
inner_cycles: List of nested cycles detected in the scope
|
||||
is_first_iteration: Whether this is the first iteration (affects initial node handling)
|
||||
|
||||
Returns:
|
||||
List of execution layers, each containing execution items
|
||||
"""
|
||||
# Build scoped nodes WITHOUT clearing entry node
|
||||
# We want to keep all edges intact for execution
|
||||
scoped_nodes = self._build_scoped_nodes(scope_nodes, clear_entry_node=None)
|
||||
|
||||
# Handle entry points based on iteration:
|
||||
# - First iteration: manually clear initial node's predecessors (for in_degree calculation only)
|
||||
# - Subsequent iterations: clear predecessors for all triggered nodes
|
||||
if is_first_iteration:
|
||||
# Clear initial node's predecessors to make it an entry point
|
||||
if initial_node_id in scoped_nodes:
|
||||
scoped_nodes[initial_node_id].predecessors = []
|
||||
else:
|
||||
# Subsequent iterations: clear predecessors for all triggered nodes
|
||||
for node_id in scope_nodes:
|
||||
if self.nodes[node_id].is_triggered():
|
||||
scoped_nodes[node_id].predecessors = []
|
||||
|
||||
# Extract scoped edges from scoped_nodes (not original nodes)
|
||||
# This ensures consistency with the filtered graph structure
|
||||
scoped_edges = []
|
||||
|
||||
# Collect nodes whose incoming edges should be excluded
|
||||
# (to break cycles in topological sorting)
|
||||
exclude_targets = set()
|
||||
if is_first_iteration:
|
||||
# First iteration: exclude edges to initial_node
|
||||
exclude_targets.add(initial_node_id)
|
||||
else:
|
||||
# Subsequent iterations: exclude edges to all triggered nodes
|
||||
for node_id in scope_nodes:
|
||||
if self.nodes[node_id].is_triggered():
|
||||
exclude_targets.add(node_id)
|
||||
|
||||
for node_id in scope_nodes:
|
||||
# Use scoped_node to get filtered edges
|
||||
scoped_node = scoped_nodes.get(node_id)
|
||||
if scoped_node:
|
||||
for edge_link in scoped_node.iter_outgoing_edges():
|
||||
# Exclude edges pointing to nodes in exclude_targets
|
||||
if edge_link.target.id in exclude_targets:
|
||||
continue
|
||||
scoped_edges.append({
|
||||
"from": node_id,
|
||||
"to": edge_link.target.id,
|
||||
"trigger": edge_link.trigger,
|
||||
"condition": edge_link.condition
|
||||
})
|
||||
|
||||
# Use GraphTopologyBuilder to build execution order
|
||||
if not inner_cycles:
|
||||
# No nested cycles, use DAG layers
|
||||
layers = GraphTopologyBuilder.build_dag_layers(scoped_nodes)
|
||||
return layers
|
||||
else:
|
||||
# Has nested cycles, use super-node approach
|
||||
super_graph = GraphTopologyBuilder.create_super_node_graph(
|
||||
scoped_nodes, scoped_edges, inner_cycles
|
||||
)
|
||||
layers = GraphTopologyBuilder.topological_sort_super_nodes(
|
||||
super_graph, inner_cycles
|
||||
)
|
||||
return layers
|
||||
|
||||
def _execute_scope_layers(
|
||||
self,
|
||||
execution_layers: List[List[Dict[str, Any]]],
|
||||
parent_cycle_id: str,
|
||||
parent_cycle_nodes: List[str],
|
||||
initial_node_id: Optional[str] = None,
|
||||
is_first_iteration: bool = False
|
||||
) -> Set[str]:
|
||||
"""
|
||||
Execute scoped layers with parallelism, supporting nested cycles.
|
||||
|
||||
Args:
|
||||
execution_layers: List of execution layers
|
||||
parent_cycle_id: Parent cycle ID
|
||||
parent_cycle_nodes: List of nodes in the parent cycle
|
||||
initial_node_id: Initial node ID (for first iteration special handling)
|
||||
is_first_iteration: Whether this is the first iteration
|
||||
|
||||
Returns:
|
||||
external_nodes: subset of exit_nodes outside parent_cycle_nodes_set
|
||||
"""
|
||||
scope_node_set = set(parent_cycle_nodes)
|
||||
external_nodes: Set[str] = set()
|
||||
stop_event = threading.Event()
|
||||
result_lock = threading.Lock()
|
||||
|
||||
def record_external(nodes: Set[str]) -> None:
|
||||
nonlocal external_nodes
|
||||
if not nodes:
|
||||
return
|
||||
with result_lock:
|
||||
if nodes:
|
||||
external_nodes.update(nodes)
|
||||
stop_event.set()
|
||||
|
||||
def item_desc(item: Dict[str, Any]) -> str:
|
||||
if item["type"] == "node":
|
||||
return f"node {item['node_id']}"
|
||||
if item["type"] == "cycle":
|
||||
return f"cycle {item['cycle_id']}"
|
||||
return "layer_item"
|
||||
|
||||
for layer in execution_layers:
|
||||
if stop_event.is_set():
|
||||
break
|
||||
|
||||
def executor_func(item: Dict[str, Any]) -> None:
|
||||
if stop_event.is_set():
|
||||
return
|
||||
|
||||
if item["type"] == "node":
|
||||
_node_id = item["node_id"]
|
||||
force_execute = is_first_iteration and (_node_id == initial_node_id)
|
||||
targets = self._execute_single_cycle_node_in_scope(
|
||||
_node_id,
|
||||
scope_node_set,
|
||||
force_execute=force_execute
|
||||
)
|
||||
if targets:
|
||||
record_external(targets)
|
||||
|
||||
elif item["type"] == "cycle":
|
||||
inner_cycle_nodes = item["nodes"]
|
||||
inner_cycle_id = item["cycle_id"]
|
||||
|
||||
self.log_manager.debug(
|
||||
f"Executing nested cycle {inner_cycle_id} within cycle {parent_cycle_id}"
|
||||
)
|
||||
|
||||
try:
|
||||
inner_initial_node = self._validate_cycle_entry(
|
||||
inner_cycle_id, inner_cycle_nodes
|
||||
)
|
||||
except ValueError as e:
|
||||
self.log_manager.error(str(e))
|
||||
raise
|
||||
|
||||
if inner_initial_node is None:
|
||||
self.log_manager.debug(
|
||||
f"Nested cycle {inner_cycle_id} has no triggered entry; skipping"
|
||||
)
|
||||
return
|
||||
|
||||
inner_external_nodes = self._execute_cycle_with_iterations(
|
||||
inner_cycle_id,
|
||||
inner_cycle_nodes,
|
||||
inner_initial_node,
|
||||
max_iterations=100,
|
||||
)
|
||||
|
||||
if inner_external_nodes:
|
||||
filtered = {
|
||||
node
|
||||
for node in inner_external_nodes
|
||||
if node not in scope_node_set
|
||||
}
|
||||
if filtered:
|
||||
record_external(filtered)
|
||||
|
||||
self.parallel_executor.execute_items_parallel(
|
||||
layer,
|
||||
executor_func,
|
||||
item_desc
|
||||
)
|
||||
|
||||
if stop_event.is_set():
|
||||
break
|
||||
|
||||
if external_nodes:
|
||||
for node_id in scope_node_set:
|
||||
self.nodes[node_id].reset_triggers()
|
||||
|
||||
return external_nodes
|
||||
|
||||
def _execute_single_cycle_node_in_scope(
|
||||
self,
|
||||
node_id: str,
|
||||
scope_node_set: Set[str],
|
||||
force_execute: bool = False
|
||||
) -> Set[str]:
|
||||
"""
|
||||
Execute a single node within a cycle scope.
|
||||
|
||||
Args:
|
||||
node_id: Node ID to execute
|
||||
scope_node_set: Nodes that belong to the current scoped cycle
|
||||
force_execute: If True, execute even if not triggered (for initial node in first iteration)
|
||||
|
||||
Returns:
|
||||
Set of node IDs triggered outside the current scoped cycle
|
||||
"""
|
||||
node = self.nodes[node_id]
|
||||
|
||||
# Check if node is triggered (unless force_execute is True)
|
||||
if not force_execute:
|
||||
if not node.is_triggered():
|
||||
return set()
|
||||
|
||||
# Reset edge triggers
|
||||
for edge_link in node.iter_outgoing_edges():
|
||||
edge_link.triggered = False
|
||||
|
||||
# Execute the node
|
||||
self.execute_node_func(node)
|
||||
|
||||
# Check if any external node was triggered
|
||||
external_targets: Set[str] = set()
|
||||
for edge_link in node.iter_outgoing_edges():
|
||||
if edge_link.target.id not in scope_node_set and edge_link.triggered:
|
||||
self.log_manager.debug(
|
||||
f"Node {node_id} triggered external node {edge_link.target.id}"
|
||||
)
|
||||
external_targets.add(edge_link.target.id)
|
||||
|
||||
return external_targets
|
||||
|
||||
def _is_initial_node_retriggered(
|
||||
self,
|
||||
initial_node_id: str,
|
||||
cycle_nodes: List[str]
|
||||
) -> bool:
|
||||
"""
|
||||
Check if the initial node is retriggered by any internal edge (from within the cycle).
|
||||
|
||||
Args:
|
||||
initial_node_id: Initial node ID
|
||||
cycle_nodes: List of nodes in the cycle
|
||||
|
||||
Returns:
|
||||
True if the initial node is retriggered by an internal edge
|
||||
"""
|
||||
initial_node = self.nodes[initial_node_id]
|
||||
|
||||
for predecessor in initial_node.predecessors:
|
||||
# Only check predecessors within the cycle
|
||||
if predecessor.id in cycle_nodes:
|
||||
edge = predecessor.find_outgoing_edge(initial_node_id)
|
||||
if edge and edge.trigger and edge.triggered:
|
||||
return True
|
||||
|
||||
return False
|
||||
Executable
+55
@@ -0,0 +1,55 @@
|
||||
"""Executor for DAG (Directed Acyclic Graph) workflows."""
|
||||
|
||||
from typing import Dict, List, Callable
|
||||
|
||||
from entity.configs import Node
|
||||
from utils.log_manager import LogManager
|
||||
from workflow.executor.parallel_executor import ParallelExecutor
|
||||
|
||||
|
||||
class DAGExecutor:
|
||||
"""Execute DAG workflows.
|
||||
|
||||
Features:
|
||||
- Execute layer by layer following the topology
|
||||
- Support parallel execution inside a layer
|
||||
- Serialize Human nodes automatically
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
log_manager: LogManager,
|
||||
nodes: Dict[str, Node],
|
||||
layers: List[List[str]],
|
||||
execute_node_func: Callable[[Node], None]
|
||||
):
|
||||
"""Initialize the executor.
|
||||
|
||||
Args:
|
||||
log_manager: Logger instance
|
||||
nodes: Mapping of node ids to ``Node`` objects
|
||||
layers: Topological layers
|
||||
execute_node_func: Callable used to execute a single node
|
||||
"""
|
||||
self.log_manager = log_manager
|
||||
self.nodes = nodes
|
||||
self.layers = layers
|
||||
self.execute_node_func = execute_node_func
|
||||
self.parallel_executor = ParallelExecutor(log_manager, nodes)
|
||||
|
||||
def execute(self) -> None:
|
||||
"""Execute the DAG workflow."""
|
||||
for layer_idx, layer_nodes in enumerate(self.layers):
|
||||
self.log_manager.debug(f"Executing Layer {layer_idx} with nodes: {layer_nodes}")
|
||||
self._execute_layer(layer_nodes)
|
||||
|
||||
def _execute_layer(self, layer_nodes: List[str]) -> None:
|
||||
"""Execute a single topological layer."""
|
||||
def execute_if_triggered(node_id: str) -> None:
|
||||
node = self.nodes[node_id]
|
||||
if node.is_triggered():
|
||||
self.execute_node_func(node)
|
||||
else:
|
||||
self.log_manager.debug(f"Node {node_id} skipped - not triggered")
|
||||
|
||||
self.parallel_executor.execute_nodes_parallel(layer_nodes, execute_if_triggered)
|
||||
Executable
+399
@@ -0,0 +1,399 @@
|
||||
"""Dynamic edge executor for edge-level Map and Tree execution.
|
||||
|
||||
Handles dynamic node expansion based on edge-level dynamic configuration.
|
||||
When a message passes through an edge with dynamic config, the target node
|
||||
is virtually expanded into multiple instances based on split results.
|
||||
"""
|
||||
|
||||
import concurrent.futures
|
||||
from typing import Callable, Dict, List, Optional
|
||||
|
||||
from entity.configs import Node
|
||||
from entity.configs.edge.dynamic_edge_config import DynamicEdgeConfig
|
||||
from entity.messages import Message, MessageRole
|
||||
from runtime.node.splitter import create_splitter_from_config, group_messages
|
||||
from utils.log_manager import LogManager
|
||||
|
||||
|
||||
class DynamicEdgeExecutor:
|
||||
"""Execute edge-level dynamic expansion.
|
||||
|
||||
When an edge has dynamic configuration, this executor:
|
||||
1. Splits the payload passing through the edge
|
||||
2. Executes the target node for each split unit
|
||||
3. Collects and returns results (flat for Map, reduced for Tree)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
log_manager: LogManager,
|
||||
node_executor_func: Callable[[Node, List[Message]], List[Message]],
|
||||
):
|
||||
"""Initialize the dynamic edge executor.
|
||||
|
||||
Args:
|
||||
log_manager: Logger instance
|
||||
node_executor_func: Function to execute a node with inputs
|
||||
"""
|
||||
self.log_manager = log_manager
|
||||
self.node_executor_func = node_executor_func
|
||||
|
||||
def execute(
|
||||
self,
|
||||
target_node: Node,
|
||||
payload: Message,
|
||||
dynamic_config: DynamicEdgeConfig,
|
||||
static_inputs: Optional[List[Message]] = None,
|
||||
) -> List[Message]:
|
||||
"""Execute dynamic expansion for an edge.
|
||||
|
||||
Args:
|
||||
target_node: The node to execute (will be used as template)
|
||||
payload: The message passing through the edge
|
||||
dynamic_config: Edge dynamic configuration
|
||||
static_inputs: Optional static inputs from non-dynamic edges
|
||||
|
||||
Returns:
|
||||
List of output messages from all executions
|
||||
"""
|
||||
split_config = dynamic_config.split
|
||||
|
||||
# Create splitter based on config
|
||||
splitter = create_splitter_from_config(split_config)
|
||||
|
||||
# Split the payload into execution units
|
||||
execution_units = splitter.split([payload])
|
||||
|
||||
if not execution_units:
|
||||
self.log_manager.debug(
|
||||
f"Dynamic edge -> {target_node.id}: no execution units after split"
|
||||
)
|
||||
return []
|
||||
|
||||
self.log_manager.info(
|
||||
f"Dynamic edge -> {target_node.id}: splitting into {len(execution_units)} parallel units"
|
||||
)
|
||||
|
||||
if dynamic_config.is_map():
|
||||
return self._execute_map(
|
||||
target_node, execution_units, dynamic_config, static_inputs
|
||||
)
|
||||
elif dynamic_config.is_tree():
|
||||
return self._execute_tree(
|
||||
target_node, execution_units, dynamic_config, static_inputs
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown dynamic type: {dynamic_config.type}")
|
||||
|
||||
def execute_from_inputs(
|
||||
self,
|
||||
target_node: Node,
|
||||
inputs: List[Message],
|
||||
dynamic_config: DynamicEdgeConfig,
|
||||
static_inputs: Optional[List[Message]] = None,
|
||||
) -> List[Message]:
|
||||
"""Execute dynamic expansion using all collected inputs.
|
||||
|
||||
This method is called from _execute_node when a node has incoming edges
|
||||
with dynamic configuration. All inputs are already collected and passed here.
|
||||
|
||||
Args:
|
||||
target_node: The node to execute
|
||||
inputs: Dynamic edge inputs to be split
|
||||
dynamic_config: Edge dynamic configuration
|
||||
static_inputs: Non-dynamic edge inputs to be replicated to all units
|
||||
|
||||
Returns:
|
||||
List of output messages from all executions
|
||||
"""
|
||||
split_config = dynamic_config.split
|
||||
static_inputs = static_inputs or []
|
||||
|
||||
# Create splitter based on config
|
||||
splitter = create_splitter_from_config(split_config)
|
||||
|
||||
# Split only dynamic inputs into execution units
|
||||
execution_units = splitter.split(inputs)
|
||||
|
||||
if not execution_units:
|
||||
self.log_manager.debug(
|
||||
f"Dynamic node {target_node.id}: no execution units after split"
|
||||
)
|
||||
# If no dynamic inputs but have static inputs, execute once with static inputs
|
||||
if static_inputs:
|
||||
return self.node_executor_func(target_node, static_inputs)
|
||||
return []
|
||||
|
||||
self.log_manager.info(
|
||||
f"Dynamic node {target_node.id}: splitting {len(inputs)} dynamic inputs into "
|
||||
f"{len(execution_units)} parallel units ({dynamic_config.type} mode)"
|
||||
+ (f", with {len(static_inputs)} static inputs replicated to each" if static_inputs else "")
|
||||
)
|
||||
|
||||
if dynamic_config.is_map():
|
||||
return self._execute_map(
|
||||
target_node, execution_units, dynamic_config, static_inputs
|
||||
)
|
||||
elif dynamic_config.is_tree():
|
||||
return self._execute_tree(
|
||||
target_node, execution_units, dynamic_config, static_inputs
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown dynamic type: {dynamic_config.type}")
|
||||
|
||||
def _execute_map(
|
||||
self,
|
||||
target_node: Node,
|
||||
execution_units: List[List[Message]],
|
||||
dynamic_config: DynamicEdgeConfig,
|
||||
static_inputs: Optional[List[Message]] = None,
|
||||
) -> List[Message]:
|
||||
"""Execute in Map mode (fan-out only).
|
||||
|
||||
Args:
|
||||
target_node: Target node template
|
||||
execution_units: Split message units
|
||||
dynamic_config: Dynamic configuration
|
||||
static_inputs: Static inputs to copy to all units
|
||||
|
||||
Returns:
|
||||
Flat list of all output messages
|
||||
"""
|
||||
map_config = dynamic_config.as_map_config()
|
||||
max_parallel = map_config.max_parallel
|
||||
all_outputs: List[Message] = []
|
||||
static_inputs = static_inputs or []
|
||||
|
||||
if len(execution_units) == 1:
|
||||
# Single unit - execute directly
|
||||
unit_inputs = list(static_inputs) + execution_units[0]
|
||||
outputs = self._execute_unit(target_node, unit_inputs, 0)
|
||||
all_outputs.extend(outputs)
|
||||
else:
|
||||
# Multiple units - parallel execution
|
||||
effective_workers = min(len(execution_units), max_parallel)
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=effective_workers) as executor:
|
||||
futures: Dict[concurrent.futures.Future, int] = {}
|
||||
|
||||
for idx, unit in enumerate(execution_units):
|
||||
unit_inputs = list(static_inputs) + unit
|
||||
future = executor.submit(
|
||||
self._execute_unit, target_node, unit_inputs, idx
|
||||
)
|
||||
futures[future] = idx
|
||||
|
||||
results_by_idx: Dict[int, List[Message]] = {}
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
idx = futures[future]
|
||||
try:
|
||||
result = future.result()
|
||||
results_by_idx[idx] = result
|
||||
self.log_manager.debug(
|
||||
f"Dynamic edge -> {target_node.id}#{idx}: "
|
||||
f"completed with {len(result)} outputs"
|
||||
)
|
||||
except Exception as e:
|
||||
self.log_manager.error(
|
||||
f"Dynamic edge -> {target_node.id}#{idx}: "
|
||||
f"failed with error: {e}"
|
||||
)
|
||||
raise
|
||||
|
||||
# Combine results in original order
|
||||
for idx in range(len(execution_units)):
|
||||
if idx in results_by_idx:
|
||||
all_outputs.extend(results_by_idx[idx])
|
||||
|
||||
self.log_manager.info(
|
||||
f"Dynamic edge -> {target_node.id}: "
|
||||
f"Map completed with {len(all_outputs)} total outputs"
|
||||
)
|
||||
|
||||
return all_outputs
|
||||
|
||||
def _execute_tree(
|
||||
self,
|
||||
target_node: Node,
|
||||
execution_units: List[List[Message]],
|
||||
dynamic_config: DynamicEdgeConfig,
|
||||
static_inputs: Optional[List[Message]] = None,
|
||||
) -> List[Message]:
|
||||
"""Execute in Tree mode (fan-out + reduce).
|
||||
|
||||
Args:
|
||||
target_node: Target node template
|
||||
execution_units: Split message units
|
||||
dynamic_config: Dynamic configuration
|
||||
static_inputs: Static inputs (used in first layer only)
|
||||
|
||||
Returns:
|
||||
Single-element list with the final reduced result
|
||||
"""
|
||||
tree_config = dynamic_config.as_tree_config()
|
||||
if tree_config is None:
|
||||
raise ValueError(f"Invalid tree configuration for edge -> {target_node.id}")
|
||||
|
||||
group_size = tree_config.group_size
|
||||
max_parallel = tree_config.max_parallel
|
||||
static_inputs = static_inputs or []
|
||||
|
||||
# Flatten execution units to individual messages
|
||||
current_messages: List[Message] = []
|
||||
for unit in execution_units:
|
||||
current_messages.extend(unit)
|
||||
|
||||
if not current_messages:
|
||||
return []
|
||||
|
||||
self.log_manager.info(
|
||||
f"Dynamic edge -> {target_node.id}: "
|
||||
f"Tree starting with {len(current_messages)} inputs, group_size={group_size}"
|
||||
)
|
||||
|
||||
layer = 0
|
||||
is_first_layer = True
|
||||
|
||||
# Reduction loop
|
||||
while len(current_messages) > 1:
|
||||
layer += 1
|
||||
|
||||
# Group messages
|
||||
groups = group_messages(current_messages, group_size)
|
||||
|
||||
self.log_manager.debug(
|
||||
f"Dynamic edge -> {target_node.id} layer {layer}: "
|
||||
f"processing {len(groups)} groups"
|
||||
)
|
||||
|
||||
layer_outputs: List[Message] = []
|
||||
|
||||
if len(groups) == 1:
|
||||
# Single group - execute directly
|
||||
group_inputs = groups[0]
|
||||
if is_first_layer:
|
||||
group_inputs = list(static_inputs) + group_inputs
|
||||
outputs = self._execute_group(target_node, group_inputs, layer, 0)
|
||||
layer_outputs.extend(outputs)
|
||||
else:
|
||||
# Multiple groups - parallel execution
|
||||
effective_workers = min(len(groups), max_parallel)
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=effective_workers) as executor:
|
||||
futures: Dict[concurrent.futures.Future, int] = {}
|
||||
|
||||
for idx, group in enumerate(groups):
|
||||
group_inputs = group
|
||||
if is_first_layer:
|
||||
group_inputs = list(static_inputs) + group_inputs
|
||||
future = executor.submit(
|
||||
self._execute_group, target_node, group_inputs, layer, idx
|
||||
)
|
||||
futures[future] = idx
|
||||
|
||||
results_by_idx: Dict[int, List[Message]] = {}
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
idx = futures[future]
|
||||
try:
|
||||
result = future.result()
|
||||
results_by_idx[idx] = result
|
||||
except Exception as e:
|
||||
self.log_manager.error(
|
||||
f"Dynamic edge -> {target_node.id}#{layer}-{idx}: "
|
||||
f"failed with error: {e}"
|
||||
)
|
||||
raise
|
||||
|
||||
for idx in range(len(groups)):
|
||||
if idx in results_by_idx:
|
||||
layer_outputs.extend(results_by_idx[idx])
|
||||
|
||||
self.log_manager.debug(
|
||||
f"Dynamic edge -> {target_node.id} layer {layer}: "
|
||||
f"produced {len(layer_outputs)} outputs"
|
||||
)
|
||||
|
||||
current_messages = layer_outputs
|
||||
is_first_layer = False
|
||||
|
||||
# Safety check
|
||||
if layer > 100:
|
||||
self.log_manager.error(
|
||||
f"Dynamic edge -> {target_node.id}: exceeded maximum layers"
|
||||
)
|
||||
break
|
||||
|
||||
self.log_manager.info(
|
||||
f"Dynamic edge -> {target_node.id}: "
|
||||
f"Tree completed after {layer} layers with {len(current_messages)} output(s)"
|
||||
)
|
||||
|
||||
return current_messages
|
||||
|
||||
def _execute_unit(
|
||||
self,
|
||||
node: Node,
|
||||
unit_inputs: List[Message],
|
||||
unit_index: int,
|
||||
) -> List[Message]:
|
||||
"""Execute a single map unit."""
|
||||
self.log_manager.debug(
|
||||
f"Dynamic edge -> {node.id}#{unit_index}: "
|
||||
f"executing with {len(unit_inputs)} inputs"
|
||||
)
|
||||
|
||||
# Tag inputs with unit index
|
||||
# Clone messages first to avoid mutating shared inputs in parallel threads
|
||||
unit_inputs = [msg.clone() for msg in unit_inputs]
|
||||
for msg in unit_inputs:
|
||||
metadata = dict(msg.metadata)
|
||||
metadata["dynamic_edge_unit_index"] = unit_index
|
||||
msg.metadata = metadata
|
||||
|
||||
# Execute using node executor
|
||||
outputs = self.node_executor_func(node, unit_inputs)
|
||||
|
||||
# Tag outputs with unit index
|
||||
for msg in outputs:
|
||||
metadata = dict(msg.metadata)
|
||||
metadata["dynamic_edge_unit_index"] = unit_index
|
||||
msg.metadata = metadata
|
||||
|
||||
return outputs
|
||||
|
||||
def _execute_group(
|
||||
self,
|
||||
node: Node,
|
||||
group_inputs: List[Message],
|
||||
layer: int,
|
||||
group_index: int,
|
||||
) -> List[Message]:
|
||||
"""Execute a single tree group."""
|
||||
instance_id = f"{node.id}#{layer}-{group_index}"
|
||||
|
||||
self.log_manager.debug(
|
||||
f"Dynamic edge -> {instance_id}: executing with {len(group_inputs)} inputs"
|
||||
)
|
||||
|
||||
# Tag inputs
|
||||
# Clone messages first to avoid mutating shared inputs in parallel threads
|
||||
group_inputs = [msg.clone() for msg in group_inputs]
|
||||
for msg in group_inputs:
|
||||
metadata = dict(msg.metadata)
|
||||
metadata["dynamic_edge_tree_layer"] = layer
|
||||
metadata["dynamic_edge_tree_group"] = group_index
|
||||
msg.metadata = metadata
|
||||
|
||||
# Execute
|
||||
outputs = self.node_executor_func(node, group_inputs)
|
||||
|
||||
# Tag outputs
|
||||
for msg in outputs:
|
||||
metadata = dict(msg.metadata)
|
||||
metadata["dynamic_edge_tree_layer"] = layer
|
||||
metadata["dynamic_edge_tree_group"] = group_index
|
||||
metadata["dynamic_edge_instance_id"] = instance_id
|
||||
msg.metadata = metadata
|
||||
msg.role = MessageRole.USER # Mark as user-generated
|
||||
|
||||
return outputs
|
||||
Executable
+141
@@ -0,0 +1,141 @@
|
||||
"""Parallel execution helpers that eliminate duplicated code."""
|
||||
|
||||
import concurrent.futures
|
||||
from typing import Any, Callable, List, Tuple
|
||||
|
||||
from utils.log_manager import LogManager
|
||||
|
||||
|
||||
class ParallelExecutor:
|
||||
"""Manage parallel execution for workflow nodes.
|
||||
|
||||
Provides shared logic for parallel batches and serializes Human nodes when needed.
|
||||
"""
|
||||
|
||||
def __init__(self, log_manager: LogManager, nodes_dict: dict):
|
||||
"""Initialize the parallel executor.
|
||||
|
||||
Args:
|
||||
log_manager: Logger instance
|
||||
nodes_dict: Mapping of ``node_id`` to ``Node``
|
||||
"""
|
||||
self.log_manager = log_manager
|
||||
self.nodes_dict = nodes_dict
|
||||
|
||||
def execute_items_parallel(
|
||||
self,
|
||||
items: List[Any],
|
||||
executor_func: Callable,
|
||||
item_desc_func: Callable[[Any], str],
|
||||
has_blocking_func: Callable[[Any], bool] | None = None,
|
||||
) -> None:
|
||||
"""Execute a list of items in parallel when possible.
|
||||
|
||||
Args:
|
||||
items: Items to execute
|
||||
executor_func: Callable that executes a single item
|
||||
item_desc_func: Callable for logging a human-readable description
|
||||
has_blocking_func: Optional callable to decide if an item requires serialization
|
||||
"""
|
||||
blocking_items, parallel_items = self._partition_blocking_items(items, has_blocking_func)
|
||||
|
||||
if parallel_items:
|
||||
self._execute_parallel_batch(parallel_items, executor_func, item_desc_func)
|
||||
|
||||
if blocking_items:
|
||||
self._execute_sequential_batch(blocking_items, executor_func, item_desc_func)
|
||||
|
||||
def execute_nodes_parallel(
|
||||
self,
|
||||
node_ids: List[str],
|
||||
executor_func: Callable[[str], None]
|
||||
) -> None:
|
||||
"""Execute a list of nodes in parallel.
|
||||
|
||||
Convenience wrapper around ``execute_items_parallel`` specialized for nodes.
|
||||
|
||||
Args:
|
||||
node_ids: List of node identifiers
|
||||
executor_func: Callable that executes a single node
|
||||
"""
|
||||
def item_desc_func(node_id: str) -> str:
|
||||
return f"node {node_id}"
|
||||
|
||||
def has_blocking_func(node_id: str) -> bool:
|
||||
return False
|
||||
|
||||
self.execute_items_parallel(
|
||||
node_ids,
|
||||
executor_func,
|
||||
item_desc_func,
|
||||
has_blocking_func
|
||||
)
|
||||
|
||||
def _partition_blocking_items(
|
||||
self,
|
||||
items: List[Any],
|
||||
has_blocking_func: Callable[[Any], bool] | None
|
||||
) -> Tuple[List[Any], List[Any]]:
|
||||
"""Split items into blocking and parallelizable lists."""
|
||||
blocking_items = []
|
||||
parallel_items = []
|
||||
|
||||
for item in items:
|
||||
if has_blocking_func and has_blocking_func(item):
|
||||
blocking_items.append(item)
|
||||
else:
|
||||
parallel_items.append(item)
|
||||
|
||||
return blocking_items, parallel_items
|
||||
|
||||
def _execute_parallel_batch(
|
||||
self,
|
||||
items: List[Any],
|
||||
executor_func: Callable,
|
||||
item_desc_func: Callable[[Any], str]
|
||||
) -> None:
|
||||
"""Execute a batch of items in parallel.
|
||||
|
||||
Args:
|
||||
items: Items to execute
|
||||
executor_func: Callable per item
|
||||
item_desc_func: Callable returning a readable description
|
||||
"""
|
||||
self.log_manager.debug(f"Executing {len(items)} items in parallel")
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=len(items)) as executor:
|
||||
futures = []
|
||||
for item in items:
|
||||
future = executor.submit(executor_func, item)
|
||||
futures.append((item, future))
|
||||
|
||||
# Wait for every future to finish
|
||||
for item, future in futures:
|
||||
try:
|
||||
future.result()
|
||||
self.log_manager.debug(f"{item_desc_func(item)} completed successfully")
|
||||
except Exception as e:
|
||||
self.log_manager.error(f"{item_desc_func(item)} failed: {str(e)}")
|
||||
raise
|
||||
|
||||
def _execute_sequential_batch(
|
||||
self,
|
||||
items: List[Any],
|
||||
executor_func: Callable,
|
||||
item_desc_func: Callable[[Any], str]
|
||||
) -> None:
|
||||
"""Execute a batch of items sequentially.
|
||||
|
||||
Args:
|
||||
items: Items to execute
|
||||
executor_func: Callable per item
|
||||
item_desc_func: Callable returning a readable description
|
||||
"""
|
||||
for item in items:
|
||||
self.log_manager.debug(f"Executing {item_desc_func(item)} (sequential)")
|
||||
try:
|
||||
executor_func(item)
|
||||
self.log_manager.debug(f"{item_desc_func(item)} completed successfully")
|
||||
except Exception as e:
|
||||
self.log_manager.error(f"{item_desc_func(item)} failed: {str(e)}")
|
||||
raise
|
||||
Executable
+82
@@ -0,0 +1,82 @@
|
||||
"""Resource coordination helpers for workflow node execution."""
|
||||
|
||||
import threading
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Iterable, List, Tuple
|
||||
|
||||
from entity.configs import Node
|
||||
from runtime.node.registry import get_node_registration
|
||||
from utils.log_manager import LogManager
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ResourceRequest:
|
||||
"""Represents a single resource requirement."""
|
||||
|
||||
key: str
|
||||
limit: int
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _ResourceSlot:
|
||||
semaphore: threading.Semaphore
|
||||
limit: int
|
||||
|
||||
|
||||
class ResourceManager:
|
||||
"""Coordinates shared resource usage across nodes."""
|
||||
|
||||
def __init__(self, log_manager: LogManager | None = None):
|
||||
self.log_manager = log_manager
|
||||
self._lock = threading.Lock()
|
||||
self._resources: Dict[str, _ResourceSlot] = {}
|
||||
|
||||
@contextmanager
|
||||
def guard_node(self, node: Node):
|
||||
"""Acquire all resources required by the given node."""
|
||||
requests = self._resolve_node_requests(node)
|
||||
with self._acquire_resources(requests):
|
||||
yield
|
||||
|
||||
def _resolve_node_requests(self, node: Node) -> List[ResourceRequest]:
|
||||
registration = get_node_registration(node.node_type)
|
||||
caps = registration.capabilities
|
||||
requests: List[ResourceRequest] = []
|
||||
key = caps.resource_key
|
||||
limit = caps.resource_limit
|
||||
if key and limit and limit > 0:
|
||||
requests.append(ResourceRequest(key=key, limit=limit))
|
||||
return requests
|
||||
|
||||
@contextmanager
|
||||
def _acquire_resources(self, requests: Iterable[ResourceRequest]):
|
||||
acquired: List[Tuple[str, threading.Semaphore]] = []
|
||||
try:
|
||||
for request in sorted(requests, key=lambda item: item.key):
|
||||
semaphore = self._get_or_create_resource(request)
|
||||
self._log_debug(f"Acquiring resource {request.key}")
|
||||
semaphore.acquire()
|
||||
acquired.append((request.key, semaphore))
|
||||
yield
|
||||
finally:
|
||||
for key, semaphore in reversed(acquired):
|
||||
semaphore.release()
|
||||
self._log_debug(f"Released resource {key}")
|
||||
|
||||
def _get_or_create_resource(self, request: ResourceRequest) -> threading.Semaphore:
|
||||
with self._lock:
|
||||
slot = self._resources.get(request.key)
|
||||
if slot and slot.limit != request.limit:
|
||||
slot = None
|
||||
if not slot:
|
||||
slot = _ResourceSlot(
|
||||
semaphore=threading.Semaphore(request.limit),
|
||||
limit=request.limit,
|
||||
)
|
||||
self._resources[request.key] = slot
|
||||
return slot.semaphore
|
||||
|
||||
def _log_debug(self, message: str) -> None:
|
||||
if self.log_manager:
|
||||
self.log_manager.debug(message)
|
||||
Reference in New Issue
Block a user