Coverage for src/dataknobs_fsm/execution/history.py: 16%
286 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-20 16:46 -0600
1"""Execution history tracking for FSM state machines."""
3import time
4from enum import Enum
5from typing import Any, Dict, List
7from dataknobs_structures import Tree
9from dataknobs_fsm.core.data_modes import DataHandlingMode
12class ExecutionStatus(Enum):
13 """Status of an execution step."""
14 PENDING = "pending"
15 IN_PROGRESS = "in_progress"
16 COMPLETED = "completed"
17 FAILED = "failed"
18 SKIPPED = "skipped"
21class ExecutionStep:
22 """Represents a single step in the execution history."""
24 def __init__(
25 self,
26 step_id: str,
27 state_name: str,
28 network_name: str,
29 timestamp: float,
30 data_mode: DataHandlingMode = DataHandlingMode.COPY,
31 status: ExecutionStatus = ExecutionStatus.PENDING
32 ):
33 """Initialize execution step.
35 Args:
36 step_id: Unique identifier for this step.
37 state_name: Name of the state.
38 network_name: Name of the network.
39 timestamp: Unix timestamp when step was created.
40 data_mode: Data mode used for this step.
41 status: Current status of the step.
42 """
43 self.step_id = step_id
44 self.state_name = state_name
45 self.network_name = network_name
46 self.timestamp = timestamp
47 self.start_time: float | None = None
48 self.end_time: float | None = None
49 self.data_mode = data_mode
50 self.status = status
52 # Execution details
53 self.arc_taken: str | None = None
54 self.data_snapshot: Any | None = None
55 self.error: Exception | None = None
56 self.metrics: Dict[str, Any] = {}
57 self.resource_usage: Dict[str, Any] = {}
59 # Stream tracking
60 self.stream_progress: Dict[str, Any] | None = None
61 self.chunks_processed: int = 0
62 self.records_processed: int = 0
64 def start(self) -> None:
65 """Mark step as started."""
66 self.start_time = time.time()
67 self.status = ExecutionStatus.IN_PROGRESS
69 def complete(self, arc_taken: str | None = None) -> None:
70 """Mark step as completed.
72 Args:
73 arc_taken: The arc that was taken from this state.
74 """
75 self.end_time = time.time()
76 self.status = ExecutionStatus.COMPLETED
77 self.arc_taken = arc_taken
79 def fail(self, error: Exception) -> None:
80 """Mark step as failed.
82 Args:
83 error: The exception that caused the failure.
84 """
85 self.end_time = time.time()
86 self.status = ExecutionStatus.FAILED
87 self.error = error
89 def skip(self, reason: str) -> None:
90 """Mark step as skipped.
92 Args:
93 reason: Reason for skipping.
94 """
95 self.end_time = time.time()
96 self.status = ExecutionStatus.SKIPPED
97 self.metrics['skip_reason'] = reason
99 def add_metric(self, key: str, value: Any) -> None:
100 """Add a metric to this step.
102 Args:
103 key: Metric key.
104 value: Metric value.
105 """
106 self.metrics[key] = value
108 def add_resource_usage(self, resource_type: str, usage: Dict[str, Any]) -> None:
109 """Track resource usage for this step.
111 Args:
112 resource_type: Type of resource (e.g., "database", "llm").
113 usage: Usage metrics.
114 """
115 self.resource_usage[resource_type] = usage
117 def update_stream_progress(
118 self,
119 chunks: int,
120 records: int,
121 current_position: int | None = None
122 ) -> None:
123 """Update streaming progress.
125 Args:
126 chunks: Number of chunks processed.
127 records: Number of records processed.
128 current_position: Current position in stream.
129 """
130 self.chunks_processed = chunks
131 self.records_processed = records
132 if current_position is not None:
133 if self.stream_progress is None:
134 self.stream_progress = {}
135 self.stream_progress['position'] = current_position
137 @property
138 def duration(self) -> float | None:
139 """Get execution duration in seconds."""
140 if self.start_time and self.end_time:
141 return self.end_time - self.start_time
142 return None
144 def to_dict(self) -> Dict[str, Any]:
145 """Convert to dictionary representation."""
146 return {
147 'step_id': self.step_id,
148 'state_name': self.state_name,
149 'network_name': self.network_name,
150 'timestamp': self.timestamp,
151 'start_time': self.start_time,
152 'end_time': self.end_time,
153 'duration': self.duration,
154 'data_mode': self.data_mode.value,
155 'status': self.status.value,
156 'arc_taken': self.arc_taken,
157 'error': str(self.error) if self.error else None,
158 'metrics': self.metrics,
159 'resource_usage': self.resource_usage,
160 'stream_progress': self.stream_progress,
161 'chunks_processed': self.chunks_processed,
162 'records_processed': self.records_processed
163 }
165 @classmethod
166 def from_dict(cls, data: Dict[str, Any]) -> 'ExecutionStep':
167 """Create ExecutionStep from dictionary representation.
169 Args:
170 data: Dictionary with step data.
172 Returns:
173 ExecutionStep instance.
174 """
175 step = cls(
176 step_id=data['step_id'],
177 state_name=data['state_name'],
178 network_name=data['network_name'],
179 timestamp=data['timestamp'],
180 data_mode=DataHandlingMode(data['data_mode']),
181 status=ExecutionStatus(data['status'])
182 )
184 # Restore timing info
185 step.start_time = data.get('start_time')
186 step.end_time = data.get('end_time')
188 # Restore execution details
189 step.arc_taken = data.get('arc_taken')
191 # Restore error (as string - can't fully reconstruct Exception)
192 error_str = data.get('error')
193 if error_str:
194 step.error = Exception(error_str)
196 # Restore metrics and usage
197 step.metrics = data.get('metrics', {})
198 step.resource_usage = data.get('resource_usage', {})
200 # Restore stream progress
201 step.stream_progress = data.get('stream_progress')
202 step.chunks_processed = data.get('chunks_processed', 0)
203 step.records_processed = data.get('records_processed', 0)
205 return step
208class ExecutionHistory:
209 """Tracks execution history using a tree structure.
211 This class manages the execution history of an FSM, tracking:
212 - State transitions
213 - Data modifications based on mode
214 - Resource usage
215 - Stream progress
216 - Execution metrics
217 """
219 def __init__(
220 self,
221 fsm_name: str,
222 execution_id: str,
223 data_mode: DataHandlingMode = DataHandlingMode.COPY,
224 max_depth: int | None = None,
225 enable_data_snapshots: bool = False
226 ):
227 """Initialize execution history.
229 Args:
230 fsm_name: Name of the FSM.
231 execution_id: Unique execution identifier.
232 data_mode: Default data mode for execution.
233 max_depth: Maximum tree depth (for pruning).
234 enable_data_snapshots: Whether to store data snapshots.
235 """
236 self.fsm_name = fsm_name
237 self.execution_id = execution_id
238 self.data_mode = data_mode
239 self.max_depth = max_depth
240 self.enable_data_snapshots = enable_data_snapshots
242 # Create tree structure - Tree stores root nodes
243 self.tree_roots: List[Tree] = []
244 self.current_node: Tree | None = None
246 # Tracking
247 self.start_time = time.time()
248 self.end_time: float | None = None
249 self.total_steps = 0
250 self.failed_steps = 0
251 self.skipped_steps = 0
253 # Mode-specific storage
254 self._mode_storage: Dict[DataHandlingMode, List[ExecutionStep]] = {
255 DataHandlingMode.COPY: [],
256 DataHandlingMode.REFERENCE: [],
257 DataHandlingMode.DIRECT: []
258 }
260 # Resource tracking
261 self.resource_summary: Dict[str, Dict[str, Any]] = {}
263 # Stream tracking
264 self.stream_summary: Dict[str, Any] = {
265 'total_chunks': 0,
266 'total_records': 0,
267 'streams_processed': 0
268 }
270 def add_step(
271 self,
272 state_name: str,
273 network_name: str,
274 data: Any | None = None,
275 parent_step_id: str | None = None
276 ) -> ExecutionStep:
277 """Add a new execution step.
279 Args:
280 state_name: Name of the state.
281 network_name: Name of the network.
282 data: Optional data snapshot.
283 parent_step_id: Parent step ID for branching.
285 Returns:
286 The created ExecutionStep.
287 """
288 import uuid
290 step_id = str(uuid.uuid4())
291 step = ExecutionStep(
292 step_id=step_id,
293 state_name=state_name,
294 network_name=network_name,
295 timestamp=time.time(),
296 data_mode=self.data_mode
297 )
299 # Store data snapshot if enabled
300 if self.enable_data_snapshots and data is not None:
301 step.data_snapshot = self._snapshot_data(data)
303 # Add to tree
304 if parent_step_id:
305 parent_node = self._find_node_by_step_id(parent_step_id)
306 if parent_node:
307 self.current_node = Tree(step, parent=parent_node)
308 elif self.current_node:
309 self.current_node = Tree(step, parent=self.current_node)
310 else:
311 # Create root node
312 self.current_node = Tree(step)
313 self.tree_roots.append(self.current_node)
315 # Track in mode-specific storage
316 self._mode_storage[self.data_mode].append(step)
318 self.total_steps += 1
320 # Prune if needed
321 if self.max_depth and self._get_max_depth() > self.max_depth:
322 self._prune_old_branches()
324 return step
326 def update_step(
327 self,
328 step_id: str,
329 status: ExecutionStatus | None = None,
330 arc_taken: str | None = None,
331 error: Exception | None = None,
332 metrics: Dict[str, Any] | None = None
333 ) -> bool:
334 """Update an existing step.
336 Args:
337 step_id: Step ID to update.
338 status: New status.
339 arc_taken: Arc taken from this step.
340 error: Error if failed.
341 metrics: Additional metrics.
343 Returns:
344 True if step was found and updated.
345 """
346 node = self._find_node_by_step_id(step_id)
347 if not node:
348 return False
350 step = node.data
352 if status == ExecutionStatus.COMPLETED:
353 step.complete(arc_taken)
354 elif status == ExecutionStatus.FAILED:
355 if error:
356 step.fail(error)
357 self.failed_steps += 1
358 elif status == ExecutionStatus.SKIPPED:
359 step.skip(metrics.get('reason', 'Unknown') if metrics else 'Unknown')
360 self.skipped_steps += 1
361 elif status:
362 step.status = status
364 if metrics:
365 for key, value in metrics.items():
366 step.add_metric(key, value)
368 return True
370 def get_path_to_current(self) -> List[ExecutionStep]:
371 """Get the path from root to current step.
373 Returns:
374 List of steps from root to current.
375 """
376 if not self.current_node:
377 return []
379 path = []
380 node = self.current_node
381 while node:
382 path.insert(0, node.data)
383 node = node.parent
385 return path
387 def get_all_paths(self) -> List[List[ExecutionStep]]:
388 """Get all execution paths in the tree.
390 Returns:
391 List of all paths from root to leaves.
392 """
393 paths = []
395 def collect_paths(node: Tree, current_path: List[ExecutionStep]):
396 current_path.append(node.data)
398 if not node.children:
399 # Leaf node - save path
400 paths.append(current_path.copy())
401 else:
402 # Continue down each branch
403 for child in node.children:
404 collect_paths(child, current_path.copy())
406 for root in self.tree_roots:
407 collect_paths(root, [])
409 return paths
411 @property
412 def steps(self) -> List[ExecutionStep]:
413 """Get all execution steps from the history tree.
415 Returns:
416 List of all execution steps in order.
417 """
418 all_steps = []
420 def collect_steps(node: Tree):
421 all_steps.append(node.data)
422 if node.children:
423 for child in node.children:
424 collect_steps(child)
426 for root in self.tree_roots:
427 collect_steps(root)
429 return all_steps
431 def get_steps_by_state(self, state_name: str) -> List[ExecutionStep]:
432 """Get all steps for a specific state.
434 Args:
435 state_name: Name of the state.
437 Returns:
438 List of steps for that state.
439 """
440 steps = []
442 def collect_steps(node: Tree):
443 if node.data.state_name == state_name:
444 steps.append(node.data)
445 if node.children:
446 for child in node.children:
447 collect_steps(child)
449 for root in self.tree_roots:
450 collect_steps(root)
452 return steps
454 def get_steps_by_mode(self, mode: DataHandlingMode) -> List[ExecutionStep]:
455 """Get all steps executed in a specific data mode.
457 Args:
458 mode: Data mode to filter by.
460 Returns:
461 List of steps in that mode.
462 """
463 return self._mode_storage.get(mode, []).copy()
465 def get_resource_usage(self) -> Dict[str, Any]:
466 """Get aggregated resource usage.
468 Returns:
469 Resource usage summary.
470 """
471 usage = {}
473 def aggregate_usage(node: Tree):
474 step = node.data
475 for resource_type, metrics in step.resource_usage.items():
476 if resource_type not in usage:
477 usage[resource_type] = {
478 'total_calls': 0,
479 'total_duration': 0,
480 'steps': []
481 }
483 usage[resource_type]['total_calls'] += 1 # type: ignore
484 if 'duration' in metrics:
485 usage[resource_type]['total_duration'] += metrics['duration']
486 usage[resource_type]['steps'].append(step.step_id)
488 if node.children:
489 for child in node.children:
490 aggregate_usage(child)
492 for root in self.tree_roots:
493 aggregate_usage(root)
495 return usage
497 def get_stream_progress(self) -> Dict[str, Any]:
498 """Get streaming progress summary.
500 Returns:
501 Stream progress information.
502 """
503 total_chunks = 0
504 total_records = 0
506 def aggregate_stream(node: Tree):
507 nonlocal total_chunks, total_records
508 step = node.data
509 total_chunks += step.chunks_processed
510 total_records += step.records_processed
512 if node.children:
513 for child in node.children:
514 aggregate_stream(child)
516 for root in self.tree_roots:
517 aggregate_stream(root)
519 return {
520 'total_chunks': total_chunks,
521 'total_records': total_records,
522 'streams_processed': self.stream_summary['streams_processed']
523 }
525 def finalize(self) -> None:
526 """Mark execution as complete."""
527 self.end_time = time.time()
529 def get_summary(self) -> Dict[str, Any]:
530 """Get execution summary.
532 Returns:
533 Summary of the execution.
534 """
535 return {
536 'fsm_name': self.fsm_name,
537 'execution_id': self.execution_id,
538 'start_time': self.start_time,
539 'end_time': self.end_time,
540 'duration': self.end_time - self.start_time if self.end_time else None,
541 'total_steps': self.total_steps,
542 'failed_steps': self.failed_steps,
543 'skipped_steps': self.skipped_steps,
544 'completed_steps': self.total_steps - self.failed_steps - self.skipped_steps,
545 'data_mode': self.data_mode.value,
546 'tree_depth': self._get_max_depth(),
547 'total_paths': len(self.get_all_paths()),
548 'resource_usage': self.get_resource_usage(),
549 'stream_progress': self.get_stream_progress()
550 }
552 def to_dict(self) -> Dict[str, Any]:
553 """Convert to dictionary representation.
555 Returns:
556 Dictionary representation of history.
557 """
558 paths = []
559 for path in self.get_all_paths():
560 paths.append([step.to_dict() for step in path])
562 return {
563 'summary': self.get_summary(),
564 'paths': paths,
565 'mode_storage': {
566 mode.value: [s.to_dict() for s in steps]
567 for mode, steps in self._mode_storage.items()
568 }
569 }
571 @classmethod
572 def from_dict(cls, data: Dict[str, Any]) -> 'ExecutionHistory':
573 """Create ExecutionHistory from dictionary representation.
575 Args:
576 data: Dictionary with history data.
578 Returns:
579 ExecutionHistory instance.
580 """
581 summary = data['summary']
583 # Create history instance
584 history = cls(
585 fsm_name=summary['fsm_name'],
586 execution_id=summary['execution_id'],
587 data_mode=DataHandlingMode(summary['data_mode']),
588 max_depth=None, # Will be inferred
589 enable_data_snapshots=False # Will be inferred from data
590 )
592 # Restore properties from summary
593 history.start_time = summary['start_time']
594 history.end_time = summary.get('end_time')
595 history.total_steps = summary['total_steps']
596 history.failed_steps = summary['failed_steps']
597 history.skipped_steps = summary['skipped_steps']
599 # Rebuild tree structure from paths
600 for path_data in data.get('paths', []):
601 parent_node = None
602 for step_dict in path_data:
603 step = ExecutionStep.from_dict(step_dict)
605 # Create tree node
606 if parent_node is None:
607 # This is a root step
608 node = Tree(step)
609 history.tree_roots.append(node)
610 else:
611 # Child step
612 node = Tree(step, parent=parent_node)
614 parent_node = node
616 # Track in mode-specific storage
617 history._mode_storage[step.data_mode].append(step)
619 # Set current node to the last node in this path
620 if parent_node:
621 history.current_node = parent_node
623 return history
625 def _find_node_by_step_id(self, step_id: str) -> Tree | None:
626 """Find a node by step ID.
628 Args:
629 step_id: Step ID to find.
631 Returns:
632 Tree node if found, None otherwise.
633 """
634 def search_node(node: Tree) -> Tree | None:
635 if node.data.step_id == step_id:
636 return node
637 if node.children:
638 for child in node.children:
639 result = search_node(child)
640 if result:
641 return result
642 return None
644 for root in self.tree_roots:
645 result = search_node(root)
646 if result:
647 return result
649 return None
651 def _snapshot_data(self, data: Any) -> Any:
652 """Create a snapshot of data based on mode.
654 Args:
655 data: Data to snapshot.
657 Returns:
658 Snapshot of the data.
659 """
660 import copy
662 if self.data_mode == DataHandlingMode.COPY:
663 # Deep copy for COPY mode
664 return copy.deepcopy(data)
665 elif self.data_mode == DataHandlingMode.REFERENCE:
666 # Store reference info only
667 return {
668 'type': type(data).__name__,
669 'id': id(data),
670 'size': len(data) if hasattr(data, '__len__') else None
671 }
672 else: # DIRECT mode
673 # Store minimal info
674 return {'type': type(data).__name__}
676 def _get_max_depth(self) -> int:
677 """Get the maximum depth of the tree.
679 Returns:
680 Maximum depth across all trees.
681 """
682 max_depth = 0
684 def get_depth(node: Tree, depth: int = 0):
685 nonlocal max_depth
686 depth += 1
687 max_depth = max(max_depth, depth)
688 if node.children:
689 for child in node.children:
690 get_depth(child, depth)
692 for root in self.tree_roots:
693 get_depth(root)
695 return max_depth
697 def _prune_old_branches(self) -> None:
698 """Prune old branches based on mode-specific strategy."""
699 if self.data_mode == DataHandlingMode.COPY:
700 # Keep all branches for COPY mode (full history)
701 pass
702 elif self.data_mode == DataHandlingMode.REFERENCE:
703 # Prune branches not on critical path
704 # Keep only paths that lead to current node
705 pass
706 else: # DIRECT mode
707 # Aggressive pruning - keep only recent history
708 # Remove all but the current path
709 current_path = self.get_path_to_current()
710 if current_path and len(current_path) > self.max_depth: # type: ignore
711 # Keep only last max_depth steps
712 steps_to_keep = current_path[-self.max_depth:] # type: ignore
713 # Rebuild tree with only these steps
714 self.tree_roots = []
715 self.current_node = None
716 for step in steps_to_keep:
717 if not self.current_node:
718 self.current_node = Tree(step)
719 self.tree_roots.append(self.current_node)
720 else:
721 self.current_node = Tree(step, parent=self.current_node)