Coverage for smartmdao / solvers.py: 100%

162 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-07-05 17:29 +0200

1import logging 

2from collections import defaultdict, deque 

3from dataclasses import dataclass 

4from typing import List, Dict, Any, Set, Protocol, Optional 

5 

6from .models import Step 

7from .executor import StepExecutor 

8from .graph import map_producers as _map_producers, build_dependency_graph as _build_dependency_graph 

9from .validation import TypeChecker 

10 

11# Initialize module-level logger 

12logger = logging.getLogger(__name__) 

13 

14class Solver(Protocol): 

15 """Interface for execution logic.""" 

16 def solve(self, steps: List[Step], inputs: Dict[str, Any], type_checker: Optional[TypeChecker] = None) -> Dict[str, Any]: 

17 ... 

18 

19class DAGSolver: 

20 """ 

21 Standard Topological Sort Solver.  

22 Ideal for linear workflows. 

23 """ 

24 def solve(self, steps: List[Step], inputs: Dict[str, Any], type_checker: Optional[TypeChecker] = None) -> Dict[str, Any]: 

25 logger.info("DAGSolver started.") 

26 execution_order = self._topological_sort(steps, set(inputs.keys())) 

27 logger.debug(f"Topological sort order: {[s.name for s in execution_order]}") 

28 

29 memory = inputs.copy() 

30 

31 for step in execution_order: 

32 StepExecutor.run_step(step, memory, type_checker=type_checker) 

33 

34 return memory 

35 

36 def _topological_sort(self, steps: List[Step], input_keys: Set[str]) -> List[Step]: 

37 producers_map = _map_producers(steps) 

38 adj_list, indegree = _build_dependency_graph(steps, input_keys, producers_map) 

39 

40 # Kahn's Algorithm 

41 queue = deque([s for s, deg in indegree.items() if deg == 0]) 

42 sorted_steps = [] 

43 

44 while queue: 

45 current = queue.popleft() 

46 sorted_steps.append(current) 

47 

48 for neighbor in adj_list[current]: 

49 indegree[neighbor] -= 1 

50 if indegree[neighbor] == 0: 

51 queue.append(neighbor) 

52 

53 if len(sorted_steps) != len(steps): 

54 logger.error("Cycle detected in DAGSolver.") 

55 raise ValueError("Cycle detected in pipeline. Use HybridSolver or IterativeSolver.") 

56 

57 return sorted_steps 

58 

59@dataclass 

60class IterativeSolver: 

61 """ 

62 Solves systems with feedback loops. 

63 """ 

64 max_iterations: int = 100 

65 tolerance: float = 1e-6 

66 target_var: Optional[str] = None 

67 execution_order: Optional[List[str]] = None 

68 

69 def solve(self, steps: List[Step], inputs: Dict[str, Any], type_checker: Optional[TypeChecker] = None) -> Dict[str, Any]: 

70 memory = inputs.copy() 

71 residuals = [] 

72 

73 run_sequence = self._determine_execution_order(steps) 

74 logger.info(f"IterativeSolver started. Sequence: {[s.name for s in run_sequence]}") 

75 

76 # Identify variables produced by these steps (for auto-convergence) 

77 produced_vars = set() 

78 for s in steps: 

79 produced_vars.update(s.resolve_output_names()) 

80 

81 for i in range(self.max_iterations): 

82 # Snapshot state for convergence check 

83 prev_state = {k: memory.get(k) for k in produced_vars if k in memory} 

84 

85 # Execute 

86 for step in run_sequence: 

87 StepExecutor.run_step(step, memory, type_checker=type_checker) 

88 

89 # Check Convergence 

90 diff = self._calculate_residual(prev_state, memory, produced_vars) 

91 residuals.append(diff) 

92 

93 # Only break if we actually calculated a numeric difference (not inf) 

94 if diff != float('inf') and diff < self.tolerance: 

95 logger.info(f"Converged at iteration {i+1} with residual {diff:.6e}") 

96 break 

97 

98 logger.debug(f"Iteration {i+1}: residual {diff:.6e}") 

99 else: 

100 logger.warning(f"Reached max_iterations ({self.max_iterations}) without converging. Last residual: {residuals[-1]:.6e}") 

101 

102 # Store residuals (append to potentially existing history from other cycles) 

103 memory.setdefault('residual_history', []).append(residuals) 

104 return memory 

105 

106 def _calculate_residual(self, prev_state: Dict, current_memory: Dict, produced_vars: Set[str]) -> float: 

107 """ 

108 Calculates the maximum change in variables.  

109 """ 

110 if self.target_var: 

111 p = prev_state.get(self.target_var) 

112 c = current_memory.get(self.target_var) 

113 return abs(c - p) if (isinstance(p, (int, float)) and isinstance(c, (int, float))) else float('inf') 

114 

115 max_diff = 0.0 

116 numeric_vars_found = False 

117 

118 for k in produced_vars: 

119 p = prev_state.get(k) 

120 c = current_memory.get(k) 

121 

122 # Strictly require both to be numeric 

123 if isinstance(p, (int, float)) and isinstance(c, (int, float)): 

