Coverage for src/dataknobs_fsm/execution/context.py: 69%

1"""Execution context for FSM state machines."""

2

3 import time

4 from dataclasses import dataclass, field

5 from enum import Enum

6 from typing import Any, Dict, List, Tuple, Union

7

8 from dataknobs_data.database import AsyncDatabase, SyncDatabase

9

10 from dataknobs_fsm.core.modes import ProcessingMode, TransactionMode

11 from dataknobs_fsm.streaming.core import StreamChunk, StreamContext

12

13

14 class ResourceStatus(Enum):

15 """Status of a resource in the execution context."""

16 AVAILABLE = "available"

17 ALLOCATED = "allocated"

18 BUSY = "busy"

19 ERROR = "error"

20

21

22@dataclass

23 class 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)

30

31

32@dataclass

33 class 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

41

42

43 class ExecutionContext:

44 """Execution context for FSM processing with full mode support.

45

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 """

54

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.

64

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

75

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] = []

81

82 # Data management

83 self.data: Any = None

84 self.metadata: Dict[str, Any] = {}

85 self.variables: Dict[str, Any] = {}

86

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)

93

94 # Transaction management

95 self.database = database

96 self.current_transaction: TransactionInfo | None = None

97 self.transaction_history: List[TransactionInfo] = []

98

99 # Stream coordination

100 self.stream_context = stream_context

101 self.current_chunk: StreamChunk | None = None

102 self.processed_chunks: int = 0

103

104 # Batch processing

105 self.batch_data: List[Any] = []

106 self.batch_results: List[Any] = []

107 self.batch_errors: List[Tuple[int, Exception]] = []

108

109 # Parallel execution

110 self.parallel_paths: Dict[str, ExecutionContext] = {}

111 self.is_child_context: bool = False

112 self.parent_context: ExecutionContext | None = None

113

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] = {}

118

119 # State instance tracking for debugging

120 self.current_state_instance: Any = None

121

122 def push_network(self, network_name: str, return_state: str | None = None) -> None:

123 """Push a network onto the execution stack.

124

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))

130

131 def pop_network(self) -> Tuple[str, str | None]:

132 """Pop a network from the execution stack.

133

134 Returns:

135 Tuple of (network_name, return_state).

136 """

137 if self.network_stack:

138 return self.network_stack.pop()

139 return ("", None)

140

141 def set_state(self, state_name: str) -> None:

142 """Set the current state.

143

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()

152

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.

160

161 Args:

162 resource_type: Type of resource.

163 resource_id: Unique resource identifier.

164 metadata: Optional resource metadata.

165

166 Returns:

167 True if allocation successful.

168 """

169 key = f"{resource_type}:{resource_id}"

170

171 if key in self.resources:

172 if self.resources[key].status != ResourceStatus.AVAILABLE:

173 return False

174

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

183

184 def release_resource(self, resource_type: str, resource_id: str) -> bool:

185 """Release an allocated resource.

186

187 Args:

188 resource_type: Type of resource.

189 resource_id: Resource identifier.

190

191 Returns:

192 True if release successful.

193 """

194 key = f"{resource_type}:{resource_id}"

195

196 if key in self.resources:

197 self.resources[key].status = ResourceStatus.AVAILABLE

198 return True

199 return False

200

201 def start_transaction(self, transaction_id: str | None = None) -> bool:

202 """Start a new transaction.

203

204 Args:

205 transaction_id: Optional transaction ID.

206

207 Returns:

208 True if transaction started.

209 """

210 if self.transaction_mode == TransactionMode.NONE:

211 return False

212

213 if self.current_transaction and not self.current_transaction.is_committed:

214 return False

215

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 )

221

222 # Start database transaction if available

223 if self.database and hasattr(self.database, 'begin_transaction'):

224 self.database.begin_transaction()

225

226 return True

227

228 def commit_transaction(self) -> bool:

229 """Commit the current transaction.

230

231 Returns:

232 True if commit successful.

