Coverage for src/dataknobs_fsm/execution/context.py: 53%
187 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-20 16:51 -0600
1"""Execution context for FSM state machines."""
3import time
4from dataclasses import dataclass, field
5from enum import Enum
6from typing import Any, Dict, List, Tuple, Union
8from dataknobs_data.database import AsyncDatabase, SyncDatabase
10from dataknobs_fsm.core.modes import ProcessingMode, TransactionMode
11from dataknobs_fsm.streaming.core import StreamChunk, StreamContext
14class ResourceStatus(Enum):
15 """Status of a resource in the execution context."""
16 AVAILABLE = "available"
17 ALLOCATED = "allocated"
18 BUSY = "busy"
19 ERROR = "error"
22@dataclass
23class ResourceAllocation:
24 """Resource allocation information."""
25 resource_type: str
26 resource_id: str
27 status: ResourceStatus
28 allocated_at: float
29 metadata: Dict[str, Any] = field(default_factory=dict)
32@dataclass
33class TransactionInfo:
34 """Transaction tracking information."""
35 transaction_id: str
36 mode: TransactionMode
37 started_at: float
38 operations: List[Dict[str, Any]] = field(default_factory=list)
39 is_committed: bool = False
40 is_rolled_back: bool = False
43class ExecutionContext:
44 """Execution context for FSM processing with full mode support.
46 This context manages:
47 - Data mode configuration (single, batch, stream)
48 - Transaction management
49 - Resource allocation and tracking
50 - Stream coordination
51 - State tracking and network stack
52 - Parallel execution paths
53 """
55 def __init__(
56 self,
57 data_mode: ProcessingMode = ProcessingMode.SINGLE,
58 transaction_mode: TransactionMode = TransactionMode.NONE,
59 resources: Dict[str, Any] | None = None,
60 database: Union[SyncDatabase, AsyncDatabase] | None = None,
61 stream_context: StreamContext | None = None
62 ):
63 """Initialize execution context.
65 Args:
66 data_mode: Data processing mode.
67 transaction_mode: Transaction handling mode.
68 resources: Initial resource configurations.
69 database: Database connection for transactions.
70 stream_context: Stream context for stream mode.
71 """
72 # Mode configuration
73 self.data_mode = data_mode
74 self.transaction_mode = transaction_mode
76 # State tracking
77 self.current_state: str | None = None
78 self.previous_state: str | None = None
79 self.network_stack: List[Tuple[str, str | None]] = []
80 self.state_history: List[str] = []
82 # Data management
83 self.data: Any = None
84 self.metadata: Dict[str, Any] = {}
85 self.variables: Dict[str, Any] = {}
87 # Resource management
88 self.resources: Dict[str, ResourceAllocation] = {}
89 self.resource_limits: Dict[str, Any] = resources or {}
90 self.resource_manager: Any = None # ResourceManager instance
91 self.current_state_resources: Dict[str, Any] = {} # Resources allocated to current state
92 self.parent_state_resources: Dict[str, Any] = {} # Resources from parent state (in subnetworks)
94 # Transaction management
95 self.database = database
96 self.current_transaction: TransactionInfo | None = None
97 self.transaction_history: List[TransactionInfo] = []
99 # Stream coordination
100 self.stream_context = stream_context
101 self.current_chunk: StreamChunk | None = None
102 self.processed_chunks: int = 0
104 # Batch processing
105 self.batch_data: List[Any] = []
106 self.batch_results: List[Any] = []
107 self.batch_errors: List[Tuple[int, Exception]] = []
109 # Parallel execution
110 self.parallel_paths: Dict[str, ExecutionContext] = {}
111 self.is_child_context: bool = False
112 self.parent_context: ExecutionContext | None = None
114 # Performance tracking
115 self.start_time: float = time.time()
116 self.state_timings: Dict[str, float] = {}
117 self.function_call_count: Dict[str, int] = {}
119 # State instance tracking for debugging
120 self.current_state_instance: Any = None
122 def push_network(self, network_name: str, return_state: str | None = None) -> None:
123 """Push a network onto the execution stack.
125 Args:
126 network_name: Name of network to push.
127 return_state: State to return to after network completes.
128 """
129 self.network_stack.append((network_name, return_state))
131 def pop_network(self) -> Tuple[str, str | None]:
132 """Pop a network from the execution stack.
134 Returns:
135 Tuple of (network_name, return_state).
136 """
137 if self.network_stack:
138 return self.network_stack.pop()
139 return ("", None)
141 def set_state(self, state_name: str) -> None:
142 """Set the current state.
144 Args:
145 state_name: Name of the new state.
146 """
147 if self.current_state:
148 self.previous_state = self.current_state
149 self.state_history.append(self.current_state)
150 self.current_state = state_name
151 self.state_timings[state_name] = time.time()
153 def allocate_resource(
154 self,
155 resource_type: str,
156 resource_id: str,
157 metadata: Dict[str, Any] | None = None
158 ) -> bool:
159 """Allocate a resource.
161 Args:
162 resource_type: Type of resource.
163 resource_id: Unique resource identifier.
164 metadata: Optional resource metadata.
166 Returns:
167 True if allocation successful.
168 """
169 key = f"{resource_type}:{resource_id}"
171 if key in self.resources:
172 if self.resources[key].status != ResourceStatus.AVAILABLE:
173 return False
175 self.resources[key] = ResourceAllocation(
176 resource_type=resource_type,
177 resource_id=resource_id,
178 status=ResourceStatus.ALLOCATED,
179 allocated_at=time.time(),
180 metadata=metadata or {}
181 )
182 return True
184 def release_resource(self, resource_type: str, resource_id: str) -> bool:
185 """Release an allocated resource.
187 Args:
188 resource_type: Type of resource.
189 resource_id: Resource identifier.
191 Returns:
192 True if release successful.
193 """
194 key = f"{resource_type}:{resource_id}"
196 if key in self.resources:
197 self.resources[key].status = ResourceStatus.AVAILABLE
198 return True
199 return False
201 def start_transaction(self, transaction_id: str | None = None) -> bool:
202 """Start a new transaction.
204 Args:
205 transaction_id: Optional transaction ID.
207 Returns:
208 True if transaction started.
209 """
210 if self.transaction_mode == TransactionMode.NONE:
211 return False
213 if self.current_transaction and not self.current_transaction.is_committed:
214 return False
216 self.current_transaction = TransactionInfo(
217 transaction_id=transaction_id or str(time.time()),
218 mode=self.transaction_mode,
219 started_at=time.time()
220 )
222 # Start database transaction if available
223 if self.database and hasattr(self.database, 'begin_transaction'):
224 self.database.begin_transaction()
226 return True
228 def commit_transaction(self) -> bool:
229 """Commit the current transaction.
231 Returns:
232 True if commit successful.
233 """
234 if not self.current_transaction:
235 return False
237 # Commit database transaction
238 if self.database and hasattr(self.database, 'commit'):
239 self.database.commit()
241 self.current_transaction.is_committed = True
242 self.transaction_history.append(self.current_transaction)
243 self.current_transaction = None
245 return True
247 def rollback_transaction(self) -> bool:
248 """Rollback the current transaction.
250 Returns:
251 True if rollback successful.
252 """
253 if not self.current_transaction:
254 return False
256 # Rollback database transaction
257 if self.database and hasattr(self.database, 'rollback'):
258 self.database.rollback()
260 self.current_transaction.is_rolled_back = True
261 self.transaction_history.append(self.current_transaction)
262 self.current_transaction = None
264 return True
266 def log_operation(self, operation: str, details: Dict[str, Any]) -> None:
267 """Log an operation in the current transaction.
269 Args:
270 operation: Operation name.
271 details: Operation details.
272 """
273 if self.current_transaction:
274 self.current_transaction.operations.append({
275 'operation': operation,
276 'details': details,
277 'timestamp': time.time()
278 })
280 def set_stream_chunk(self, chunk: StreamChunk) -> None:
281 """Set the current stream chunk for processing.
283 Args:
284 chunk: Stream chunk to process.
285 """
286 self.current_chunk = chunk
287 self.processed_chunks += 1
289 def add_batch_item(self, item: Any) -> None:
290 """Add an item to the batch.
292 Args:
293 item: Item to add to batch.
294 """
295 if self.data_mode == ProcessingMode.BATCH:
296 self.batch_data.append(item)
298 def add_batch_result(self, result: Any) -> None:
299 """Add a result to batch results.
301 Args:
302 result: Processing result.
303 """
304 if self.data_mode == ProcessingMode.BATCH:
305 self.batch_results.append(result)
307 def add_batch_error(self, index: int, error: Exception) -> None:
308 """Add an error to batch errors.
310 Args:
311 index: Batch item index.
312 error: Error that occurred.
313 """
314 if self.data_mode == ProcessingMode.BATCH:
315 self.batch_errors.append((index, error))
317 def create_child_context(self, path_id: str) -> 'ExecutionContext':
318 """Create a child context for parallel execution.
320 Args:
321 path_id: Unique identifier for the execution path.
323 Returns:
324 New child execution context.
325 """
326 child = ExecutionContext(
327 data_mode=self.data_mode,
328 transaction_mode=self.transaction_mode,
329 resources=self.resource_limits.copy(),
330 database=self.database,
331 stream_context=self.stream_context
332 )
334 child.is_child_context = True
335 child.parent_context = self
336 child.variables = self.variables.copy()
338 self.parallel_paths[path_id] = child
339 return child
341 def merge_child_context(self, path_id: str) -> bool:
342 """Merge a child context back into parent.
344 Args:
345 path_id: Path identifier to merge.
347 Returns:
348 True if merge successful.
349 """
350 if path_id not in self.parallel_paths:
351 return False
353 child = self.parallel_paths[path_id]
355 # Merge results
356 if self.data_mode == ProcessingMode.BATCH:
357 self.batch_results.extend(child.batch_results)
358 self.batch_errors.extend(child.batch_errors)
360 # Merge metadata
361 self.metadata.update(child.metadata)
363 # Update function call counts
364 for func, count in child.function_call_count.items():
365 self.function_call_count[func] = self.function_call_count.get(func, 0) + count
367 del self.parallel_paths[path_id]
368 return True
370 def get_resource_usage(self) -> Dict[str, Any]:
371 """Get current resource usage statistics.
373 Returns:
374 Resource usage information.
375 """
376 allocated = sum(1 for r in self.resources.values()
377 if r.status == ResourceStatus.ALLOCATED)
378 busy = sum(1 for r in self.resources.values()
379 if r.status == ResourceStatus.BUSY)
381 return {
382 'total_resources': len(self.resources),
383 'allocated': allocated,
384 'busy': busy,
385 'available': len(self.resources) - allocated - busy,
386 'by_type': self._group_resources_by_type()
387 }
389 def _group_resources_by_type(self) -> Dict[str, int]:
390 """Group resources by type.
392 Returns:
393 Resource counts by type.
394 """
395 by_type: Dict[str, int] = {}
396 for resource in self.resources.values():
397 by_type[resource.resource_type] = by_type.get(resource.resource_type, 0) + 1
398 return by_type
400 def get_performance_stats(self) -> Dict[str, Any]:
401 """Get performance statistics.
403 Returns:
404 Performance statistics.
405 """
406 elapsed_time = time.time() - self.start_time
408 return {
409 'elapsed_time': elapsed_time,
410 'states_visited': len(self.state_history),
411 'current_state': self.current_state,
412 'transactions': len(self.transaction_history),
413 'chunks_processed': self.processed_chunks,
414 'batch_items': len(self.batch_data),
415 'batch_results': len(self.batch_results),
416 'batch_errors': len(self.batch_errors),
417 'function_calls': dict(self.function_call_count),
418 'parallel_paths': len(self.parallel_paths)
419 }
421 def get_complete_path(self) -> List[str]:
422 """Get the complete state traversal path including current state.
424 Returns:
425 List of state names in traversal order.
426 """
427 path = self.state_history.copy() if self.state_history else []
429 # Add current state if not already in path and if it exists
430 if self.current_state and (not path or path[-1] != self.current_state):
431 path.append(self.current_state)
433 return path
435 def clone(self) -> 'ExecutionContext':
436 """Create a clone of this context.
438 Returns:
439 Cloned execution context.
440 """
441 clone = ExecutionContext(
442 data_mode=self.data_mode,
443 transaction_mode=self.transaction_mode,
444 resources=self.resource_limits.copy(),
445 database=self.database,
446 stream_context=self.stream_context
447 )
449 clone.current_state = self.current_state
450 clone.previous_state = self.previous_state
451 clone.state_history = self.state_history.copy()
452 clone.data = self.data
453 clone.metadata = self.metadata.copy()
454 clone.variables = self.variables.copy()
456 return clone
458 def is_complete(self) -> bool:
459 """Check if the FSM execution has reached an end state.
461 Returns:
462 True if in an end state or no current state, False otherwise.
463 """
464 if not self.current_state:
465 return True
467 # Check if current state is marked as ended
468 return self.metadata.get('is_end_state', False)
470 def get_current_state(self) -> str | None:
471 """Get the name of the current state.
473 Returns:
474 Current state name or None if not set.
475 """
476 return self.current_state
478 def get_data_snapshot(self) -> Dict[str, Any]:
479 """Get a snapshot of the current data.
481 Returns:
482 Copy of the current data dictionary.
483 """
484 if isinstance(self.data, dict):
485 return self.data.copy()
486 elif hasattr(self.data, '__dict__'):
487 return vars(self.data).copy()
488 else:
489 return {'value': self.data}
491 def get_execution_stats(self) -> Dict[str, Any]:
492 """Get execution statistics.
494 Returns:
495 Dictionary with execution metrics.
496 """
497 return {
498 'states_visited': len(self.state_history),
499 'current_state': self.current_state,
500 'previous_state': self.previous_state,
501 'transition_count': self.transition_count,
502 'execution_id': self.execution_id,
503 'data_mode': self.data_mode.value if self.data_mode else None,
504 'transaction_mode': self.transaction_mode.value if self.transaction_mode else None
505 }
507 def get_current_state_instance(self) -> Any:
508 """Get the current state instance object.
510 Returns:
511 The StateInstance object for the current state, or None if not set.
512 """
513 return self.current_state_instance