Coverage for async-durable-execution/src/async_durable_execution/composite/parallel.py: 99%

1"""Concurrent executor for parallel and map operations.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import json 

7import logging 

8import time 

9from collections import Counter 

10from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence 

11from dataclasses import dataclass, field as dataclass_field 

12from enum import Enum 

13from typing import TYPE_CHECKING, Any, Generic, TypeAlias, TypeVar, cast 

14 

15from ..exceptions import ( 

16 CallableRuntimeError, 

17 InvalidStateError, 

18 SuspendExecution, 

19 TimedSuspendExecution, 

20) 

21from ..models import ( 

22 ErrorObject, 

23 Operation, 

24 OperationIdentifier, 

25 OperationStatus, 

26 OperationSubType, 

27 SerializableModel, 

28 _metadata, 

29) 

30from ..primitive.base import OperationExecutor 

31from ..primitive.child import ( 

32 ChildOperationExecutor, 

33 OrphanedChildException, 

34 _create_child_context_task as _run_in_child_context, 

35 get_durable_context, 

36) 

37from ..context import bind_current_context 

38from ..execution import durable_callable 

39from ..serdes import deserialize 

40 

41if TYPE_CHECKING: 

42 from ..primitive.child import DurableContext 

43 from ..serdes import SerDes 

44 from ..state import ExecutionState 

45 

46 

47logger = logging.getLogger(__name__) 

48 

49CallableType = TypeVar("CallableType") 

50ResultType = TypeVar("ResultType") 

51R = TypeVar("R") 

52T = TypeVar("T") 

53C_contra = TypeVar("C_contra", contravariant=True) 

54 

55SummaryGenerator: TypeAlias = Callable[[C_contra], str] 

56"""Create a compact JSON summary for oversized checkpoint payloads.""" 

57 

58 

59class NestingType(Enum): 

60 """Control how child contexts are created for batch operations.""" 

61 

62 NESTED = "NESTED" 

63 FLAT = "FLAT" 

64 

65 

66@dataclass(frozen=True) 

67class CompletionConfig: 

68 """Configuration for determining when parallel/map operations complete.""" 

69 

70 min_successful: int | None = None 

71 tolerated_failure_count: int | None = None 

72 

73 @staticmethod 

74 def first_successful(): 

75 return CompletionConfig( 

76 min_successful=1, 

77 tolerated_failure_count=None, 

78 ) 

79 

80 @staticmethod 

81 def all_completed(): 

82 return CompletionConfig( 

83 min_successful=None, 

84 tolerated_failure_count=None, 

85 ) 

86 

87 @staticmethod 

88 def all_successful(): 

89 return CompletionConfig( 

90 min_successful=None, 

91 tolerated_failure_count=0, 

92 ) 

93 

94 

95class BatchItemStatus(Enum): 

96 """Status of one item or branch inside a batch-style operation.""" 

97 

98 SUCCEEDED = "SUCCEEDED" 

99 FAILED = "FAILED" 

100 STARTED = "STARTED" 

101 

102 

103class CompletionReason(Enum): 

104 """Why a map or parallel operation stopped collecting results.""" 

105 

106 ALL_COMPLETED = "ALL_COMPLETED" 

107 MIN_SUCCESSFUL_REACHED = "MIN_SUCCESSFUL_REACHED" 

108 FAILURE_TOLERANCE_EXCEEDED = "FAILURE_TOLERANCE_EXCEEDED" 

109 

110 

111@dataclass(frozen=True) 

112class SuspendResult: 

113 """Internal helper describing whether an executor should suspend.""" 

114 

115 should_suspend: bool 

116 exception: SuspendExecution | None = None 

117 

118 @staticmethod 

119 def do_not_suspend() -> SuspendResult: 

120 return SuspendResult(should_suspend=False) 

121 

122 @staticmethod 

123 def suspend(exception: SuspendExecution) -> SuspendResult: 

124 return SuspendResult(should_suspend=True, exception=exception) 

125 

126 

127@dataclass(frozen=True) 

128class BatchItem(SerializableModel, Generic[R]): 

129 """Result record for one branch or iteration in `BatchResult`.""" 

130 

131 index: int 

132 status: BatchItemStatus 

133 result: R | None = dataclass_field( 

134 default=None, metadata=_metadata(alias="result", omit_if_none=False) 

135 ) 

136 error: ErrorObject | None = dataclass_field( 

137 default=None, 

138 metadata=_metadata(alias="error", omit_if_none=False), 

139 ) 

140 

141 

142@dataclass(frozen=True) 

143class BatchResult(SerializableModel, Generic[R]): # noqa: PYI059 

144 """Aggregated outcome of a `map()` or `parallel()` operation.""" 

145 

146 all: list[BatchItem[R]] 

147 completion_reason: CompletionReason = dataclass_field( 

148 metadata=_metadata(alias="completionReason") 

149 ) 

150 

151 @classmethod 

152 def from_dict( 

153 cls, data: Mapping[str, Any], completion_config: CompletionConfig | None = None 

154 ) -> BatchResult[R]: 

155 batch_items = [BatchItem.from_dict(item) for item in data["all"]] 

156 

157 completion_reason_value = data.get("completionReason") 

158 if completion_reason_value is None: 

159 result = cls.from_items(batch_items, completion_config) 

160 logger.warning( 

161 "Missing completionReason in BatchResult deserialization, " 

162 "inferred '%s' from batch item statuses. " 

163 "This may indicate incomplete serialization data.", 

164 result.completion_reason.value, 

165 ) 

166 return result 

167 

168 return cls( 

169 all=batch_items, 

170 completion_reason=CompletionReason(completion_reason_value), 

171 ) 

172 

173 @staticmethod 

174 def _get_completion_reason( 

175 failure_count: int, 

176 success_count: int, 

177 completed_count: int, 

178 total_count: int, 

179 completion_config: CompletionConfig | None, 

180 ) -> CompletionReason: 

181 if completion_config is None: 

182 if failure_count > 0: 

183 return CompletionReason.FAILURE_TOLERANCE_EXCEEDED 

184 else: 

185 has_any_completion_criteria = ( 

186 completion_config.min_successful is not None 

187 or completion_config.tolerated_failure_count is not None 

188 ) 

189 

190 if not has_any_completion_criteria: 

191 if failure_count > 0: 

192 return CompletionReason.FAILURE_TOLERANCE_EXCEEDED 

193 else: 

194 if ( 

195 completion_config.tolerated_failure_count is not None 

196 and failure_count > completion_config.tolerated_failure_count 

197 ): 

198 return CompletionReason.FAILURE_TOLERANCE_EXCEEDED 

199 

200 if completed_count == total_count: 

201 return CompletionReason.ALL_COMPLETED 

202 

203 if ( 

204 completion_config is not None 

205 and completion_config.min_successful is not None 

206 and success_count >= completion_config.min_successful 

207 ): 

208 return CompletionReason.MIN_SUCCESSFUL_REACHED 

209 

210 return CompletionReason.ALL_COMPLETED 

211 

212 @classmethod 

213 def from_items( 

214 cls, 

215 items: list[BatchItem[R]], 

216 completion_config: CompletionConfig | None = None, 

217 ) -> BatchResult[R]: 

218 statuses = (item.status for item in items) 

219 counts = Counter(statuses) 

220 succeeded_count = counts.get(BatchItemStatus.SUCCEEDED, 0) 

221 failed_count = counts.get(BatchItemStatus.FAILED, 0) 

222 started_count = counts.get(BatchItemStatus.STARTED, 0) 

223 

224 completed_count = succeeded_count + failed_count 

225 total_count = started_count + completed_count 

226 

227 completion_reason = cls._get_completion_reason( 

228 failure_count=failed_count, 

229 success_count=succeeded_count, 

230 completed_count=completed_count, 

231 total_count=total_count, 

232 completion_config=completion_config, 

233 ) 

234 

235 return cls(all=items, completion_reason=completion_reason) 

236 

237 def succeeded(self) -> list[BatchItem[R]]: 

238 return [ 

239 item 

240 for item in self.all 

241 if item.status is BatchItemStatus.SUCCEEDED and item.result is not None 

242 ] 

243 

244 def failed(self) -> list[BatchItem[R]]: 

245 return [ 

246 item 

247 for item in self.all 

248 if item.status is BatchItemStatus.FAILED and item.error is not None 

249 ] 

250 

251 def started(self) -> list[BatchItem[R]]: 

252 return [item for item in self.all if item.status is BatchItemStatus.STARTED] 

253 

254 @property 

255 def status(self) -> BatchItemStatus: 

256 return BatchItemStatus.FAILED if self.has_failure else BatchItemStatus.SUCCEEDED 

257 

258 @property 

259 def has_failure(self) -> bool: 

260 return any(item.status is BatchItemStatus.FAILED for item in self.all) 

261 

262 def throw_if_error(self) -> None: 

263 first_error = next( 

264 (item.error for item in self.all if item.status is BatchItemStatus.FAILED), 

265 None, 

266 ) 

267 if first_error: 

268 raise CallableRuntimeError.from_error_object(first_error) 

269 

270 def get_results(self) -> list[R]: 

271 return [ 

272 item.result 

273 for item in self.all 

274 if item.status is BatchItemStatus.SUCCEEDED and item.result is not None 

275 ] 

276 

277 def get_errors(self) -> list[ErrorObject]: 

278 return [item.error for item in self.failed() if item.error is not None] 

279 

280 @property 

281 def success_count(self) -> int: 

282 return sum(1 for item in self.all if item.status is BatchItemStatus.SUCCEEDED) 

283 

284 @property 

285 def failure_count(self) -> int: 

286 return sum(1 for item in self.all if item.status is BatchItemStatus.FAILED) 

287 

288 @property 

289 def started_count(self) -> int: 

290 return sum(1 for item in self.all if item.status is BatchItemStatus.STARTED) 

291 

292 @property 

293 def total_count(self) -> int: 

294 return len(self.all) 

295 

296 

297@dataclass(frozen=True) 

298class Executable(Generic[CallableType]): 

299 """Index plus callable payload used by the concurrent executors.""" 

300 

301 index: int 

302 func: CallableType 

303 

304 

305class BranchStatus(Enum): 

306 """In-memory lifecycle state for a concurrently scheduled branch.""" 

307 

308 PENDING = "pending" 

309 RUNNING = "running" 

310 COMPLETED = "completed" 

311 SUSPENDED = "suspended" 

312 SUSPENDED_WITH_TIMEOUT = "suspended_with_timeout" 

313 FAILED = "failed" 

314 

315 

316class ExecutableWithState(Generic[CallableType, ResultType]): 

317 """Manages the execution state and lifecycle of an executable.""" 

318 

319 def __init__(self, executable: Executable[CallableType]): 

320 self.executable = executable 

321 self._status = BranchStatus.PENDING 

322 self._future: asyncio.Task[ResultType] | None = None 

323 self._suspend_until: float | None = None 

324 self._result: ResultType | None = None 

325 self._is_result_set = False 

326 self._error: Exception | None = None 

327 

328 @property 

329 def future(self) -> asyncio.Task[ResultType]: 

330 if self._future is None: 

331 msg = f"ExecutableWithState was never started. {self.executable.index}" 

332 raise InvalidStateError(msg) 

333 return self._future 

334 

335 @property 

336 def status(self) -> BranchStatus: 

337 return self._status 

338 

339 @property 

340 def result(self) -> ResultType: 

341 if not self._is_result_set or self._status != BranchStatus.COMPLETED: 

342 msg = f"result not available in status {self._status}" 

343 raise InvalidStateError(msg) 

344 return cast("ResultType", self._result) 

345 

346 @property 

347 def error(self) -> Exception: 

348 if self._error is None or self._status != BranchStatus.FAILED: 

349 msg = f"error not available in status {self._status}" 

350 raise InvalidStateError(msg) 

351 return self._error 

352 

353 @property 

354 def suspend_until(self) -> float | None: 

355 return self._suspend_until 

356 

357 @property 

358 def is_running(self) -> bool: 

359 return self._status is BranchStatus.RUNNING 

360 

361 @property 

362 def can_resume(self) -> bool: 

363 return self._status is BranchStatus.SUSPENDED or ( 

364 self._status is BranchStatus.SUSPENDED_WITH_TIMEOUT 

365 and self._suspend_until is not None 

366 and time.time() >= self._suspend_until 

367 ) 

368 

369 @property 

370 def index(self) -> int: 

371 return self.executable.index 

372 

373 @property 

374 def callable(self) -> CallableType: 

375 return self.executable.func 

376 

377 def run(self, future: asyncio.Task[ResultType]) -> None: 

378 if self._status != BranchStatus.PENDING: 

379 msg = f"Cannot start running from {self._status}" 

380 raise InvalidStateError(msg) 

381 self._status = BranchStatus.RUNNING 

382 self._future = future 

383 

384 def suspend(self) -> None: 

385 self._status = BranchStatus.SUSPENDED 

386 self._suspend_until = None 

387 

388 def suspend_with_timeout(self, timestamp: float) -> None: 

389 self._status = BranchStatus.SUSPENDED_WITH_TIMEOUT 

390 self._suspend_until = timestamp 

391 

392 def complete(self, result: ResultType) -> None: 

393 self._status = BranchStatus.COMPLETED 

394 self._result = result 

395 self._is_result_set = True 

396 

397 def fail(self, error: Exception) -> None: 

398 self._status = BranchStatus.FAILED 

399 self._error = error 

400 

401 def reset_to_pending(self) -> None: 

402 self._status = BranchStatus.PENDING 

403 self._future = None 

404 self._suspend_until = None 

405 

406 

407class ExecutionCounters: 

408 """Counters for tracking execution state on a single event loop.""" 

409 

410 def __init__( 

411 self, 

412 total_tasks: int, 

413 min_successful: int, 

414 tolerated_failure_count: int | None, 

415 ): 

416 self.total_tasks = total_tasks 

417 self.min_successful = min_successful 

418 self.tolerated_failure_count = tolerated_failure_count 

419 self.success_count = 0 

420 self.failure_count = 0 

421 

422 def complete_task(self) -> None: 

423 self.success_count += 1 

424 

425 def fail_task(self) -> None: 

426 self.failure_count += 1 

427 

428 def should_continue(self) -> bool: 

429 if self.tolerated_failure_count is None: 

430 return self.failure_count == 0 

431 

432 if ( 

433 self.tolerated_failure_count is not None 

434 and self.failure_count > self.tolerated_failure_count 

435 ): 

436 return False 

437 

438 return True 

439 

440 def is_complete(self) -> bool: 

441 completed_count = self.success_count + self.failure_count 

442 

443 if completed_count == self.total_tasks: 

444 return True 

445 

446 return self.success_count >= self.min_successful 

447 

448 def should_complete(self) -> bool: 

449 return self.is_complete() or not self.should_continue() 

450 

451 def is_all_completed(self) -> bool: 

452 return self.success_count == self.total_tasks 

453 

454 def is_min_successful_reached(self) -> bool: 

455 return self.success_count >= self.min_successful 

456 

457 def is_failure_tolerance_exceeded(self) -> bool: 

458 return self._is_failure_condition_reached( 

459 tolerated_count=self.tolerated_failure_count, 

460 failure_count=self.failure_count, 

461 ) 

462 

463 def _is_failure_condition_reached( 

464 self, 

465 tolerated_count: int | None, 

466 failure_count: int, 

467 ) -> bool: 

468 if tolerated_count is not None and failure_count > tolerated_count: 

469 return True 

470 

471 return False 

472 

473 

474class TimerScheduler: 

475 """Manage timed suspend tasks with event-loop tasks.""" 

476 

477 def __init__( 

478 self, 

479 resubmit_callback: Callable[[ExecutableWithState[Any, Any]], Awaitable[None]], 

480 ) -> None: 

481 self.resubmit_callback = resubmit_callback 

482 self._resume_tasks: set[asyncio.Task[None]] = set() 

483 

484 async def __aenter__(self) -> TimerScheduler: 

485 return self 

486 

487 async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: 

488 await self.shutdown() 

489 

490 def schedule_resume( 

491 self, 

492 exe_state: ExecutableWithState[CallableType, ResultType], 

493 resume_time: float, 

494 ) -> None: 

495 async def runner() -> None: 

496 await asyncio.sleep(max(0.0, resume_time - time.time())) 

497 if exe_state.can_resume: 

498 exe_state.reset_to_pending() 

499 await self.resubmit_callback(exe_state) 

500 

501 task = asyncio.create_task(runner()) 

502 self._resume_tasks.add(task) 

503 task.add_done_callback(self._resume_tasks.discard) 

504 

505 async def shutdown(self) -> None: 

506 tasks = list(self._resume_tasks) 

507 for task in tasks: 

508 task.cancel() 

509 if tasks: 

510 await asyncio.gather(*tasks, return_exceptions=True) 

511 self._resume_tasks.clear() 

512 

513 

514class ParallelExecutor( 

515 OperationExecutor[BatchResult[ResultType]], 

516 Generic[CallableType, ResultType], 

517): 

518 """Execute durable operations concurrently using asyncio tasks.""" 

519 

520 def __init__( 

521 self, 

522 execution_state: ExecutionState, 

523 operation_identifier: OperationIdentifier, 

524 executor_context: DurableContext, 

525 executables: list[Executable[CallableType]], 

526 max_concurrency: int | None, 

527 completion_config: CompletionConfig, 

528 top_level_sub_type: OperationSubType, 

529 iteration_sub_type: OperationSubType, 

530 name_prefix: str, 

531 serdes: SerDes | None, 

532 item_serdes: SerDes | None = None, 

533 summary_generator: SummaryGenerator | None = None, 

534 nesting_type: NestingType = NestingType.NESTED, 

535 branch_namer: Callable[[int], str] | None = None, 

536 ): 

537 super().__init__( 

538 state=execution_state, 

539 operation_identifier=operation_identifier, 

540 ) 

541 self.executor_context = executor_context 

542 self.executables = executables 

543 self.max_concurrency = max_concurrency 

544 self.completion_config = completion_config 

545 self.sub_type_top = top_level_sub_type 

546 self.sub_type_iteration = iteration_sub_type 

547 self.name_prefix = name_prefix 

548 self.summary_generator = summary_generator 

549 self.nesting_type = nesting_type 

550 self._branch_namer = branch_namer 

551 self._completion_event = asyncio.Event() 

552 self._suspend_exception: SuspendExecution | None = None 

553 self._running_tasks: set[asyncio.Task[ResultType]] = set() 

554 

555 min_successful = self.completion_config.min_successful or len(self.executables) 

556 self.counters = ExecutionCounters( 

557 len(executables), 

558 min_successful, 

559 self.completion_config.tolerated_failure_count, 

560 ) 

561 self.executables_with_state: list[ExecutableWithState] = [] 

562 self.serdes = serdes 

563 self.item_serdes = item_serdes 

564 

565 async def execute_item( 

566 self, 

567 child_context: DurableContext, 

568 executable: Executable[CallableType], 

569 ) -> ResultType: 

570 func = cast("Callable[[], Awaitable[ResultType]]", executable.func) 

571 with bind_current_context(child_context): 

572 result: ResultType = await func() 

573 return result 

574 

575 def get_iteration_name(self, index: int) -> str: 

576 if self._branch_namer is not None: 

577 return self._branch_namer(index) 

578 return f"{self.name_prefix}{index}" 

579 

580 async def start(self) -> BatchResult[ResultType]: 

581 return await self.execute() 

582 

583 async def replay(self, operation: Operation) -> BatchResult[ResultType]: 

584 if operation.status is OperationStatus.SUCCEEDED: 

585 return await self.replay_completed(self.state, self.executor_context) 

586 return await self.execute() 

587 

588 async def execute(self) -> BatchResult[ResultType]: 

589 logger.debug( 

590 "▶️ Executing concurrent operation, items: %d", len(self.executables) 

591 ) 

592 

593 if not self.executables: 

594 logger.debug("No items to execute, returning empty result") 

595 return self._create_result() 

596 

597 max_workers = self.max_concurrency or len(self.executables) 

598 semaphore = asyncio.Semaphore(max_workers) 

599 self.executables_with_state = [ 

600 ExecutableWithState(executable=exe) for exe in self.executables 

601 ] 

602 self._completion_event.clear() 

603 self._suspend_exception = None 

604 self._running_tasks.clear() 

605 

606 async def submit_task( 

607 executable_with_state: ExecutableWithState[CallableType, ResultType], 

608 ) -> None: 

609 if self._completion_event.is_set(): 

610 return 

611 

612 async def run_task() -> ResultType: 

613 async with semaphore: 

614 return await self._execute_item_in_child_context( 

615 self.executor_context, 

616 executable_with_state.executable, 

617 ) 

618 

619 task = asyncio.create_task(run_task()) 

620 executable_with_state.run(task) 

621 self._running_tasks.add(task) 

622 

623 def on_done(done_task: asyncio.Task[ResultType]) -> None: 

624 self._running_tasks.discard(done_task) 

625 asyncio.create_task( 

626 self._on_task_complete( 

627 executable_with_state, 

628 done_task, 

629 scheduler, 

630 ) 

631 ) 

632 

633 task.add_done_callback(on_done) 

634 

635 async def resubmitter( 

636 executable_with_state: ExecutableWithState[CallableType, ResultType], 

637 ) -> None: 

638 await self.state.create_checkpoint(is_sync=False) 

639 await submit_task(executable_with_state) 

640 

641 async with TimerScheduler(resubmitter) as scheduler: 

642 for exe_state in self.executables_with_state: 

643 await submit_task(exe_state) 

644 

645 await self._completion_event.wait() 

646 

647 for task in list(self._running_tasks): 

648 if not task.done(): 

649 task.cancel() 

650 if self._running_tasks: 

651 await asyncio.gather(*self._running_tasks, return_exceptions=True) 

652 

653 if self._suspend_exception: 

654 raise self._suspend_exception 

655 

656 return self._create_result() 

657 

658 def should_execution_suspend(self) -> SuspendResult: 

659 earliest_timestamp: float = float("inf") 

660 indefinite_suspend_task: ( 

661 ExecutableWithState[CallableType, ResultType] | None 

662 ) = None 

663 

664 for exe_state in self.executables_with_state: 

665 if exe_state.status in {BranchStatus.PENDING, BranchStatus.RUNNING}: 

666 return SuspendResult.do_not_suspend() 

667 if exe_state.status is BranchStatus.SUSPENDED_WITH_TIMEOUT: 

668 if ( 

669 exe_state.suspend_until 

670 and exe_state.suspend_until < earliest_timestamp 

671 ): 

672 earliest_timestamp = cast(float, exe_state.suspend_until) 

673 elif exe_state.status is BranchStatus.SUSPENDED: 

674 indefinite_suspend_task = exe_state 

675 

676 if earliest_timestamp != float("inf"): 

677 return SuspendResult.suspend( 

678 TimedSuspendExecution( 

679 "All concurrent work complete or suspended pending retry.", 

680 earliest_timestamp, 

681 ) 

682 ) 

683 if indefinite_suspend_task: 

684 return SuspendResult.suspend( 

685 SuspendExecution( 

686 "All concurrent work complete or suspended and pending external callback." 

687 ) 

688 ) 

689 

690 return SuspendResult.do_not_suspend() 

691 

692 async def _on_task_complete( 

693 self, 

694 exe_state: ExecutableWithState[CallableType, ResultType], 

695 task: asyncio.Task[ResultType], 

696 scheduler: TimerScheduler, 

697 ) -> None: 

698 if task.cancelled(): 

699 exe_state.suspend() 

700 return 

701 

702 try: 

703 result = task.result() 

704 exe_state.complete(result) 

705 self.counters.complete_task() 

706 except OrphanedChildException: 

707 logger.debug( 

708 "Terminating orphaned branch %s without error because parent has completed already", 

709 exe_state.index, 

710 ) 

711 return 

712 except TimedSuspendExecution as tse: 

713 exe_state.suspend_with_timeout(tse.scheduled_timestamp) 

714 scheduler.schedule_resume(exe_state, tse.scheduled_timestamp) 

715 except SuspendExecution: 

716 exe_state.suspend() 

717 except Exception as e: # noqa: BLE001 

718 exe_state.fail(e) 

719 self.counters.fail_task() 

720 

721 if self.counters.should_complete(): 

722 self._completion_event.set() 

723 else: 

724 suspend_result = self.should_execution_suspend() 

725 if suspend_result.should_suspend: 

726 self._suspend_exception = suspend_result.exception 

727 self._completion_event.set() 

728 

729 def _create_result(self) -> BatchResult[ResultType]: 

730 batch_items: list[BatchItem[ResultType]] = [] 

731 for executable in self.executables_with_state: 

732 match executable.status: 

733 case BranchStatus.COMPLETED: 

734 batch_items.append( 

735 BatchItem( 

736 executable.index, 

737 BatchItemStatus.SUCCEEDED, 

738 executable.result, 

739 ) 

740 ) 

741 case BranchStatus.FAILED: 

742 batch_items.append( 

743 BatchItem( 

744 executable.index, 

745 BatchItemStatus.FAILED, 

746 error=ErrorObject.from_exception(executable.error), 

747 ) 

748 ) 

749 case ( 

750 BranchStatus.PENDING 

751 | BranchStatus.RUNNING 

752 | BranchStatus.SUSPENDED 

753 | BranchStatus.SUSPENDED_WITH_TIMEOUT 

754 ): 

755 batch_items.append( 

756 BatchItem(executable.index, BatchItemStatus.STARTED) 

757 ) 

758 

759 return BatchResult.from_items(batch_items, self.completion_config) 

760 

761 async def _execute_item_in_child_context( 

762 self, 

763 executor_context: DurableContext, 

764 executable: Executable[CallableType], 

765 ) -> ResultType: 

766 operation_id: str = ( 

767 executor_context.step_counter._create_step_id_for_logical_step( # noqa: SLF001 

768 executable.index 

769 ) 

770 ) 

771 name: str = self.get_iteration_name(executable.index) 

772 is_virtual: bool = self.nesting_type is NestingType.FLAT 

773 

774 child_context: DurableContext = executor_context.create_child_context( 

775 operation_id, is_virtual=is_virtual 

776 ) 

777 operation_identifier = OperationIdentifier( 

778 operation_id=operation_id, 

779 sub_type=self.sub_type_iteration, 

780 parent_id=executor_context.parent_id, # noqa: SLF001 

781 name=name, 

782 ) 

783 

784 async def run_child_operation() -> ResultType: 

785 return await self.execute_item(child_context, executable) 

786 

787 executor: ChildOperationExecutor[ResultType] = ChildOperationExecutor( 

788 run_child_operation, 

789 child_context.execution_state, 

790 operation_identifier, 

791 serdes=self.item_serdes or self.serdes, 

792 summary_generator=self.summary_generator, 

793 is_virtual=is_virtual, 

794 ) 

795 return await executor.process() 

796 

797 async def replay_completed( 

798 self, execution_state: ExecutionState, executor_context: DurableContext 

799 ) -> BatchResult[ResultType]: 

800 items: list[BatchItem[ResultType]] = [] 

801 for executable in self.executables: 

802 operation_id = ( 

803 executor_context.step_counter._create_step_id_for_logical_step( # noqa: SLF001 

804 executable.index 

805 ) 

806 ) 

807 operation = execution_state.operations.get(operation_id) 

808 

809 result: ResultType | None = None 

810 error = None 

811 status: BatchItemStatus 

812 if operation is not None and operation.status is OperationStatus.SUCCEEDED: 

813 status = BatchItemStatus.SUCCEEDED 

814 operation_details = operation.context_details 

815 if operation_details is not None and operation_details.replay_children: 

816 result = await self._execute_item_in_child_context( 

817 executor_context, executable 

818 ) 

819 elif ( 

820 operation_details is not None 

821 and operation_details.result is not None 

822 ): 

823 result = await deserialize( 

824 serdes=self.item_serdes or self.serdes, 

825 data=operation_details.result, 

826 operation_id=operation_id, 

827 durable_execution_arn=execution_state.durable_execution_arn, 

828 ) 

829 elif operation is not None and operation.status is OperationStatus.FAILED: