Coverage for async_durable_execution/_operation/parallel.py: 94%
717 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +0000
1"""Concurrent executor for parallel and map operations."""
3from __future__ import annotations
5import asyncio
6import json
7import logging
8import time
9from collections import Counter
10from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
11from dataclasses import dataclass, field as dataclass_field
12from enum import Enum
13from typing import TYPE_CHECKING, Any, Generic, TypeAlias, TypeVar, cast
15from .._core import (
16 CallableRuntimeError,
17 DurableContext,
18 EncodedValue,
19 ErrorObject,
20 ExecutionState,
21 ExtendedTypeSerDes,
22 InvalidStateError,
23 InvocationError,
24 Operation,
25 OperationIdentifier,
26 OperationStatus,
27 OperationSubType,
28 OperationType,
29 OrphanedChildException,
30 SerDes,
31 SerDesError,
32 SuspendExecution,
33 TimedSuspendExecution,
34 TypeTag,
35 ValidationError,
36 MappingModel,
37 bind_current_context,
38 deserialize,
39 durable_callable,
40 get_durable_context,
41)
42from .._primitive.base import OperationExecutor
43from .._primitive.child import ChildOperationExecutor
44from ..extension import ExtensionContext, ExtensionOperation, get_extension_context
46if TYPE_CHECKING:
47 from .child import SummaryGenerator
50logger = logging.getLogger(__name__)
52CallableType = TypeVar("CallableType")
53ResultType = TypeVar("ResultType")
54R = TypeVar("R")
55T = TypeVar("T")
58def _run_in_child_context(
59 func: Callable[[], Awaitable[T]],
60 *,
61 sub_type: OperationSubType,
62 name: str | None = None,
63 serdes: SerDes | None = None,
64 summary_generator: SummaryGenerator | None = None,
65 is_virtual: bool = False,
66) -> asyncio.Task[T]:
67 """Run an SDK-owned child operation through the stable operation SPI."""
68 return (
69 get_extension_context()
70 ._reserve_sdk_operation(name) # noqa: SLF001
71 ._run_in_child_context( # noqa: SLF001
72 func,
73 sub_type=sub_type,
74 serdes=serdes,
75 summary_generator=summary_generator,
76 is_virtual=is_virtual,
77 )
78 )
81class CompletionReason(Enum):
82 """Why a `map()` or `parallel()` operation stopped collecting results.
84 Values:
85 ALL_COMPLETED: Every item or branch reached a terminal state and no
86 earlier success, failure, or custom completion condition applied.
87 MIN_SUCCESSFUL_REACHED: The configured `min_successful` threshold was
88 reached.
89 FAILURE_TOLERANCE_EXCEEDED: The number of failures exceeded the
90 configured `tolerated_failure_count`, or no failure tolerance was
91 configured and at least one failure was observed.
92 CUSTOM_COMPLETION_SUCCEEDED: A custom completion function completed the
93 operation successfully.
94 CUSTOM_COMPLETION_FAILED: A custom completion function completed the
95 operation as failed.
96 """
98 ALL_COMPLETED = "ALL_COMPLETED"
99 MIN_SUCCESSFUL_REACHED = "MIN_SUCCESSFUL_REACHED"
100 FAILURE_TOLERANCE_EXCEEDED = "FAILURE_TOLERANCE_EXCEEDED"
101 CUSTOM_COMPLETION_SUCCEEDED = "CUSTOM_COMPLETION_SUCCEEDED"
102 CUSTOM_COMPLETION_FAILED = "CUSTOM_COMPLETION_FAILED"
104 @property
105 def is_succeeded(self) -> bool:
106 """Whether this completion reason represents successful completion."""
107 return self in {
108 CompletionReason.ALL_COMPLETED,
109 CompletionReason.MIN_SUCCESSFUL_REACHED,
110 CompletionReason.CUSTOM_COMPLETION_SUCCEEDED,
111 }
114@dataclass(frozen=True)
115class CompletionStatus:
116 """Live completion progress passed to a custom completion function.
118 Attributes:
119 success_count: Number of items or branches that have completed
120 successfully.
121 failure_count: Number of items or branches that have failed.
122 total_count: Total number of items or branches registered for the
123 operation.
124 completed_count: Calculated number of terminal items or branches,
125 equal to `success_count + failure_count`.
126 all_completed: Whether every registered item or branch is terminal.
128 Raises:
129 ValueError: If any count is negative, or if completed count exceeds
130 `total_count`.
131 """
133 success_count: int
134 failure_count: int
135 total_count: int
137 def __post_init__(self) -> None:
138 counts = (
139 self.success_count,
140 self.failure_count,
141 self.total_count,
142 )
143 if any(count < 0 for count in counts): 143 ↛ 144line 143 didn't jump to line 144 because the condition on line 143 was never true
144 msg = "completion counts must be non-negative"
145 raise ValueError(msg)
146 if self.completed_count > self.total_count:
147 msg = "completed_count cannot exceed total_count"
148 raise ValueError(msg)
150 @property
151 def completed_count(self) -> int:
152 """Number of items that have reached a terminal state."""
153 return self.success_count + self.failure_count
155 @property
156 def all_completed(self) -> bool:
157 """Whether all items have reached a terminal state."""
158 return self.completed_count == self.total_count
161@dataclass(frozen=True)
162class CompletionDecision:
163 """Decision returned by a completion condition.
165 Args:
166 should_complete: Whether the operation should stop collecting results.
167 completion_reason: Required when `should_complete` is `True`, and must
168 be `None` when `should_complete` is `False`.
170 Raises:
171 ValueError: If `completion_reason` is missing for a complete decision,
172 or present for a continue decision.
173 """
175 should_complete: bool
176 completion_reason: CompletionReason | None = None
178 def __post_init__(self) -> None:
179 if self.should_complete and self.completion_reason is None:
180 msg = "completion_reason is required when should_complete is true"
181 raise ValueError(msg)
182 if not self.should_complete and self.completion_reason is not None:
183 msg = "completion_reason must be None when should_complete is false"
184 raise ValueError(msg)
186 @staticmethod
187 def complete(completion_reason: CompletionReason) -> CompletionDecision:
188 """Create a decision that completes the operation.
190 Args:
191 completion_reason: Reason to store on the resulting `BatchResult`.
193 Returns:
194 A `CompletionDecision` with `should_complete=True`.
195 """
196 return CompletionDecision(True, completion_reason)
198 @staticmethod
199 def continue_execution() -> CompletionDecision:
200 """Create a decision that keeps collecting item or branch results.
202 Returns:
203 A `CompletionDecision` with `should_complete=False`.
204 """
205 return CompletionDecision(False)
207 @property
208 def is_succeeded(self) -> bool:
209 """Whether this decision completes the operation successfully."""
210 return (
211 self.should_complete
212 and self.completion_reason is not None
213 and self.completion_reason.is_succeeded
214 )
217ShouldComplete: TypeAlias = Callable[[CompletionStatus], CompletionDecision]
218"""Callable used by `CompletionConfig.custom()` to decide batch completion."""
221class NestingType(Enum):
222 """Control how child contexts are created for batch operations."""
224 NESTED = "NESTED"
225 FLAT = "FLAT"
228@dataclass(frozen=True)
229class CompletionConfig:
230 """Configuration for determining when parallel/map operations complete.
232 Without `should_complete`, completion is evaluated in this order:
234 1. Complete successfully when `success_count >= min_successful`, if
235 `min_successful` is configured.
236 2. Complete as failed when `failure_count > tolerated_failure_count`, if
237 `tolerated_failure_count` is configured.
238 3. Complete as failed when `tolerated_failure_count` is `None` and at least
239 one failure is observed.
240 4. Complete successfully when every item or branch has completed.
242 If `should_complete` is configured, it fully controls the completion
243 decision and must return a `CompletionDecision`.
245 Args:
246 min_successful: Optional success threshold. Reaching this count
247 completes the operation successfully.
248 tolerated_failure_count: Optional failure tolerance. Failures complete
249 the operation as failed only after they exceed this count. When this
250 is `None`, any observed failure fails the operation unless the
251 success threshold has already been reached.
252 should_complete: Optional custom completion function. This is mutually
253 exclusive with `min_successful` and `tolerated_failure_count`.
255 Raises:
256 TypeError: If `should_complete` is provided but is not callable.
257 ValueError: If `should_complete` is combined with threshold fields.
258 """
260 min_successful: int | None = None
261 tolerated_failure_count: int | None = None
262 should_complete: ShouldComplete | None = None
264 def __post_init__(self) -> None:
265 if self.should_complete is not None and not callable(self.should_complete): 265 ↛ 266line 265 didn't jump to line 266 because the condition on line 265 was never true
266 msg = "should_complete must be callable"
267 raise TypeError(msg)
268 if self.should_complete is not None and (
269 self.min_successful is not None or self.tolerated_failure_count is not None
270 ):
271 msg = (
272 "should_complete is mutually exclusive with min_successful "
273 "and tolerated_failure_count"
274 )
275 raise ValueError(msg)
277 @classmethod
278 def thresholds(
279 cls,
280 *,
281 min_successful: int | None = None,
282 tolerated_failure_count: int | None = None,
283 ) -> CompletionConfig:
284 """Create a threshold-based completion configuration.
286 Args:
287 min_successful: Optional success threshold. The operation completes
288 successfully once this many items or branches succeed.
289 tolerated_failure_count: Optional failure tolerance. The operation
290 completes as failed once failures exceed this count.
292 Returns:
293 A `CompletionConfig` using the supplied threshold fields.
294 """
295 return cls(
296 min_successful=min_successful,
297 tolerated_failure_count=tolerated_failure_count,
298 )
300 @classmethod
301 def first_successful(cls) -> CompletionConfig:
302 """Create a configuration that completes after the first success.
304 Returns:
305 A `CompletionConfig` with `min_successful=1` and no explicit failure
306 tolerance. If a failure is observed before any success, the
307 operation completes as failed.
308 """
309 return cls(
310 min_successful=1,
311 tolerated_failure_count=None,
312 )
314 @classmethod
315 def all_completed(cls) -> CompletionConfig:
316 """Create a configuration with no explicit thresholds.
318 Returns:
319 A `CompletionConfig` with both threshold fields set to `None`. The
320 operation completes successfully when all work completes without
321 failures, and completes as failed when any failure is observed.
322 """
323 return cls(
324 min_successful=None,
325 tolerated_failure_count=None,
326 )
328 @classmethod
329 def all_successful(cls) -> CompletionConfig:
330 """Create a configuration that requires every item or branch to succeed.
332 Returns:
333 A `CompletionConfig` with `tolerated_failure_count=0`. The first
334 failure exceeds the zero-failure tolerance and completes the
335 operation as failed.
336 """
337 return cls(
338 min_successful=None,
339 tolerated_failure_count=0,
340 )
342 @classmethod
343 def custom(cls, should_complete: ShouldComplete) -> CompletionConfig:
344 """Create a configuration that delegates completion to a callback.
346 Args:
347 should_complete: Deterministic callable that receives a
348 `CompletionStatus` and returns a `CompletionDecision`.
350 Returns:
351 A `CompletionConfig` that uses the supplied callback.
352 """
353 return cls(should_complete=should_complete)
355 @property
356 def has_custom_should_complete(self) -> bool:
357 """Whether a custom completion function is configured."""
358 return self.should_complete is not None
360 def completion_decision(self, status: CompletionStatus) -> CompletionDecision:
361 """Evaluate whether the supplied progress status should complete.
363 Args:
364 status: Current completion progress for a `map()` or `parallel()`
365 operation.
367 Returns:
368 A `CompletionDecision` describing whether execution should continue
369 and, if complete, why.
371 Raises:
372 TypeError: If a custom completion callback returns `None`.
373 """
374 if self.should_complete is not None: 374 ↛ 381line 374 didn't jump to line 381 because the condition on line 374 was always true
375 decision = self.should_complete(status)
376 if decision is None:
377 msg = "should_complete must return a CompletionDecision"
378 raise TypeError(msg)
379 return decision
381 if (
382 self.min_successful is not None
383 and status.success_count >= self.min_successful
384 ):
385 return CompletionDecision.complete(CompletionReason.MIN_SUCCESSFUL_REACHED)
387 if (
388 self.tolerated_failure_count is not None
389 and status.failure_count > self.tolerated_failure_count
390 ):
391 return CompletionDecision.complete(
392 CompletionReason.FAILURE_TOLERANCE_EXCEEDED
393 )
395 if self.tolerated_failure_count is None and status.failure_count > 0:
396 return CompletionDecision.complete(
397 CompletionReason.FAILURE_TOLERANCE_EXCEEDED
398 )
400 if status.all_completed:
401 return CompletionDecision.complete(CompletionReason.ALL_COMPLETED)
403 return CompletionDecision.continue_execution()
406def _validate_max_concurrency(max_concurrency: int | None) -> None:
407 if max_concurrency is not None and (
408 isinstance(max_concurrency, bool)
409 or not isinstance(max_concurrency, int)
410 or max_concurrency < 1
411 ):
412 msg = "max_concurrency must be a positive integer or None"
413 raise ValidationError(msg)
416class BatchItemStatus(Enum):
417 """Status of one item or branch inside a batch-style operation.
419 A ``CANCELLED`` item started but did not finish before the parent reached
420 an early completion condition. Its cancellation is stored in the parent
421 ``BatchResult`` rather than checkpointed as a child result.
422 """
424 SUCCEEDED = "SUCCEEDED"
425 FAILED = "FAILED"
426 CANCELLED = "CANCELLED"
427 STARTED = "STARTED"
430@dataclass(frozen=True)
431class SuspendResult:
432 """Internal helper describing whether an executor should suspend."""
434 should_suspend: bool
435 exception: SuspendExecution | None = None
437 @staticmethod
438 def do_not_suspend() -> SuspendResult:
439 return SuspendResult(should_suspend=False)
441 @staticmethod
442 def suspend(exception: SuspendExecution) -> SuspendResult:
443 return SuspendResult(should_suspend=True, exception=exception)
446@dataclass(frozen=True)
447class BatchItem(MappingModel, Generic[R]):
448 """Result record for one branch or iteration in `BatchResult`."""
450 index: int
451 status: BatchItemStatus
452 result: R | None = dataclass_field(
453 default=None,
454 metadata={"omit_if_none": False},
455 )
456 error: ErrorObject | None = dataclass_field(
457 default=None,
458 metadata={"omit_if_none": False},
459 )
462@dataclass(frozen=True)
463class BatchResult(MappingModel, Generic[R]):
464 """Aggregated outcome of a `map()` or `parallel()` operation."""
466 all: list[BatchItem[R]]
467 completion_reason: CompletionReason = dataclass_field(
468 metadata={"alias": "completionReason"}
469 )
471 @classmethod
472 def from_dict(
473 cls, data: Mapping[str, Any], completion_config: CompletionConfig | None = None
474 ) -> BatchResult[R]:
475 batch_items: list[BatchItem[R]] = [
476 BatchItem.from_dict(item) for item in data["all"]
477 ]
479 completion_reason_value = data.get("completionReason")
480 if completion_reason_value is None:
481 result = cls.from_items(batch_items, completion_config)
482 logger.warning(
483 "Missing completionReason in BatchResult deserialization, "
484 "inferred '%s' from batch item statuses. "
485 "This may indicate incomplete serialization data.",
486 result.completion_reason.value,
487 )
488 return result
490 return cls(
491 all=batch_items,
492 completion_reason=CompletionReason(completion_reason_value),
493 )
495 @staticmethod
496 def _get_completion_reason(
497 failure_count: int,
498 success_count: int,
499 completed_count: int,
500 total_count: int,
501 completion_config: CompletionConfig | None,
502 ) -> CompletionReason:
503 if completion_config is None:
504 if failure_count > 0:
505 return CompletionReason.FAILURE_TOLERANCE_EXCEEDED
506 else:
507 if completion_config.has_custom_should_complete:
508 status = CompletionStatus(
509 success_count=success_count,
510 failure_count=failure_count,
511 total_count=total_count,
512 )
513 decision = completion_config.completion_decision(status)
514 if decision.should_complete and decision.completion_reason is not None: 514 ↛ 517line 514 didn't jump to line 517 because the condition on line 514 was always true
515 return decision.completion_reason
517 return CompletionReason.ALL_COMPLETED
519 has_any_completion_criteria = (
520 completion_config.min_successful is not None
521 or completion_config.tolerated_failure_count is not None
522 )
524 if not has_any_completion_criteria:
525 if failure_count > 0:
526 return CompletionReason.FAILURE_TOLERANCE_EXCEEDED
527 else:
528 if (
529 completion_config.tolerated_failure_count is not None
530 and failure_count > completion_config.tolerated_failure_count
531 ):
532 return CompletionReason.FAILURE_TOLERANCE_EXCEEDED
534 if completed_count == total_count:
535 return CompletionReason.ALL_COMPLETED
537 if (
538 completion_config is not None
539 and completion_config.min_successful is not None
540 and success_count >= completion_config.min_successful
541 ):
542 return CompletionReason.MIN_SUCCESSFUL_REACHED
544 return CompletionReason.ALL_COMPLETED
546 @classmethod
547 def from_items(
548 cls,
549 items: list[BatchItem[R]],
550 completion_config: CompletionConfig | None = None,
551 ) -> BatchResult[R]:
552 statuses = (item.status for item in items)
553 counts = Counter(statuses)
554 succeeded_count = counts.get(BatchItemStatus.SUCCEEDED, 0)
555 failed_count = counts.get(BatchItemStatus.FAILED, 0)
556 cancelled_count = counts.get(BatchItemStatus.CANCELLED, 0)
557 started_count = counts.get(BatchItemStatus.STARTED, 0)
559 completed_count = succeeded_count + failed_count
560 total_count = completed_count + started_count + cancelled_count
562 completion_reason = cls._get_completion_reason(
563 failure_count=failed_count,
564 success_count=succeeded_count,
565 completed_count=completed_count,
566 total_count=total_count,
567 completion_config=completion_config,
568 )
570 return cls(all=items, completion_reason=completion_reason)
572 def succeeded(self) -> list[BatchItem[R]]:
573 return [
574 item
575 for item in self.all
576 if item.status is BatchItemStatus.SUCCEEDED and item.result is not None
577 ]
579 def failed(self) -> list[BatchItem[R]]:
580 return [
581 item
582 for item in self.all
583 if item.status is BatchItemStatus.FAILED and item.error is not None
584 ]
586 def started(self) -> list[BatchItem[R]]:
587 return [item for item in self.all if item.status is BatchItemStatus.STARTED]
589 def cancelled(self) -> list[BatchItem[R]]:
590 return [item for item in self.all if item.status is BatchItemStatus.CANCELLED]
592 @property
593 def status(self) -> BatchItemStatus:
594 return BatchItemStatus.FAILED if self.has_failure else BatchItemStatus.SUCCEEDED
596 @property
597 def has_failure(self) -> bool:
598 return any(item.status is BatchItemStatus.FAILED for item in self.all)
600 def throw_if_error(self) -> None:
601 first_error = next(
602 (item.error for item in self.all if item.status is BatchItemStatus.FAILED),
603 None,
604 )
605 if first_error:
606 raise CallableRuntimeError.from_error_object(first_error)
608 def get_results(self) -> list[R]:
609 return [
610 item.result
611 for item in self.all
612 if item.status is BatchItemStatus.SUCCEEDED and item.result is not None
613 ]
615 def get_errors(self) -> list[ErrorObject]:
616 return [item.error for item in self.failed() if item.error is not None]
618 @property
619 def success_count(self) -> int:
620 return sum(1 for item in self.all if item.status is BatchItemStatus.SUCCEEDED)
622 @property
623 def failure_count(self) -> int:
624 return sum(1 for item in self.all if item.status is BatchItemStatus.FAILED)
626 @property
627 def started_count(self) -> int:
628 return sum(1 for item in self.all if item.status is BatchItemStatus.STARTED)
630 @property
631 def cancelled_count(self) -> int:
632 return sum(1 for item in self.all if item.status is BatchItemStatus.CANCELLED)
634 @property
635 def total_count(self) -> int:
636 return len(self.all)
639_BATCH_RESULT_TAG = "br"
642def _batch_result_payload(value: BatchResult[Any]) -> dict[str, Any]:
643 return {
644 "all": [
645 {
646 "index": item.index,
647 "status": item.status.value,
648 "result": item.result,
649 "error": item.error.to_dict() if item.error is not None else None,
650 }
651 for item in value.all
652 ],
653 "completionReason": value.completion_reason.value,
654 }
657class _BatchResultCodec:
658 """Extended type codec owned by the parallel operation."""
660 tag = _BATCH_RESULT_TAG
662 @staticmethod
663 def can_encode(obj: Any) -> bool:
664 return isinstance(obj, BatchResult)
666 @staticmethod
667 def encode(
668 obj: Any,
669 encode_value: Callable[[Any], EncodedValue],
670 ) -> Any:
671 encoded = encode_value(_batch_result_payload(cast("BatchResult[Any]", obj)))
672 if encoded.tag != "m": 672 ↛ 673line 672 didn't jump to line 673 because the condition on line 672 was never true
673 msg = "Serialized BatchResult value must contain a mapping."
674 raise SerDesError(msg)
675 return encoded.value
677 @staticmethod
678 def decode(
679 value: Any,
680 decode_value: Callable[[TypeTag | str, Any], Any],
681 ) -> BatchResult[Any]:
682 decoded = decode_value("m", value)
683 if not isinstance(decoded, Mapping): 683 ↛ 684line 683 didn't jump to line 684 because the condition on line 683 was never true
684 msg = "Serialized BatchResult value must contain a mapping."
685 raise SerDesError(msg)
686 return BatchResult.from_dict(decoded)
689class _BatchResultSerDes(ExtendedTypeSerDes[Any]):
690 """Operation-owned serializer for BatchResult values."""
692 def __init__(self) -> None:
693 super().__init__(type_codecs=(_BatchResultCodec(),))
695 def _check_circular_references(
696 self,
697 obj: Any,
698 seen: set[int] | None = None,
699 ) -> None:
700 if not isinstance(obj, BatchResult):
701 super()._check_circular_references(obj, seen)
702 return
704 if seen is None:
705 seen = set()
706 obj_id = id(obj)
707 if obj_id in seen: 707 ↛ 708line 707 didn't jump to line 708 because the condition on line 707 was never true
708 msg = "Circular references are not supported"
709 raise SerDesError(msg)
711 seen.add(obj_id)
712 try:
713 super()._check_circular_references(_batch_result_payload(obj), seen)
714 finally:
715 seen.remove(obj_id)
718_BATCH_RESULT_SERDES = _BatchResultSerDes()
721@dataclass(frozen=True)
722class Executable(Generic[CallableType]):
723 """Index plus callable payload used by the concurrent executors."""
725 index: int
726 func: CallableType
729class BranchStatus(Enum):
730 """In-memory lifecycle state for a concurrently scheduled branch.
732 Values:
733 NOT_STARTED: The branch has not started and does not occupy a
734 concurrency slot.
735 PENDING: A previously suspended branch is being resubmitted. It has no
736 active task but continues to occupy its original concurrency slot.
737 RUNNING: The branch has an active asyncio task and occupies a
738 concurrency slot.
739 COMPLETED: The branch completed successfully. This is a terminal state
740 and releases its concurrency slot.
741 SUSPENDED: The branch is waiting indefinitely, such as for an external
742 callback. It has no active task but continues to occupy its slot.
743 SUSPENDED_WITH_TIMEOUT: The branch is waiting until a scheduled
744 timestamp, such as for a wait or retry. It has no active task but
745 continues to occupy its slot.
746 FAILED: The branch completed with an error. This is a terminal state
747 and releases its concurrency slot.
748 CANCELLED: The branch was cancelled after the parent reached an early
749 completion condition. This is a terminal state.
751 Typical state transitions::
753 NOT_STARTED -> RUNNING -> COMPLETED
754 -> FAILED
755 -> CANCELLED
756 -> SUSPENDED
757 -> SUSPENDED_WITH_TIMEOUT
758 SUSPENDED_WITH_TIMEOUT -> PENDING -> RUNNING
760 A timed suspension transitions through ``PENDING`` when its branch is
761 resubmitted in the same invocation. An indefinitely suspended branch waits
762 for a later durable invocation, which rebuilds this in-memory state before
763 replaying the branch.
764 """
766 NOT_STARTED = "not_started"
767 PENDING = "pending"
768 RUNNING = "running"
769 COMPLETED = "completed"
770 SUSPENDED = "suspended"
771 SUSPENDED_WITH_TIMEOUT = "suspended_with_timeout"
772 FAILED = "failed"
773 CANCELLED = "cancelled"
776class ExecutableWithState(Generic[CallableType, ResultType]):
777 """Manages the execution state and lifecycle of an executable."""
779 def __init__(self, executable: Executable[CallableType]) -> None:
780 self.executable = executable
781 self._status = BranchStatus.NOT_STARTED
782 self._future: asyncio.Task[ResultType] | None = None
783 self._suspend_until: float | None = None
784 self._result: ResultType | None = None
785 self._is_result_set = False
786 self._error: Exception | None = None
788 @property
789 def future(self) -> asyncio.Task[ResultType]:
790 if self._future is None:
791 msg = f"ExecutableWithState has no active task. {self.executable.index}"
792 raise InvalidStateError(msg)
793 return self._future
795 @property
796 def status(self) -> BranchStatus:
797 return self._status
799 @property
800 def result(self) -> ResultType:
801 if not self._is_result_set or self._status != BranchStatus.COMPLETED:
802 msg = f"result not available in status {self._status}"
803 raise InvalidStateError(msg)
804 return cast("ResultType", self._result)
806 @property
807 def error(self) -> Exception:
808 if self._error is None or self._status != BranchStatus.FAILED:
809 msg = f"error not available in status {self._status}"
810 raise InvalidStateError(msg)
811 return self._error
813 @property
814 def suspend_until(self) -> float | None:
815 return self._suspend_until
817 @property
818 def is_running(self) -> bool:
819 return self._status is BranchStatus.RUNNING
821 @property
822 def can_resume(self) -> bool:
823 return self._status is BranchStatus.SUSPENDED or (
824 self._status is BranchStatus.SUSPENDED_WITH_TIMEOUT
825 and self._suspend_until is not None
826 and time.time() >= self._suspend_until
827 )
829 @property
830 def index(self) -> int:
831 return self.executable.index
833 @property
834 def callable(self) -> CallableType:
835 return self.executable.func
837 def run(self, future: asyncio.Task[ResultType]) -> None:
838 if self._status not in {BranchStatus.NOT_STARTED, BranchStatus.PENDING}:
839 msg = f"Cannot start running from {self._status}"
840 raise InvalidStateError(msg)
841 self._status = BranchStatus.RUNNING
842 self._future = future
844 def suspend(self) -> None:
845 self._status = BranchStatus.SUSPENDED
846 self._suspend_until = None
848 def suspend_with_timeout(self, timestamp: float) -> None:
849 self._status = BranchStatus.SUSPENDED_WITH_TIMEOUT
850 self._suspend_until = timestamp
852 def complete(self, result: ResultType) -> None:
853 self._status = BranchStatus.COMPLETED
854 self._result = result
855 self._is_result_set = True
857 def fail(self, error: Exception) -> None:
858 self._status = BranchStatus.FAILED
859 self._error = error
861 def cancel(self) -> None:
862 self._status = BranchStatus.CANCELLED
864 def reset_to_pending(self) -> None:
865 self._status = BranchStatus.PENDING
866 self._future = None
867 self._suspend_until = None
870class ExecutionCounters:
871 """Counters for tracking execution state on a single event loop."""
873 def __init__(
874 self,
875 total_tasks: int,
876 completion_config: CompletionConfig,
877 ) -> None:
878 self.total_tasks = total_tasks
879 self.completion_config = completion_config
880 self.success_count = 0
881 self.failure_count = 0
883 def complete_task(self) -> None:
884 self.success_count += 1
886 def fail_task(self) -> None:
887 self.failure_count += 1
889 def should_continue(self) -> bool:
890 tolerated_failure_count = self.completion_config.tolerated_failure_count
892 if tolerated_failure_count is None:
893 return self.failure_count == 0
895 if (
896 tolerated_failure_count is not None
897 and self.failure_count > tolerated_failure_count
898 ):
899 return False
901 return True
903 def is_complete(self) -> bool:
904 completed_count = self.success_count + self.failure_count
906 if completed_count == self.total_tasks:
907 return True
909 min_successful = self.completion_config.min_successful
910 return min_successful is not None and self.success_count >= min_successful
912 def should_complete(self) -> bool:
913 return self.completion_decision().should_complete
915 def completion_status(self) -> CompletionStatus:
916 return CompletionStatus(
917 success_count=self.success_count,
918 failure_count=self.failure_count,
919 total_count=self.total_tasks,
920 )
922 def completion_decision(self) -> CompletionDecision:
923 if self.completion_config.has_custom_should_complete:
924 return self.completion_config.completion_decision(self.completion_status())
926 if self.is_complete() or not self.should_continue():
927 return CompletionDecision.complete(
928 BatchResult._get_completion_reason(
929 failure_count=self.failure_count,
930 success_count=self.success_count,
931 completed_count=self.success_count + self.failure_count,
932 total_count=self.total_tasks,
933 completion_config=self.completion_config,
934 )
935 )
937 return CompletionDecision.continue_execution()
939 def is_all_completed(self) -> bool:
940 return self.success_count == self.total_tasks
942 def is_min_successful_reached(self) -> bool:
943 min_successful = self.completion_config.min_successful
944 return min_successful is not None and self.success_count >= min_successful
946 def is_failure_tolerance_exceeded(self) -> bool:
947 return self._is_failure_condition_reached(
948 tolerated_count=self.completion_config.tolerated_failure_count,
949 failure_count=self.failure_count,
950 )
952 def _is_failure_condition_reached(
953 self,
954 tolerated_count: int | None,
955 failure_count: int,
956 ) -> bool:
957 if tolerated_count is not None and failure_count > tolerated_count:
958 return True
960 return False
963class TimerScheduler:
964 """Manage timed suspend tasks with event-loop tasks."""
966 def __init__(
967 self,
968 resubmit_callback: Callable[[ExecutableWithState[Any, Any]], Awaitable[None]],
969 ) -> None:
970 self.resubmit_callback = resubmit_callback
971 self._resume_tasks: set[asyncio.Task[None]] = set()
973 async def __aenter__(self) -> TimerScheduler:
974 return self
976 async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
977 await self.shutdown()
979 def schedule_resume(
980 self,
981 exe_state: ExecutableWithState[CallableType, ResultType],
982 resume_time: float,
983 ) -> None:
984 async def runner() -> None:
985 await asyncio.sleep(max(0.0, resume_time - time.time()))
986 if exe_state.can_resume:
987 exe_state.reset_to_pending()
988 await self.resubmit_callback(exe_state)
990 task = asyncio.create_task(runner())
991 self._resume_tasks.add(task)
992 task.add_done_callback(self._resume_tasks.discard)
994 async def shutdown(self) -> None:
995 tasks = list(self._resume_tasks)
996 for task in tasks:
997 task.cancel()
998 if tasks:
999 await asyncio.gather(*tasks, return_exceptions=True)
1000 self._resume_tasks.clear()
1003class ParallelExecutor(
1004 OperationExecutor[BatchResult[ResultType]],
1005 Generic[CallableType, ResultType],
1006):
1007 """Execute durable operations concurrently using asyncio tasks."""
1009 def __init__(
1010 self,
1011 execution_state: ExecutionState,
1012 operation_identifier: OperationIdentifier,
1013 executor_context: DurableContext,
1014 executables: list[Executable[CallableType]],
1015 max_concurrency: int | None,
1016 completion_config: CompletionConfig,
1017 top_level_sub_type: OperationSubType,
1018 iteration_sub_type: OperationSubType,
1019 name_prefix: str,
1020 serdes: SerDes | None,
1021 item_serdes: SerDes | None = None,
1022 summary_generator: SummaryGenerator | None = None,
1023 nesting_type: NestingType = NestingType.NESTED,
1024 branch_namer: Callable[[int], str] | None = None,
1025 branch_operations: Mapping[int, ExtensionOperation] | None = None,
1026 ) -> None:
1027 super().__init__(
1028 state=execution_state,
1029 operation_identifier=operation_identifier,
1030 )
1031 self.executor_context = executor_context
1032 self.executables = executables
1033 self.max_concurrency = max_concurrency
1034 self.completion_config = completion_config
1035 self.sub_type_top = top_level_sub_type
1036 self.sub_type_iteration = iteration_sub_type
1037 self.name_prefix = name_prefix
1038 self.summary_generator = summary_generator
1039 self.nesting_type = nesting_type
1040 self._branch_namer = branch_namer
1041 self._branch_operations = branch_operations
1042 self._started_branch_operations: set[int] = set()
1043 self._completion_event = asyncio.Event()
1044 self._suspend_exception: SuspendExecution | None = None
1045 self._completion_exception: Exception | None = None
1046 self._completion_decision: CompletionDecision | None = None
1047 self._running_tasks: set[asyncio.Task[ResultType]] = set()
1048 self._completion_tasks: set[asyncio.Task[None]] = set()
1050 self.counters = ExecutionCounters(
1051 len(executables),
1052 self.completion_config,
1053 )
1054 self.executables_with_state: list[ExecutableWithState] = []
1055 self.serdes = serdes
1056 self.item_serdes = item_serdes
1058 async def execute_item(
1059 self,
1060 child_context: DurableContext,
1061 executable: Executable[CallableType],
1062 ) -> ResultType:
1063 func = cast("Callable[[], Awaitable[ResultType]]", executable.func)
1064 with bind_current_context(child_context):
1065 result: ResultType = await func()
1066 return result
1068 def get_iteration_name(self, index: int) -> str:
1069 if self._branch_namer is not None:
1070 return self._branch_namer(index)
1071 return f"{self.name_prefix}{index}"
1073 async def start(self) -> BatchResult[ResultType]:
1074 return await self.execute()
1076 async def replay(self, operation: Operation) -> BatchResult[ResultType]:
1077 if operation.status is OperationStatus.SUCCEEDED:
1078 return await self.replay_completed(self.state, self.executor_context)
1079 return await self.execute()
1081 async def execute(self) -> BatchResult[ResultType]:
1082 logger.debug(
1083 "▶️ Executing concurrent operation, items: %d", len(self.executables)
1084 )
1086 if not self.executables:
1087 logger.debug("No items to execute, returning empty result")
1088 return self._create_result()
1090 max_workers = self.max_concurrency or len(self.executables)
1091 semaphore = asyncio.Semaphore(max_workers)
1092 self.executables_with_state = [
1093 ExecutableWithState(executable=exe) for exe in self.executables
1094 ]
1095 self._completion_event.clear()
1096 self._suspend_exception = None
1097 self._completion_exception = None
1098 self._completion_decision = None
1099 self._running_tasks.clear()
1100 self._completion_tasks.clear()
1101 next_executable_index = 0
1103 async def submit_task(
1104 executable_with_state: ExecutableWithState[CallableType, ResultType],
1105 ) -> None:
1106 if self._completion_event.is_set(): 1106 ↛ 1107line 1106 didn't jump to line 1107 because the condition on line 1106 was never true
1107 return
1109 async def run_task() -> ResultType:
1110 async with semaphore:
1111 return await self._execute_item_in_child_context(
1112 self.executor_context,
1113 executable_with_state.executable,
1114 )
1116 task = asyncio.create_task(run_task())
1117 executable_with_state.run(task)
1118 self._running_tasks.add(task)
1120 def on_done(done_task: asyncio.Task[ResultType]) -> None:
1121 self._running_tasks.discard(done_task)
1122 completion_task = asyncio.create_task(handle_task_completion(done_task))
1123 self._completion_tasks.add(completion_task)
1124 completion_task.add_done_callback(self._completion_tasks.discard)
1126 async def handle_task_completion(
1127 done_task: asyncio.Task[ResultType],
1128 ) -> None:
1129 await self._on_task_complete(
1130 executable_with_state,
1131 done_task,
1132 scheduler,
1133 )
1134 if (
1135 not self._completion_event.is_set()
1136 and executable_with_state.status
1137 in {BranchStatus.COMPLETED, BranchStatus.FAILED}
1138 ):
1139 await submit_next_task()
1140 if not self._completion_event.is_set():
1141 self._complete_if_execution_cannot_progress()
1143 task.add_done_callback(on_done)
1145 async def submit_next_task() -> None:
1146 nonlocal next_executable_index
1147 if self._completion_event.is_set() or next_executable_index >= len(
1148 self.executables_with_state
1149 ):
1150 return
1152 executable_with_state = self.executables_with_state[next_executable_index]
1153 next_executable_index += 1
1154 await submit_task(executable_with_state)
1156 async def resubmitter(
1157 executable_with_state: ExecutableWithState[CallableType, ResultType],
1158 ) -> None:
1159 await self.state.create_checkpoint(is_sync=False)
1160 await submit_task(executable_with_state)
1162 async with TimerScheduler(resubmitter) as scheduler:
1163 for _ in range(min(max_workers, len(self.executables_with_state))):
1164 await submit_next_task()
1166 await self._completion_event.wait()
1168 for task in list(self._running_tasks):
1169 if not task.done(): 1169 ↛ 1168line 1169 didn't jump to line 1168 because the condition on line 1169 was always true
1170 task.cancel()
1171 if self._running_tasks:
1172 await asyncio.gather(*self._running_tasks, return_exceptions=True)
1173 await asyncio.sleep(0)
1174 while self._completion_tasks: 1174 ↛ 1175line 1174 didn't jump to line 1175 because the condition on line 1174 was never true
1175 await asyncio.gather(
1176 *list(self._completion_tasks),
1177 return_exceptions=True,
1178 )
1179 if self._completion_decision is not None:
1180 self._cancel_unfinished_executables()
1182 if self._suspend_exception:
1183 raise self._suspend_exception
1184 if self._completion_exception: 1184 ↛ 1187line 1184 didn't jump to line 1187
1185 raise self._completion_exception
1187 return self._create_result()
1189 def should_execution_suspend(self) -> SuspendResult:
1190 earliest_timestamp: float = float("inf")
1191 indefinite_suspend_task: (
1192 ExecutableWithState[CallableType, ResultType] | None
1193 ) = None
1195 for exe_state in self.executables_with_state:
1196 if exe_state.status in {BranchStatus.PENDING, BranchStatus.RUNNING}:
1197 return SuspendResult.do_not_suspend()
1198 if exe_state.status is BranchStatus.NOT_STARTED:
1199 continue
1200 if exe_state.status is BranchStatus.SUSPENDED_WITH_TIMEOUT:
1201 if (
1202 exe_state.suspend_until
1203 and exe_state.suspend_until < earliest_timestamp
1204 ):
1205 earliest_timestamp = cast(float, exe_state.suspend_until)
1206 elif exe_state.status is BranchStatus.SUSPENDED:
1207 indefinite_suspend_task = exe_state
1209 if earliest_timestamp != float("inf"):
1210 return SuspendResult.suspend(
1211 TimedSuspendExecution(
1212 "All concurrent work complete or suspended pending retry.",
1213 earliest_timestamp,
1214 )
1215 )
1216 if indefinite_suspend_task: 1216 ↛ 1223line 1216 didn't jump to line 1223 because the condition on line 1216 was always true
1217 return SuspendResult.suspend(
1218 SuspendExecution(
1219 "All concurrent work complete or suspended and pending external callback."
1220 )
1221 )
1223 return SuspendResult.do_not_suspend()
1225 async def _on_task_complete(
1226 self,
1227 exe_state: ExecutableWithState[CallableType, ResultType],
1228 task: asyncio.Task[ResultType],
1229 scheduler: TimerScheduler,
1230 ) -> None:
1231 if task.cancelled():
1232 exe_state.cancel()
1233 return
1235 try:
1236 result = task.result()
1237 exe_state.complete(result)
1238 self.counters.complete_task()
1239 except OrphanedChildException:
1240 logger.debug(
1241 "Terminating orphaned branch %s without error because parent has completed already",
1242 exe_state.index,
1243 )
1244 return
1245 except TimedSuspendExecution as tse:
1246 exe_state.suspend_with_timeout(tse.scheduled_timestamp)
1247 scheduler.schedule_resume(exe_state, tse.scheduled_timestamp)
1248 except SuspendExecution:
1249 exe_state.suspend()
1250 except InvocationError as error:
1251 if error.is_retryable(): 1251 ↛ 1257line 1251 didn't jump to line 1257 because the condition on line 1251 was always true
1252 if self._completion_decision is not None:
1253 return
1254 self._completion_exception = error
1255 self._completion_event.set()
1256 return
1257 exe_state.fail(error)
1258 self.counters.fail_task()
1259 except Exception as e: # noqa: BLE001
1260 exe_state.fail(e)
1261 self.counters.fail_task()
1263 completion_decision = self.counters.completion_decision()
1264 if completion_decision.should_complete:
1265 self._completion_decision = completion_decision
1266 self._completion_event.set()
1268 def _complete_if_execution_cannot_progress(self) -> None:
1269 suspend_result = self.should_execution_suspend()
1270 if suspend_result.should_suspend:
1271 self._suspend_exception = suspend_result.exception
1272 self._completion_event.set()
1273 elif self._all_executables_terminal(): 1273 ↛ 1274line 1273 didn't jump to line 1274 because the condition on line 1273 was never true
1274 self._completion_exception = InvalidStateError(
1275 "custom should_complete did not complete after all branches "
1276 "reached terminal states"
1277 )
1278 self._completion_event.set()
1280 def _all_executables_terminal(self) -> bool:
1281 return all(
1282 exe_state.status
1283 in {
1284 BranchStatus.COMPLETED,
1285 BranchStatus.FAILED,
1286 BranchStatus.CANCELLED,
1287 }
1288 for exe_state in self.executables_with_state
1289 )
1291 def _cancel_unfinished_executables(self) -> None:
1292 for exe_state in self.executables_with_state:
1293 if exe_state.status in {
1294 BranchStatus.PENDING,
1295 BranchStatus.RUNNING,
1296 BranchStatus.SUSPENDED,
1297 BranchStatus.SUSPENDED_WITH_TIMEOUT,
1298 }:
1299 exe_state.cancel()
1301 def _create_result(self) -> BatchResult[ResultType]:
1302 batch_items: list[BatchItem[ResultType]] = []
1303 for executable in self.executables_with_state:
1304 match executable.status:
1305 case BranchStatus.COMPLETED:
1306 batch_items.append(
1307 BatchItem(
1308 executable.index,
1309 BatchItemStatus.SUCCEEDED,
1310 executable.result,
1311 )
1312 )
1313 case BranchStatus.FAILED:
1314 batch_items.append(
1315 BatchItem(
1316 executable.index,
1317 BatchItemStatus.FAILED,
1318 error=ErrorObject.from_exception(executable.error),
1319 )
1320 )
1321 case BranchStatus.CANCELLED:
1322 batch_items.append(
1323 BatchItem(executable.index, BatchItemStatus.CANCELLED)
1324 )
1325 case (
1326 BranchStatus.PENDING
1327 | BranchStatus.RUNNING
1328 | BranchStatus.SUSPENDED
1329 | BranchStatus.SUSPENDED_WITH_TIMEOUT
1330 ):
1331 batch_items.append(
1332 BatchItem(executable.index, BatchItemStatus.STARTED)
1333 )
1334 case BranchStatus.NOT_STARTED: 1334 ↛ 1303line 1334 didn't jump to line 1303 because the pattern on line 1334 always matched
1335 continue
1337 if (
1338 self._completion_decision is not None
1339 and self._completion_decision.completion_reason is not None
1340 ):
1341 return BatchResult(
1342 all=batch_items,
1343 completion_reason=self._completion_decision.completion_reason,
1344 )
1346 return BatchResult.from_items(batch_items, self.completion_config)
1348 async def _execute_item_in_child_context(
1349 self,
1350 executor_context: DurableContext,
1351 executable: Executable[CallableType],
1352 ) -> ResultType:
1353 is_virtual: bool = self.nesting_type is NestingType.FLAT
1355 async def run_child_operation() -> ResultType:
1356 return await self.execute_item(get_durable_context(), executable)
1358 if self._branch_operations is not None:
1359 operation = self._branch_operations[executable.index]
1360 if executable.index in self._started_branch_operations: 1360 ↛ 1361line 1360 didn't jump to line 1361 because the condition on line 1360 was never true
1361 task = operation._restart_child_context( # noqa: SLF001
1362 run_child_operation,
1363 serdes=self.item_serdes or self.serdes,
1364 summary_generator=self.summary_generator,
1365 is_virtual=is_virtual,
1366 )
1367 else:
1368 self._started_branch_operations.add(executable.index)
1369 task = operation._run_in_child_context( # noqa: SLF001
1370 run_child_operation,
1371 sub_type=self.sub_type_iteration,
1372 serdes=self.item_serdes or self.serdes,
1373 summary_generator=self.summary_generator,
1374 is_virtual=is_virtual,
1375 )
1376 return await task
1378 operation_id: str = (
1379 executor_context.step_counter._create_step_id_for_logical_step( # noqa: SLF001
1380 executable.index
1381 )
1382 )
1383 name: str = self.get_iteration_name(executable.index)
1385 child_context: DurableContext = executor_context.create_child_context(
1386 operation_id, is_virtual=is_virtual
1387 )
1388 operation_identifier = OperationIdentifier(
1389 operation_id=operation_id,
1390 sub_type=self.sub_type_iteration,
1391 parent_id=executor_context.parent_id,
1392 name=name,
1393 )
1395 async def run_legacy_child_operation() -> ResultType:
1396 return await self.execute_item(child_context, executable)
1398 executor: ChildOperationExecutor[ResultType] = ChildOperationExecutor(
1399 run_legacy_child_operation,
1400 child_context.execution_state,
1401 operation_identifier,
1402 serdes=self.item_serdes or self.serdes,
1403 summary_generator=self.summary_generator,
1404 is_virtual=is_virtual,
1405 )
1406 return await executor.process()
1408 async def replay_completed(
1409 self, execution_state: ExecutionState, executor_context: DurableContext
1410 ) -> BatchResult[ResultType]:
1411 items: list[BatchItem[ResultType]] = []
1412 for executable in self.executables:
1413 if isinstance(self._branch_operations, _BranchOperationReservations):
1414 operation_id = self._branch_operations.operation_id(executable.index)
1415 elif self._branch_operations is not None:
1416 operation_id = self._branch_operations[executable.index]._operation_id # noqa: SLF001
1417 else:
1418 operation_id = (
1419 executor_context.step_counter._create_step_id_for_logical_step( # noqa: SLF001
1420 executable.index
1421 )
1422 )
1423 operation = execution_state.operations.get(operation_id)
1425 result: ResultType | None = None
1426 error = None
1427 status: BatchItemStatus
1428 if operation is not None and operation.status is OperationStatus.SUCCEEDED:
1429 status = BatchItemStatus.SUCCEEDED
1430 operation_details = operation.context_details
1431 if operation_details is not None and operation_details.replay_children:
1432 result = await self._execute_item_in_child_context(
1433 executor_context, executable
1434 )
1435 elif (
1436 operation_details is not None
1437 and operation_details.result is not None
1438 ):
1439 result = await deserialize(
1440 serdes=self.item_serdes or self.serdes,
1441 data=operation_details.result,
1442 operation_id=operation_id,
1443 durable_execution_arn=execution_state.durable_execution_arn,
1444 recursive_level=execution_state.recursive_level,
1445 operation_name=operation.name,
1446 parent_id=operation.parent_id,
1447 operation_type=operation.operation_type,
1448 operation_sub_type=operation.sub_type,
1449 )
1450 elif operation is not None and operation.status is OperationStatus.FAILED:
1451 error = (
1452 operation.context_details.error
1453 if operation.context_details is not None
1454 else None
1455 )
1456 status = BatchItemStatus.FAILED
1457 elif operation is not None and operation.status in {
1458 OperationStatus.CANCELLED,
1459 OperationStatus.STARTED,
1460 }:
1461 status = BatchItemStatus.CANCELLED
1462 else:
1463 continue
1465 items.append(
1466 BatchItem(executable.index, status, result=result, error=error)
1467 )
1468 return BatchResult.from_items(items, self.completion_config)
1471class ParallelSummaryGenerator:
1472 """Default summary generator for oversized parallel `BatchResult` payloads."""
1474 def __call__(self, result: BatchResult) -> str:
1475 fields = {
1476 "totalCount": result.total_count,
1477 "successCount": result.success_count,
1478 "failureCount": result.failure_count,
1479 "completionReason": result.completion_reason.value,
1480 "status": result.status.value,
1481 "startedCount": result.started_count,
1482 "type": "ParallelResult",
1483 }
1485 return json.dumps(fields)
1488class _BranchOperationReservations(Mapping[int, ExtensionOperation]):
1489 """Create branch reservations lazily while retaining replay checkpoints."""
1491 def __init__(
1492 self,
1493 *,
1494 context: DurableContext,
1495 count: int,
1496 sub_type: OperationSubType,
1497 name_prefix: str,
1498 branch_namer: Callable[[int], str] | None,
1499 ) -> None:
1500 self._context = context
1501 self._count = count
1502 self._sub_type = sub_type
1503 self._name_prefix = name_prefix
1504 self._branch_namer = branch_namer
1505 self._extension = ExtensionContext(context)
1506 self._parent_replaying = context.is_replaying()
1507 self._reservations: dict[int, ExtensionOperation] = {}
1508 self._register_historical_checkpoints()
1510 def __getitem__(self, index: int) -> ExtensionOperation:
1511 operation_id = self.operation_id(index)
1512 reservation = self._reservations.get(index)
1513 if reservation is None:
1514 reservation = self._extension._reserve_sdk_operation_id( # noqa: SLF001
1515 self._branch_name(index),
1516 operation_id=operation_id,
1517 parent_replaying=self._parent_replaying,
1518 )
1519 self._reservations[index] = reservation
1520 return reservation
1522 def __iter__(self) -> Iterator[int]:
1523 return iter(range(self._count))
1525 def __len__(self) -> int:
1526 return self._count
1528 def operation_id(self, index: int) -> str:
1529 """Return a branch ID without creating its reservation or name."""
1530 if index < 0 or index >= self._count: 1530 ↛ 1531line 1530 didn't jump to line 1531 because the condition on line 1530 was never true
1531 raise KeyError(index)
1532 return self._context.step_counter._create_step_id_for_logical_step( # noqa: SLF001
1533 index
1534 )
1536 def _branch_name(self, index: int) -> str:
1537 if self._branch_namer is not None:
1538 return self._branch_namer(index)
1539 return f"{self._name_prefix}{index}"
1541 def _register_historical_checkpoints(self) -> None:
1542 operations = self._context.execution_state.operations
1543 if not isinstance(operations, Mapping):
1544 return
1545 for operation in operations.values():
1546 if (
1547 operation.operation_type is OperationType.CONTEXT
1548 and operation.sub_type == self._sub_type
1549 and operation.parent_id == self._context.parent_id
1550 ):
1551 self._context.step_counter._register_reservation( # noqa: SLF001
1552 operation.operation_id,
1553 has_checkpoint=True,
1554 )
1557@durable_callable
1558async def parallel_handler(
1559 callables: Sequence[Callable[[], Awaitable[R]]],
1560 execution_state: ExecutionState,
1561 parallel_context: DurableContext,
1562 operation_identifier: OperationIdentifier,
1563 *,
1564 max_concurrency: int | None = None,
1565 completion_config: CompletionConfig | None = None,
1566 serdes: SerDes | None = None,
1567 item_serdes: SerDes | None = None,
1568 summary_generator: SummaryGenerator | None = ParallelSummaryGenerator(),
1569 nesting_type: NestingType = NestingType.NESTED,
1570 top_level_sub_type: OperationSubType = OperationSubType.PARALLEL,
1571 iteration_sub_type: OperationSubType = OperationSubType.PARALLEL_BRANCH,
1572 name_prefix: str = "parallel-branch-",
1573 branch_namer: Callable[[int], str] | None = None,
1574) -> BatchResult[R]:
1575 """Execute multiple operations in parallel."""
1576 # Summary Generator Construction (matches TypeScript implementation):
1577 # Construct the summary generator at the handler level, just like TypeScript does in parallel-handler.ts.
1578 # This matches the pattern where handlers are responsible for configuring operation-specific behavior.
1579 #
1580 # See TypeScript reference: aws-durable-execution-sdk-js/src/handlers/parallel-handler/parallel-handler.ts (~line 112)
1582 branch_operations: Mapping[int, ExtensionOperation] | None = None
1583 if isinstance(parallel_context, DurableContext):
1584 branch_operations = _BranchOperationReservations(
1585 context=parallel_context,
1586 count=len(callables),
1587 sub_type=iteration_sub_type,
1588 name_prefix=name_prefix,
1589 branch_namer=branch_namer,
1590 )
1591 executor_kwargs: dict[str, Any] = {
1592 "executables": [
1593 Executable(index=i, func=func) for i, func in enumerate(callables)
1594 ],
1595 "max_concurrency": max_concurrency,
1596 "completion_config": completion_config or CompletionConfig.all_successful(),
1597 "top_level_sub_type": top_level_sub_type,
1598 "iteration_sub_type": iteration_sub_type,
1599 "name_prefix": name_prefix,
1600 "serdes": serdes,
1601 "summary_generator": summary_generator,
1602 "item_serdes": item_serdes,
1603 "nesting_type": nesting_type,
1604 "branch_namer": branch_namer,
1605 "execution_state": execution_state,
1606 "operation_identifier": operation_identifier,
1607 "executor_context": parallel_context,
1608 }
1609 if branch_operations is not None:
1610 executor_kwargs["branch_operations"] = branch_operations
1612 executor: ParallelExecutor[Callable[[], Awaitable[R]], R] = ParallelExecutor(
1613 **executor_kwargs,
1614 )
1616 return await executor.process()
1619def parallel(
1620 branches: Iterable[Callable[[], Awaitable[T]]],
1621 *,
1622 name: str | None = None,
1623 max_concurrency: int | None = None,
1624 completion_config: CompletionConfig | None = None,
1625 serdes: SerDes | None = None,
1626 item_serdes: SerDes | None = None,
1627 summary_generator: SummaryGenerator | None = ParallelSummaryGenerator(),
1628 nesting_type: NestingType = NestingType.NESTED,
1629) -> asyncio.Task[BatchResult[T]]:
1630 """Start a durable parallel operation.
1632 Each branch is an async zero-argument callable, typically a bound durable
1633 callable such as `fetch_user(user_id)`. Branches run in child durable
1634 contexts and may contain durable operations such as `step()` or `wait()`.
1636 The returned object is an `asyncio.Task`; awaiting it yields a `BatchResult`.
1637 Calling `parallel()` without immediately awaiting it schedules the durable
1638 operation in the background, consistent with other operation helpers.
1640 By default, `parallel()` uses `CompletionConfig.all_successful()`: every
1641 branch must succeed, and the first failure completes the operation as failed.
1642 Pass `completion_config` to use threshold-based or custom completion.
1644 Args:
1645 branches: Async zero-argument branch callables to run concurrently.
1646 name: Optional durable operation name.
1647 max_concurrency: Optional limit for in-flight branches. A suspended
1648 branch retains its slot until it reaches a terminal state.
1649 completion_config: Optional completion policy. Use
1650 `CompletionConfig.thresholds()`, `first_successful()`,
1651 `all_completed()`, `all_successful()`, or `custom()`.
1652 serdes: Optional serializer for the final `BatchResult`.
1653 item_serdes: Optional serializer for each branch result.
1654 summary_generator: Optional callable used to summarize oversized
1655 checkpoint payloads.
1656 nesting_type: Whether branch operations use nested or flat operation
1657 identifiers.
1659 Returns:
1660 An `asyncio.Task` that resolves to a `BatchResult` containing one
1661 `BatchItem` per branch.
1663 Raises:
1664 ValidationError: If `max_concurrency` is not a positive integer or
1665 `None`.
1666 RuntimeError: If called outside a durable context.
1667 """
1668 _validate_max_concurrency(max_concurrency)
1669 context = get_durable_context()
1670 validated_branches: list[Callable[[], Awaitable[T]]] = []
1671 for branch in branches:
1672 validated_branches.append(branch)
1674 async def run_parallel_handler() -> BatchResult[T]:
1675 parallel_context = get_durable_context()
1676 operation_id = parallel_context.step_id_prefix
1677 if operation_id is None:
1678 msg = "parallel operation id is not available in the current context"
1679 raise RuntimeError(msg)
1680 operation_identifier = OperationIdentifier(
1681 operation_id=operation_id,
1682 sub_type=OperationSubType.PARALLEL,
1683 parent_id=parallel_context.parent_id,
1684 name=name,
1685 )
1687 handler = parallel_handler(
1688 callables=validated_branches,
1689 execution_state=context.execution_state,
1690 parallel_context=parallel_context,
1691 operation_identifier=operation_identifier,
1692 max_concurrency=max_concurrency,
1693 completion_config=completion_config or CompletionConfig.all_successful(),
1694 serdes=serdes,
1695 item_serdes=item_serdes,
1696 summary_generator=summary_generator,
1697 nesting_type=nesting_type,
1698 )
1699 return await handler()
1701 return _run_in_child_context(
1702 run_parallel_handler,
1703 sub_type=OperationSubType.PARALLEL,
1704 name=name,
1705 serdes=serdes if serdes is not None else _BATCH_RESULT_SERDES,
1706 )