124 diff = abs(c - p) 

125 max_diff = max(max_diff, diff) 

126 numeric_vars_found = True 

127 

128 if numeric_vars_found: 

129 return max_diff 

130 

131 # If no numeric variables updated, we can't judge convergence numerically. 

132 return float('inf') 

133 

134 def _determine_execution_order(self, steps: List[Step]) -> List[Step]: 

135 if not self.execution_order: 

136 return steps 

137 

138 step_map = {s.name: s for s in steps} 

139 return [step_map[name] for name in self.execution_order if name in step_map] 

140 

141 

142class HybridSolver: 

143 """ 

144 Advanced solver that automatically decomposes the pipeline into  

145 Linear (DAG) and Iterative (Cyclic) components (Strongly Connected Components). 

146 """ 

147 def __init__(self, max_iterations: int = 100, tolerance: float = 1e-6): 

148 self.max_iterations = max_iterations 

149 self.tolerance = tolerance 

150 

151 def solve(self, steps: List[Step], inputs: Dict[str, Any], type_checker: Optional[TypeChecker] = None) -> Dict[str, Any]: 

152 logger.info("HybridSolver started.") 

153 input_keys = set(inputs.keys()) 

154 producers_map = _map_producers(steps) 

155 

156 # 1. Build Adjacency Graph (Producer -> Consumer) 

157 adj_list, _ = _build_dependency_graph(steps, input_keys, producers_map) 

158 

159 # 2. Find Strongly Connected Components (SCCs) 

160 sccs = self._tarjan_scc(steps, adj_list) 

161 logger.debug(f"Detected {len(sccs)} execution blocks (SCCs).") 

162 

163 # 3. Build Condensation Graph (DAG of SCCs) 

164 scc_map = {step: i for i, cluster in enumerate(sccs) for step in cluster} 

165 scc_adj = defaultdict(set) 

166 scc_indegree = defaultdict(int) 

167 

168 for u in steps: 

169 u_scc = scc_map[u] 

170 for v in adj_list[u]: 

171 v_scc = scc_map[v] 

172 if u_scc != v_scc: 

173 if v_scc not in scc_adj[u_scc]: 

174 scc_adj[u_scc].add(v_scc) 

175 scc_indegree[v_scc] += 1 

176 

177 # Ensure all SCCs have an entry 

178 for i in range(len(sccs)): 

179 if i not in scc_indegree: 

180 scc_indegree[i] = 0 

181 

182 # 4. Topological Sort of SCCs 

183 queue = deque([i for i, deg in scc_indegree.items() if deg == 0]) 

184 execution_plan = [] 

185 

186 while queue: 

187 current_scc_idx = queue.popleft() 

188 execution_plan.append(sccs[current_scc_idx]) 

189 

190 for neighbor_scc in scc_adj[current_scc_idx]: 

191 scc_indegree[neighbor_scc] -= 1 

192 if scc_indegree[neighbor_scc] == 0: 

193 queue.append(neighbor_scc) 

194 

195 # 5. Execute 

196 memory = inputs.copy() 

197 

198 for group in execution_plan: 

199 # Case A: Linear 

200 if len(group) == 1 and group[0] not in adj_list[group[0]]: 

201 step = group[0] 

202 StepExecutor.run_step(step, memory, type_checker=type_checker) 

203 continue 

204 

205 # Case B: Cyclic 

206 # Sort alphabetically to ensure deterministic execution order within the cycle 

207 group_sorted = sorted(group, key=lambda s: s.name) 

208 

209 logger.info(f"Cyclic Block Detected: {[s.name for s in group_sorted]}") 

210 sub_solver = IterativeSolver( 

211 max_iterations=self.max_iterations, 

212 tolerance=self.tolerance 

213 ) 

214 

215 cycle_results = sub_solver.solve(group_sorted, memory, type_checker=type_checker) 

216 memory.update(cycle_results) 

217 

218 return memory 

219 

220 def _tarjan_scc(self, steps: List[Step], adj_list: Dict[Step, List[Step]]) -> List[List[Step]]: 

221 index = 0 

222 indices = {} 

223 lowlinks = {} 

224 stack = [] 

225 on_stack = set() 

226 sccs = [] 

227 

228 def strongconnect(v): 

229 nonlocal index 

230 indices[v] = index 

231 lowlinks[v] = index 

232 index += 1 

233 stack.append(v) 

234 on_stack.add(v) 

235 

236 for w in adj_list[v]: 

237 if w not in indices: 

238 strongconnect(w) 

239 lowlinks[v] = min(lowlinks[v], lowlinks[w]) 

240 elif w in on_stack: 

241 lowlinks[v] = min(lowlinks[v], indices[w]) 

242 

243 if lowlinks[v] == indices[v]: 

244 new_scc = [] 

245 while True: 

246 w = stack.pop() 

247 on_stack.remove(w) 

248 new_scc.append(w) 

249 if w == v: 

250 break 

251 sccs.append(new_scc) 

252 

253 for step in steps: 

254 if step not in indices: 

255 strongconnect(step) 

256 

257 return sccs