233 """

234 if not self.current_transaction:

235 return False

236

237 # Commit database transaction

238 if self.database and hasattr(self.database, 'commit'):

239 self.database.commit()

240

241 self.current_transaction.is_committed = True

242 self.transaction_history.append(self.current_transaction)

243 self.current_transaction = None

244

245 return True

246

247 def rollback_transaction(self) -> bool:

248 """Rollback the current transaction.

249

250 Returns:

251 True if rollback successful.

252 """

253 if not self.current_transaction:

254 return False

255

256 # Rollback database transaction

257 if self.database and hasattr(self.database, 'rollback'):

258 self.database.rollback()

259

260 self.current_transaction.is_rolled_back = True

261 self.transaction_history.append(self.current_transaction)

262 self.current_transaction = None

263

264 return True

265

266 def log_operation(self, operation: str, details: Dict[str, Any]) -> None:

267 """Log an operation in the current transaction.

268

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 })

279

280 def set_stream_chunk(self, chunk: StreamChunk) -> None:

281 """Set the current stream chunk for processing.

282

283 Args:

284 chunk: Stream chunk to process.

285 """

286 self.current_chunk = chunk

287 self.processed_chunks += 1

288

289 def add_batch_item(self, item: Any) -> None:

290 """Add an item to the batch.

291

292 Args:

293 item: Item to add to batch.

294 """

295 if self.data_mode == ProcessingMode.BATCH:

296 self.batch_data.append(item)

297

298 def add_batch_result(self, result: Any) -> None:

299 """Add a result to batch results.

300

301 Args:

302 result: Processing result.

303 """

304 if self.data_mode == ProcessingMode.BATCH:

305 self.batch_results.append(result)

306

307 def add_batch_error(self, index: int, error: Exception) -> None:

308 """Add an error to batch errors.

309

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))

316

317 def create_child_context(self, path_id: str) -> 'ExecutionContext':

318 """Create a child context for parallel execution.

319

320 Args:

321 path_id: Unique identifier for the execution path.

322

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 )

333

334 child.is_child_context = True

335 child.parent_context = self

336 child.variables = self.variables.copy()

337

338 self.parallel_paths[path_id] = child

339 return child

340

341 def merge_child_context(self, path_id: str) -> bool:

342 """Merge a child context back into parent.

343

344 Args:

345 path_id: Path identifier to merge.

346

347 Returns:

348 True if merge successful.

349 """

350 if path_id not in self.parallel_paths:

351 return False

352

353 child = self.parallel_paths[path_id]

354

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)

359

360 # Merge metadata

361 self.metadata.update(child.metadata)

362

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

366

367 del self.parallel_paths[path_id]

368 return True

369

370 def get_resource_usage(self) -> Dict[str, Any]:

371 """Get current resource usage statistics.

372

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)

380

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 }

388

389 def _group_resources_by_type(self) -> Dict[str, int]:

390 """Group resources by type.

391

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

399

400 def get_performance_stats(self) -> Dict[str, Any]:

401 """Get performance statistics.

402

403 Returns:

404 Performance statistics.

405 """

406 elapsed_time = time.time() - self.start_time

407

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 }

420

421 def get_complete_path(self) -> List[str]:

422 """Get the complete state traversal path including current state.

423

424 Returns:

425 List of state names in traversal order.

426 """

427 path = self.state_history.copy() if self.state_history else []

428

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)

432

433 return path

434

435 def clone(self) -> 'ExecutionContext':

436 """Create a clone of this context.

437

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 )

448

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()

455

456 return clone

457

458 def is_complete(self) -> bool:

459 """Check if the FSM execution has reached an end state.

460

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

466

467 # Check if current state is marked as ended

468 return self.metadata.get('is_end_state', False)

469

470 def get_current_state(self) -> str | None:

471 """Get the name of the current state.

472

473 Returns:

474 Current state name or None if not set.

475 """

476 return self.current_state

477

478 def get_data_snapshot(self) -> Dict[str, Any]:

479 """Get a snapshot of the current data.

480

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}

490

491 def get_execution_stats(self) -> Dict[str, Any]:

492 """Get execution statistics.

493

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 }

506

507 def get_current_state_instance(self) -> Any:

508 """Get the current state instance object.

509

510 Returns:

511 The StateInstance object for the current state, or None if not set.

512 """

513 return self.current_state_instance