diff --git a/packages/aws-durable-execution-sdk-python-examples/test/map/test_map_completion.py b/packages/aws-durable-execution-sdk-python-examples/test/map/test_map_completion.py index f7bd850a..ddafc0a2 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/map/test_map_completion.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/map/test_map_completion.py @@ -21,10 +21,13 @@ def test_reproduce_completion_config_behavior_with_detailed_logging(durable_runn result_data = deserialize_operation_payload(result.result) - # 5 items are processed 2 of them succeeded. We exit early because min_successful is 2. - # Additionally, failure_count shows 0 because failed items have retry strategies configured and are still retrying + # max_concurrency=3 starts items 0-2. The two failing items hold their + # concurrency slots while retrying, so only the first success frees a + # slot for item 3. min_successful=2 is reached before item 4 can start, + # so 4 items appear in the result and the fifth never starts. + # failure_count shows 0 because failed items have retry strategies configured and are still retrying # when execution completes. Failures aren't finalized until retries complete, so they don't appear in the failure_count. - assert result_data["totalItems"] == 5 + assert result_data["totalItems"] == 4 assert result_data["successfulCount"] == 2 assert result_data["failedCount"] == 0 assert result_data["hasFailures"] is False diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py index 3634e0cb..17644353 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py @@ -4,27 +4,31 @@ import heapq import logging -import threading +import queue import time from abc import ABC, abstractmethod -from concurrent.futures import Future, ThreadPoolExecutor -from typing import TYPE_CHECKING, Generic, Self, TypeVar +from collections import deque +from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING, Generic, TypeVar from aws_durable_execution_sdk_python.concurrency.models import ( BatchItem, BatchItemStatus, BatchResult, + Branch, + BranchEvent, + BranchEventKind, BranchStatus, + CompletionPolicy, + CompletionRecord, Executable, - ExecutableWithState, - ExecutionCounters, - SuspendResult, ) from aws_durable_execution_sdk_python.config import ( ChildConfig, NestingType, ) from aws_durable_execution_sdk_python.exceptions import ( + InvalidStateError, OrphanedChildException, SuspendExecution, TimedSuspendExecution, @@ -35,13 +39,14 @@ if TYPE_CHECKING: - from collections.abc import Callable - from aws_durable_execution_sdk_python.config import CompletionConfig from aws_durable_execution_sdk_python.context import DurableContext from aws_durable_execution_sdk_python.lambda_service import OperationSubType from aws_durable_execution_sdk_python.serdes import SerDes - from aws_durable_execution_sdk_python.state import ExecutionState + from aws_durable_execution_sdk_python.state import ( + CheckpointedResult, + ExecutionState, + ) from aws_durable_execution_sdk_python.types import SummaryGenerator @@ -54,99 +59,26 @@ ResultType = TypeVar("ResultType") -# region concurrency logic -class TimerScheduler: - """Manage timed suspend tasks with a background timer thread.""" - - def __init__( - self, resubmit_callback: Callable[[list[ExecutableWithState]], None] - ) -> None: - self.resubmit_callback = resubmit_callback - self._pending_resumes: list[tuple[float, int, ExecutableWithState]] = [] - self._lock = threading.Lock() - self._schedule_counter = 0 - self._shutdown = threading.Event() - self._timer_thread = threading.Thread(target=self._timer_loop, daemon=True) - self._timer_thread.start() - - def __enter__(self) -> Self: - return self - - def __exit__(self, exc_type, exc_val, exc_tb) -> None: - self.shutdown() - - def schedule_resume( - self, exe_state: ExecutableWithState, resume_time: float - ) -> None: - """Schedule a task to resume at the specified time. - - Uses a counter as a tie-breaker to ensure FIFO ordering when multiple - tasks have the same resume_time, preventing TypeError from comparing - ExecutableWithState objects. - """ - with self._lock: - heapq.heappush( - self._pending_resumes, - (resume_time, self._schedule_counter, exe_state), - ) - self._schedule_counter += 1 - - def shutdown(self) -> None: - """Shutdown the timer thread and cancel all pending resumes.""" - self._shutdown.set() - self._timer_thread.join(timeout=1.0) - with self._lock: - self._pending_resumes.clear() - - def _timer_loop(self) -> None: - """Background thread that processes timed resumes.""" - while not self._shutdown.is_set(): - next_resume_time = None - - with self._lock: - if self._pending_resumes: - next_resume_time = self._pending_resumes[0][0] - - if next_resume_time is None: - # No pending resumes, wait a bit and check again - self._shutdown.wait(timeout=0.1) - continue - - current_time = time.time() - if current_time >= next_resume_time: - # Drain every due resume under the lock, transitioning each to - # PENDING atomically with the pop. Keeping pop+reset_to_pending - # together is required: should_execution_suspend reads branch - # status without this lock, so an item that is removed from the - # heap but still SUSPENDED_WITH_TIMEOUT could trigger a spurious - # parent suspend. - ready: list[ExecutableWithState] = [] - with self._lock: - while ( - self._pending_resumes - and self._pending_resumes[0][0] <= current_time - ): - _, _, exe_state = heapq.heappop(self._pending_resumes) - if exe_state.can_resume: - exe_state.reset_to_pending() - ready.append(exe_state) - # Resubmit outside the lock. Only the heap pop and the PENDING - # transition need the lock. The checkpoint refresh is a blocking - # network call and the submit hands work to the pool, so running - # them off the lock keeps timed resumes from serializing behind - # the network round trip and keeps the timer thread from - # re-entering this non-reentrant lock when a submitted future - # completes inline and its done-callback calls schedule_resume. - if ready: - self.resubmit_callback(ready) - else: - # Wait until next resume time - wait_time = min(next_resume_time - current_time, 0.1) - self._shutdown.wait(timeout=wait_time) - - class ConcurrentExecutor(ABC, Generic[CallableType, ResultType]): - """Execute durable operations concurrently. This contains the execution logic for Map and Parallel.""" + """Execute durable operations concurrently. This contains the execution logic for Map and Parallel. + + Scheduling model: a single coordinator loop runs on the calling thread + and owns all branch state. Worker threads run branches and report each + outcome as a :class:`BranchEvent` on a queue; they never mutate shared + state, so the module needs no locks. + + ``max_concurrency`` bounds in-flight branches, not threads. A branch + that suspends (e.g. awaiting an invoke result or callback) keeps its + concurrency slot until it reaches a terminal state. New branches start + only when a slot frees up. When every in-flight branch is suspended and + no slot is available, the parent suspends too: with the earliest resume + timestamp when one exists, indefinitely otherwise. + + Branches are always started in index order and operation ids derive + from the branch index, so scheduling is deterministic across + invocations. On re-invocation the previously started branches are + admitted first and replay from their checkpoints. + """ def __init__( self, @@ -178,29 +110,14 @@ def __init__( self.name_prefix = name_prefix self.summary_generator = summary_generator self.nesting_type = nesting_type - - # Event-driven state tracking for when the executor is done - self._completion_event = threading.Event() - self._suspend_exception: SuspendExecution | None = None - self._resume_error: Exception | None = None - - # ExecutionCounters will keep track of completion criteria and on-going counters - min_successful = self.completion_config.min_successful or len(self.executables) - tolerated_failure_count = self.completion_config.tolerated_failure_count - tolerated_failure_percentage = ( - self.completion_config.tolerated_failure_percentage - ) - - self.counters: ExecutionCounters = ExecutionCounters( - len(executables), - min_successful, - tolerated_failure_count, - tolerated_failure_percentage, - ) - self.executables_with_state: list[ExecutableWithState] = [] self.serdes = serdes self.item_serdes = item_serdes + self.policy: CompletionPolicy = CompletionPolicy.from_config( + len(executables), completion_config + ) + self.branches: list[Branch[CallableType, ResultType]] = [] + @abstractmethod def execute_item( self, child_context: DurableContext, executable: Executable[CallableType] @@ -219,223 +136,233 @@ def get_iteration_name(self, index: int) -> str: def execute( self, execution_state: ExecutionState, executor_context: DurableContext ) -> BatchResult[ResultType]: - """Execute items concurrently with event-driven state management.""" + """Run the coordinator loop until the batch completes or suspends.""" logger.debug( "▶️ Executing concurrent operation, items: %d", len(self.executables) ) - # Early return for empty executables if not self.executables: logger.debug("No items to execute, returning empty result") return self._create_result() - max_workers = self.max_concurrency or len(self.executables) - - self.executables_with_state = [ - ExecutableWithState(executable=exe) for exe in self.executables - ] - self._completion_event.clear() - self._suspend_exception = None - self._resume_error = None - - def resubmitter(ready: list[ExecutableWithState]) -> None: - """Resubmit a wave of timed-suspended tasks. - - One checkpoint refresh serves the whole due wave: the fetch returns - all operations, so every resumed branch reads fresh state. The - refresh only raises when the background checkpoint subsystem has - failed, which is terminal for the whole execution, so record the - error and wake the parent to re-raise it. Catching here keeps the - single timer thread alive so a failure does not strand the other - pending resumes. - """ - try: - execution_state.create_checkpoint() - except Exception as exc: # noqa: BLE001 - # resubmitter runs only on the single timer thread, so this - # check-then-set needs no lock. First error wins: keep the - # earliest failure if several waves fail before execute() reads - # it (they are the same terminal checkpoint failure anyway). - if self._resume_error is None: # pragma: no branch - self._resume_error = exc - self._completion_event.set() - return - for executable_with_state in ready: - submit_task(executable_with_state) - - thread_executor = ThreadPoolExecutor(max_workers=max_workers) - try: - with TimerScheduler(resubmitter) as scheduler: - - def submit_task(executable_with_state: ExecutableWithState) -> Future: - """Submit task to the thread executor and mark its state as started.""" - future = thread_executor.submit( - self._execute_item_in_child_context, - executor_context, - executable_with_state.executable, - ) - executable_with_state.run(future) + max_in_flight: int = self.max_concurrency or len(self.executables) + self.branches = [Branch(executable=exe) for exe in self.executables] - def on_done(future: Future) -> None: - self._on_task_complete(executable_with_state, future, scheduler) + events: queue.Queue[BranchEvent[ResultType]] = queue.Queue() + pending: deque[Branch[CallableType, ResultType]] = deque(self.branches) + timed_resumes: list[tuple[float, int]] = [] + branch_by_index: dict[int, Branch[CallableType, ResultType]] = { + branch.index: branch for branch in self.branches + } - future.add_done_callback(on_done) - return future + in_flight: int = 0 + running: int = 0 + succeeded: int = 0 + failed: int = 0 - # Submit initial tasks - futures = [ - submit_task(exe_state) for exe_state in self.executables_with_state - ] + pool: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=max_in_flight) - # Wait for completion - self._completion_event.wait() - - # Cancel futures that haven't started yet - for future in futures: - future.cancel() - - # A timed resume failed to refresh state (terminal checkpoint - # subsystem failure). Re-raise so the invocation fails and the - # backend retries from the last durable checkpoint. - if self._resume_error is not None: - raise self._resume_error + def submit(branch: Branch[CallableType, ResultType]) -> None: + branch.start() + pool.submit( + self._branch_worker, executor_context, events, branch.executable + ) - # Suspend execution if everything done and at least one of the tasks raised a suspend exception. - if self._suspend_exception: - raise self._suspend_exception + try: + while True: + if self.policy.is_complete( + succeeded, failed + ) or not self.policy.should_continue(failed): + break + + # Start branches in index order up to the in-flight limit. + # Suspended branches keep their slot: in_flight only + # decreases on terminal events. + while ( + pending + and in_flight < max_in_flight + and self.policy.should_continue(failed) + ): + submit(pending.popleft()) + in_flight += 1 + running += 1 + + # Resume due timed suspends in-process. One checkpoint + # refresh serves the whole due wave; a failure is terminal + # for the execution and propagates from this thread. + now: float = time.time() + due: list[Branch[CallableType, ResultType]] = [] + while timed_resumes and timed_resumes[0][0] <= now: + _, index = heapq.heappop(timed_resumes) + due.append(branch_by_index[index]) + if due: + execution_state.create_checkpoint() + for branch in due: + submit(branch) + running += 1 + continue + + if running == 0: + # Every in-flight branch is suspended and no slot is + # free (or no work remains): suspend the parent. + if timed_resumes: + raise TimedSuspendExecution( + "All concurrent work complete or suspended pending retry.", + timed_resumes[0][0], + ) + raise SuspendExecution( + "All concurrent work complete or suspended and pending external callback." + ) + timeout: float | None = None + if timed_resumes: + timeout = max(timed_resumes[0][0] - time.time(), 0) + try: + event: BranchEvent[ResultType] = events.get(timeout=timeout) + except queue.Empty: + # A timed resume came due while branches were running. + continue + + applied: Branch[CallableType, ResultType] = branch_by_index[event.index] + match event.kind: + case BranchEventKind.COMPLETED: + applied.complete(event.result) + succeeded += 1 + running -= 1 + in_flight -= 1 + case BranchEventKind.FAILED if event.error is not None: + applied.fail(event.error) + failed += 1 + running -= 1 + in_flight -= 1 + case BranchEventKind.SUSPENDED: + applied.suspend() + running -= 1 + case BranchEventKind.SUSPENDED_UNTIL if event.resume_at is not None: + applied.suspend_until(event.resume_at) + heapq.heappush(timed_resumes, (event.resume_at, event.index)) + running -= 1 + case BranchEventKind.ORPHANED: + # An ancestor context already checkpointed terminal, so + # every further checkpoint under it is rejected. Stop + # scheduling: the result of this batch is discarded + # upstream by the same orphan mechanism. + break + case _: + # A dropped event would leave the counters stale and + # hang the coordinator, so fail loudly instead. + msg = f"Unhandled branch event: {event}" + raise InvalidStateError(msg) finally: - # Shutdown without waiting for running threads for early return when - # completion criteria are met (e.g., min_successful). - # Running threads will continue in background but they raise OrphanedChildException - # on the next attempt to checkpoint. - thread_executor.shutdown(wait=False, cancel_futures=True) + # Shutdown without waiting for running threads for early return + # when completion criteria are met (e.g., min_successful). + # Running threads continue in the background and raise + # OrphanedChildException on their next attempt to checkpoint. + pool.shutdown(wait=False, cancel_futures=True) + + # Apply terminal events that raced the completion decision, so a + # branch that finished just before the batch completed is reported + # with its true status instead of STARTED. Best effort: events from + # still-running branches that arrive later are not waited for. + while True: + try: + raced: BranchEvent[ResultType] = events.get_nowait() + except queue.Empty: + break + raced_branch: Branch[CallableType, ResultType] = branch_by_index[ + raced.index + ] + if raced.kind is BranchEventKind.COMPLETED: + raced_branch.complete(raced.result) + elif raced.kind is BranchEventKind.FAILED and raced.error is not None: + raced_branch.fail(raced.error) - # Build final result return self._create_result() - def should_execution_suspend(self) -> SuspendResult: - """Check if execution should suspend.""" - earliest_timestamp: float = float("inf") - indefinite_suspend_task: ( - ExecutableWithState[CallableType, ResultType] | None - ) = None - - for exe_state in self.executables_with_state: - if exe_state.status in {BranchStatus.PENDING, BranchStatus.RUNNING}: - # Exit here! Still have tasks that can make progress, don't suspend. - return SuspendResult.do_not_suspend() - if exe_state.status is BranchStatus.SUSPENDED_WITH_TIMEOUT: - if ( - exe_state.suspend_until - and exe_state.suspend_until < earliest_timestamp - ): - earliest_timestamp = exe_state.suspend_until - elif exe_state.status is BranchStatus.SUSPENDED: - indefinite_suspend_task = exe_state - - # All tasks are in final states and at least one of them is a suspend. - if earliest_timestamp != float("inf"): - return SuspendResult.suspend( - TimedSuspendExecution( - "All concurrent work complete or suspended pending retry.", - earliest_timestamp, - ) - ) - if indefinite_suspend_task: - return SuspendResult.suspend( - SuspendExecution( - "All concurrent work complete or suspended and pending external callback." - ) - ) - - return SuspendResult.do_not_suspend() - - def _on_task_complete( - self, - exe_state: ExecutableWithState, - future: Future, - scheduler: TimerScheduler, - ) -> None: - """Handle task completion, suspension, or failure.""" - - if future.cancelled(): - exe_state.suspend() - return - - try: - result = future.result() - exe_state.complete(result) - self.counters.complete_task() - except OrphanedChildException: - # Parent already completed and returned. - # State is already RUNNING, which _create_result() marked as STARTED - # Just log and exit - no state change needed - logger.debug( - "Terminating orphaned branch %s without error because parent has completed already", - exe_state.index, - ) - return - except TimedSuspendExecution as tse: - exe_state.suspend_with_timeout(tse.scheduled_timestamp) - scheduler.schedule_resume(exe_state, tse.scheduled_timestamp) - except SuspendExecution: - exe_state.suspend() - # For indefinite suspend, don't schedule resume - except Exception as e: # noqa: BLE001 - exe_state.fail(e) - self.counters.fail_task() - - # Check if execution should complete or suspend - if self.counters.should_complete(): - self._completion_event.set() - else: - suspend_result = self.should_execution_suspend() - if suspend_result.should_suspend: - self._suspend_exception = suspend_result.exception - self._completion_event.set() - def _create_result(self) -> BatchResult[ResultType]: - """ - Build the final BatchResult. + """Build the final BatchResult from branch states. - When this function executes, we've terminated the upper/parent context for whatever reason. - It follows that our items can be only in 3 states, Completed, Failed and Started (in all of the possible forms). - We tag each branch based on its observed value at the time of completion of the parent / upper context, and pass the - results to BatchResult. + Branches map to batch items by status: COMPLETED and FAILED map to + their terminal statuses, anything started but not terminal maps to + STARTED. Never-started branches (still PENDING) are omitted, + matching the TypeScript implementation. - Any inference wrt completion reason is left up to BatchResult, keeping the logic inference isolated. + The completion reason is computed against the true batch total, not + just the started branches. """ + succeeded: int = 0 + failed: int = 0 batch_items: list[BatchItem[ResultType]] = [] - for executable in self.executables_with_state: - match executable.status: + for branch in self.branches: + match branch.status: case BranchStatus.COMPLETED: + succeeded += 1 batch_items.append( BatchItem( - executable.index, + branch.index, BatchItemStatus.SUCCEEDED, - executable.result, + branch.result, ) ) - case BranchStatus.FAILED: + case BranchStatus.FAILED if branch.error is not None: + failed += 1 batch_items.append( BatchItem( - executable.index, + branch.index, BatchItemStatus.FAILED, - error=ErrorObject.from_exception(executable.error), + error=ErrorObject.from_exception(branch.error), ) ) case ( - BranchStatus.PENDING - | BranchStatus.RUNNING + BranchStatus.RUNNING | BranchStatus.SUSPENDED | BranchStatus.SUSPENDED_WITH_TIMEOUT ): - batch_items.append( - BatchItem(executable.index, BatchItemStatus.STARTED) - ) + batch_items.append(BatchItem(branch.index, BatchItemStatus.STARTED)) + case BranchStatus.PENDING: + pass + case _: + # A silently skipped branch would shrink a + # customer-visible result, so fail loudly instead. + msg = f"Branch {branch.index} in unexpected state {branch.status}" + raise InvalidStateError(msg) + + return BatchResult(batch_items, self.policy.reason(succeeded, failed)) + + def _branch_worker( + self, + executor_context: DurableContext, + events: queue.Queue[BranchEvent[ResultType]], + executable: Executable[CallableType], + ) -> None: + """Worker-thread body: run one branch and report its outcome. - return BatchResult.from_items(batch_items, self.completion_config) + Converts every outcome into a :class:`BranchEvent` on the queue and + never raises into the pool. The coordinator loop is the sole + consumer of the events. + """ + try: + result: ResultType = self._execute_item_in_child_context( + executor_context, executable + ) + except TimedSuspendExecution as tse: + events.put( + BranchEvent.suspended_until(executable.index, tse.scheduled_timestamp) + ) + except SuspendExecution: + events.put(BranchEvent.suspended(executable.index)) + except OrphanedChildException: + # Parent already completed and returned; the branch stays + # RUNNING and is reported as STARTED. + logger.debug( + "Terminating orphaned branch %s without error because parent has completed already", + executable.index, + ) + events.put(BranchEvent.orphaned(executable.index)) + except Exception as e: # noqa: BLE001 + events.put(BranchEvent.failed(executable.index, e)) + else: + events.put(BranchEvent.completed(executable.index, result)) def _execute_item_in_child_context( self, @@ -512,13 +439,96 @@ def run_in_child_handler() -> ResultType: ) return result - def replay(self, execution_state: ExecutionState, executor_context: DurableContext): + def replay( + self, + execution_state: ExecutionState, + executor_context: DurableContext, + checkpointed_result: CheckpointedResult | None = None, + ) -> BatchResult[ResultType]: + """Reconstruct the batch result while in replay_children mode. + + When the operation's summary carries a recorded completion decision, + the reconstruction obeys it so the result matches the live result + exactly: branches recorded STARTED are reported STARTED without + consulting child checkpoints, branches past the started prefix are + omitted, terminal branches are re-derived, and the completion + reason is the recorded value verbatim. + + Summaries without a record (custom summary generators, checkpoints + written before the record existed) fall back to deriving every + branch from its child checkpoint. """ - Replay rather than re-run children. + record: CompletionRecord | None = CompletionRecord.from_summary_payload( + checkpointed_result.result if checkpointed_result else None + ) + if record is None: + return self._replay_from_checkpoints(execution_state, executor_context) - if we are here, then we are in replay_children. - This will pre-generate all the operation ids for the children and collect the checkpointed - results. + items: list[BatchItem[ResultType]] = [] + for executable in self.executables: + if executable.index >= record.started_total: + continue + if executable.index in record.started_indexes: + items.append(BatchItem(executable.index, BatchItemStatus.STARTED)) + continue + items.append( + self._replay_terminal_item( + execution_state, executor_context, executable + ) + ) + return BatchResult(items, record.completion_reason) + + def _replay_terminal_item( + self, + execution_state: ExecutionState, + executor_context: DurableContext, + executable: Executable[CallableType], + ) -> BatchItem[ResultType]: + """Re-derive one branch recorded as terminal. + + Non-virtual branches have a terminal checkpoint: terminal branch + events are only emitted after the synchronous SUCCEED/FAIL + checkpoint call returned. Virtual (FLAT) branches never checkpoint + themselves, so re-executing the branch body over its inner + operations' checkpoints discriminates success from failure. + """ + operation_id: str = executor_context._create_step_id_for_logical_step( # noqa: SLF001 + executable.index + ) + checkpoint: CheckpointedResult = execution_state.get_checkpoint_result( + operation_id + ) + if checkpoint.is_succeeded(): + result: ResultType = self._execute_item_in_child_context( + executor_context, executable + ) + return BatchItem(executable.index, BatchItemStatus.SUCCEEDED, result) + if checkpoint.is_failed(): + return BatchItem( + executable.index, BatchItemStatus.FAILED, error=checkpoint.error + ) + if self.nesting_type is NestingType.FLAT: + try: + flat_result: ResultType = self._execute_item_in_child_context( + executor_context, executable + ) + except Exception as e: # noqa: BLE001 + return BatchItem( + executable.index, + BatchItemStatus.FAILED, + error=ErrorObject.from_exception(e), + ) + return BatchItem(executable.index, BatchItemStatus.SUCCEEDED, flat_result) + return BatchItem(executable.index, BatchItemStatus.STARTED) + + def _replay_from_checkpoints( + self, execution_state: ExecutionState, executor_context: DurableContext + ) -> BatchResult[ResultType]: + """Derive every branch from its child checkpoint. + + Fallback reconstruction for summaries that carry no completion + record. Every executable is represented, matching the live results + that predate the recorded decision. """ items: list[BatchItem[ResultType]] = [] for executable in self.executables: @@ -545,6 +555,3 @@ def replay(self, execution_state: ExecutionState, executor_context: DurableConte batch_item = BatchItem(executable.index, status, result=result, error=error) items.append(batch_item) return BatchResult.from_items(items, self.completion_config) - - -# endregion concurrency logic diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/models.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/models.py index a8137eb2..a1bb1ff1 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/models.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/models.py @@ -1,31 +1,24 @@ -"""Concurrent executor for parallel and map operations.""" +"""Models for concurrent map and parallel execution.""" from __future__ import annotations +import json import logging -import threading -import time from collections import Counter from dataclasses import dataclass from enum import Enum from typing import TYPE_CHECKING, Generic, TypeVar -from aws_durable_execution_sdk_python.exceptions import ( - InvalidStateError, - SuspendExecution, -) +from aws_durable_execution_sdk_python.exceptions import InvalidStateError from aws_durable_execution_sdk_python.lambda_service import ErrorObject from aws_durable_execution_sdk_python.types import BatchResult as BatchResultProtocol if TYPE_CHECKING: - from concurrent.futures import Future - from aws_durable_execution_sdk_python.config import CompletionConfig logger = logging.getLogger(__name__) -T = TypeVar("T") R = TypeVar("R") CallableType = TypeVar("CallableType") @@ -46,17 +39,158 @@ class CompletionReason(Enum): @dataclass(frozen=True) -class SuspendResult: - should_suspend: bool - exception: SuspendExecution | None = None +class CompletionPolicy: + """Evaluate completion criteria for a batch of concurrent branches. - @staticmethod - def do_not_suspend() -> SuspendResult: - return SuspendResult(should_suspend=False) + Single home for the completion logic shared by the concurrency + coordinator (scheduling decisions) and :class:`BatchResult` + (completion reason inference). A user-supplied completion predicate, + if introduced later, slots in here. + + Fail-fast semantics: when no criteria are configured at all, a single + failure exceeds tolerance. + """ + + total: int + min_successful: int | None = None + tolerated_failure_count: int | None = None + tolerated_failure_percentage: int | float | None = None + + @classmethod + def from_config( + cls, total: int, config: CompletionConfig | None + ) -> CompletionPolicy: + """Build a policy for a batch of ``total`` branches from a CompletionConfig.""" + if config is None: + return cls(total=total) + return cls( + total=total, + min_successful=config.min_successful, + tolerated_failure_count=config.tolerated_failure_count, + tolerated_failure_percentage=config.tolerated_failure_percentage, + ) + + @property + def has_criteria(self) -> bool: + """True when any completion criterion is configured.""" + return ( + self.min_successful is not None + or self.tolerated_failure_count is not None + or self.tolerated_failure_percentage is not None + ) + + def is_tolerance_exceeded(self, failed: int) -> bool: + """True when failures exceed the configured tolerance.""" + if not self.has_criteria: + return failed > 0 + if ( + self.tolerated_failure_count is not None + and failed > self.tolerated_failure_count + ): + return True + if self.tolerated_failure_percentage is not None and self.total > 0: + failure_percentage: float = (failed / self.total) * 100 + return failure_percentage > self.tolerated_failure_percentage + return False + + def should_continue(self, failed: int) -> bool: + """True while more branches may be scheduled.""" + return not self.is_tolerance_exceeded(failed) + + def is_complete(self, succeeded: int, failed: int) -> bool: + """True when the batch has met a completion criterion.""" + if succeeded + failed >= self.total: + return True + return self.min_successful is not None and succeeded >= self.min_successful + + def reason(self, succeeded: int, failed: int) -> CompletionReason: + """Infer the completion reason. Tolerance is checked first.""" + if self.is_tolerance_exceeded(failed): + return CompletionReason.FAILURE_TOLERANCE_EXCEEDED + if succeeded + failed >= self.total: + return CompletionReason.ALL_COMPLETED + if self.min_successful is not None and succeeded >= self.min_successful: + return CompletionReason.MIN_SUCCESSFUL_REACHED + return CompletionReason.ALL_COMPLETED + + +@dataclass(frozen=True) +class CompletionRecord: + """The completion decision recorded in a batch operation's summary. + + Written by the map/parallel summary generators when a large result is + summarized, and read back by replay so the reconstructed result matches + the live result exactly instead of being re-derived. + + ``started_total`` is the number of branches ever started. Branches are + admitted in index order, so the started set is exactly the index prefix + ``[0, started_total)``. ``started_indexes`` is the authoritative set of + branches that were live-reported STARTED (started but not terminal when + the batch completed). + """ + + completion_reason: CompletionReason + started_total: int + started_indexes: frozenset[int] + + @classmethod + def from_summary_payload(cls, payload: str | None) -> CompletionRecord | None: + """Parse a recorded completion decision from a summary payload. + + Returns None when the payload does not carry a decision record: + absent or empty payloads, payloads from custom summary generators, + summaries written before the record existed, and unknown + completion reasons. + """ + if not payload: + return None + try: + data = json.loads(payload) + except (json.JSONDecodeError, TypeError, UnicodeDecodeError): + return None + if not isinstance(data, dict): + return None + + reason_value = data.get("completionReason") + started_total = data.get("totalCount") + if not isinstance(reason_value, str) or not isinstance(started_total, int): + return None + try: + completion_reason = CompletionReason(reason_value) + except ValueError: + return None + + started_raw = data.get("startedIndexes") + completed_raw = data.get("completedIndexes") + if isinstance(started_raw, list): + started_indexes = frozenset(started_raw) + elif isinstance(completed_raw, list): + started_indexes = frozenset(range(started_total)) - frozenset(completed_raw) + else: + return None + + return cls( + completion_reason=completion_reason, + started_total=started_total, + started_indexes=started_indexes, + ) @staticmethod - def suspend(exception: SuspendExecution) -> SuspendResult: - return SuspendResult(should_suspend=True, exception=exception) + def summary_index_fields(items: list[BatchItem]) -> dict[str, list[int]]: + """Build the index field for a summary from live batch items. + + Records whichever of the started or completed index sets is + smaller, so the record stays compact at both early-exit extremes. + """ + started: list[int] = [ + item.index for item in items if item.status is BatchItemStatus.STARTED + ] + completed: list[int] = [ + item.index for item in items if item.status is not BatchItemStatus.STARTED + ] + if len(started) <= len(completed): + return {"startedIndexes": started} + return {"completedIndexes": completed} @dataclass(frozen=True) @@ -114,115 +248,27 @@ def from_dict( completion_reason = CompletionReason(completion_reason_value) return cls(batch_items, completion_reason) - @staticmethod - def _get_completion_reason( - failure_count: int, - success_count: int, - completed_count: int, - total_count: int, - completion_config: CompletionConfig | None, - ) -> CompletionReason: - """ - Determine completion reason based on completion counts. - - Logic order: - 1. Check failure tolerance FIRST (before checking if all completed) - 2. Check if all completed - 3. Check if minimum successful reached - 4. Default to ALL_COMPLETED - - Args: - failure_count: Number of failed items - success_count: Number of succeeded items - completed_count: Total completed (succeeded + failed) - total_count: Total number of items - completion_config: Optional completion configuration - - Returns: - CompletionReason enum value - """ - # STEP 1: Check tolerance first, before checking if all completed - - # Handle fail-fast behavior (no completion config or empty completion config) - if completion_config is None: - if failure_count > 0: - return CompletionReason.FAILURE_TOLERANCE_EXCEEDED - else: - # Check if completion config has any criteria set - has_any_completion_criteria = ( - completion_config.min_successful is not None - or completion_config.tolerated_failure_count is not None - or completion_config.tolerated_failure_percentage is not None - ) - - if not has_any_completion_criteria: - # Empty completion config - fail fast on any failure - if failure_count > 0: - return CompletionReason.FAILURE_TOLERANCE_EXCEEDED - else: - # Check specific tolerance thresholds - if ( - completion_config.tolerated_failure_count is not None - and failure_count > completion_config.tolerated_failure_count - ): - return CompletionReason.FAILURE_TOLERANCE_EXCEEDED - - if ( - completion_config.tolerated_failure_percentage is not None - and total_count > 0 - ): - failure_percentage = (failure_count / total_count) * 100 - if ( - failure_percentage - > completion_config.tolerated_failure_percentage - ): - return CompletionReason.FAILURE_TOLERANCE_EXCEEDED - - # STEP 2: Check if all completed - if completed_count == total_count: - return CompletionReason.ALL_COMPLETED - - # STEP 3: Check if minimum successful reached - if ( - completion_config is not None - and completion_config.min_successful is not None - and success_count >= completion_config.min_successful - ): - return CompletionReason.MIN_SUCCESSFUL_REACHED - - # STEP 4: Default - return CompletionReason.ALL_COMPLETED - @classmethod def from_items( cls, items: list[BatchItem[R]], completion_config: CompletionConfig | None = None, ): - """ - Infer completion reason based on batch item statuses and completion config. + """Infer completion reason from batch item statuses and completion config. - This follows the same logic as the TypeScript implementation. + The batch total is derived from the items present, so this is only + exact when every branch of the batch is represented. The concurrency + executor, which may omit never-started branches, computes the reason + against the true total instead. """ - statuses = (item.status for item in items) - counts = Counter(statuses) - succeeded_count = counts.get(BatchItemStatus.SUCCEEDED, 0) - failed_count = counts.get(BatchItemStatus.FAILED, 0) - started_count = counts.get(BatchItemStatus.STARTED, 0) - - completed_count = succeeded_count + failed_count - total_count = started_count + completed_count - - # Determine completion reason using the same logic as JavaScript SDK - completion_reason = cls._get_completion_reason( - failure_count=failed_count, - success_count=succeeded_count, - completed_count=completed_count, - total_count=total_count, - completion_config=completion_config, - ) + counts: Counter[BatchItemStatus] = Counter(item.status for item in items) + succeeded_count: int = counts.get(BatchItemStatus.SUCCEEDED, 0) + failed_count: int = counts.get(BatchItemStatus.FAILED, 0) - return cls(items, completion_reason) + policy: CompletionPolicy = CompletionPolicy.from_config( + len(items), completion_config + ) + return cls(items, policy.reason(succeeded_count, failed_count)) def to_dict(self) -> dict: return { @@ -313,228 +359,99 @@ class BranchStatus(Enum): FAILED = "failed" -class ExecutableWithState(Generic[CallableType, ResultType]): - """Manages the execution state and lifecycle of an executable.""" - - def __init__(self, executable: Executable[CallableType]): - self.executable = executable - self._status = BranchStatus.PENDING - self._future: Future | None = None - self._suspend_until: float | None = None - self._result: ResultType = None # type: ignore[assignment] - self._is_result_set: bool = False - self._error: Exception | None = None +class Branch(Generic[CallableType, ResultType]): + """Execution state of a single branch. - @property - def future(self) -> Future: - """Get the future, raising error if not available.""" - if self._future is None: - msg = f"ExecutableWithState was never started. {self.executable.index}" - raise InvalidStateError(msg) - return self._future + Owned exclusively by the coordinator thread: every transition happens + there, so no synchronization is needed. + """ - @property - def status(self) -> BranchStatus: - """Get current status.""" - return self._status - - @property - def result(self) -> ResultType: - """Get result if completed.""" - if not self._is_result_set or self._status != BranchStatus.COMPLETED: - msg = f"result not available in status {self._status}" - raise InvalidStateError(msg) - return self._result - - @property - def error(self) -> Exception: - """Get error if failed.""" - if self._error is None or self._status != BranchStatus.FAILED: - msg = f"error not available in status {self._status}" - raise InvalidStateError(msg) - return self._error - - @property - def suspend_until(self) -> float | None: - """Get suspend timestamp.""" - return self._suspend_until - - @property - def is_running(self) -> bool: - """Check if currently running.""" - return self._status is BranchStatus.RUNNING - - @property - def can_resume(self) -> bool: - """Check if can resume from suspension.""" - return self._status is BranchStatus.SUSPENDED or ( - self._status is BranchStatus.SUSPENDED_WITH_TIMEOUT - and self._suspend_until is not None - and time.time() >= self._suspend_until - ) + def __init__(self, executable: Executable[CallableType]) -> None: + self.executable = executable + self.status: BranchStatus = BranchStatus.PENDING + self.result: ResultType | None = None + self.error: Exception | None = None + self.resume_at: float | None = None @property def index(self) -> int: return self.executable.index - @property - def callable(self) -> CallableType: - return self.executable.func - - # region State transitions - def run(self, future: Future) -> None: - """Transition to RUNNING state with a future.""" - if self._status != BranchStatus.PENDING: - msg = f"Cannot start running from {self._status}" + def start(self) -> None: + """Transition to RUNNING for initial submission or timed resume.""" + if self.status not in { + BranchStatus.PENDING, + BranchStatus.SUSPENDED_WITH_TIMEOUT, + }: + msg = f"Cannot start branch {self.index} from {self.status}" raise InvalidStateError(msg) - self._status = BranchStatus.RUNNING - self._future = future + self.status = BranchStatus.RUNNING + self.resume_at = None - def suspend(self) -> None: - """Transition to SUSPENDED state (indefinite).""" - self._status = BranchStatus.SUSPENDED - self._suspend_until = None - - def suspend_with_timeout(self, timestamp: float) -> None: - """Transition to SUSPENDED_WITH_TIMEOUT state.""" - self._status = BranchStatus.SUSPENDED_WITH_TIMEOUT - self._suspend_until = timestamp - - def complete(self, result: ResultType) -> None: - """Transition to COMPLETED state.""" - self._status = BranchStatus.COMPLETED - self._result = result - self._is_result_set = True + def complete(self, result: ResultType | None) -> None: + self.status = BranchStatus.COMPLETED + self.result = result def fail(self, error: Exception) -> None: - """Transition to FAILED state.""" - self._status = BranchStatus.FAILED - self._error = error + self.status = BranchStatus.FAILED + self.error = error - def reset_to_pending(self) -> None: - """Reset to PENDING state for resubmission.""" - self._status = BranchStatus.PENDING - self._future = None - self._suspend_until = None + def suspend(self) -> None: + """Suspend indefinitely, pending an external callback.""" + self.status = BranchStatus.SUSPENDED + self.resume_at = None - # endregion State transitions + def suspend_until(self, resume_at: float) -> None: + """Suspend until a timestamp, eligible for in-process resume.""" + self.status = BranchStatus.SUSPENDED_WITH_TIMEOUT + self.resume_at = resume_at -class ExecutionCounters: - """Thread-safe counters for tracking execution state.""" +class BranchEventKind(Enum): + COMPLETED = "completed" + FAILED = "failed" + SUSPENDED = "suspended" + SUSPENDED_UNTIL = "suspended_until" + ORPHANED = "orphaned" - def __init__( - self, - total_tasks: int, - min_successful: int, - tolerated_failure_count: int | None, - tolerated_failure_percentage: float | None, - ): - self.total_tasks: int = total_tasks - self.min_successful: int = min_successful - self.tolerated_failure_count: int | None = tolerated_failure_count - self.tolerated_failure_percentage: float | None = tolerated_failure_percentage - self.success_count: int = 0 - self.failure_count: int = 0 - self._lock = threading.Lock() - - def complete_task(self) -> None: - """Task completed successfully.""" - with self._lock: - self.success_count += 1 - - def fail_task(self) -> None: - """Task failed.""" - with self._lock: - self.failure_count += 1 - - def should_continue(self) -> bool: - """ - Check if we should continue starting new tasks (based on failure tolerance). - Matches TypeScript shouldContinue() logic. - """ - with self._lock: - # If no completion config, only continue if no failures - if ( - self.tolerated_failure_count is None - and self.tolerated_failure_percentage is None - ): - return self.failure_count == 0 - - # Check failure count tolerance - if ( - self.tolerated_failure_count is not None - and self.failure_count > self.tolerated_failure_count - ): - return False - - # Check failure percentage tolerance - if self.tolerated_failure_percentage is not None and self.total_tasks > 0: - failure_percentage = (self.failure_count / self.total_tasks) * 100 - if failure_percentage > self.tolerated_failure_percentage: - return False - return True +@dataclass(frozen=True) +class BranchEvent(Generic[ResultType]): + """Outcome of one branch execution attempt. - def is_complete(self) -> bool: - """ - Check if execution should complete (based on completion criteria). - Matches TypeScript isComplete() logic. - """ - with self._lock: - completed_count = self.success_count + self.failure_count + The only message worker threads send to the coordinator. Workers never + touch branch state directly. + """ - # All tasks completed - if completed_count == self.total_tasks: - return True + index: int + kind: BranchEventKind + result: ResultType | None = None + error: Exception | None = None + resume_at: float | None = None - # when we breach min successful, we've completed - return self.success_count >= self.min_successful + @classmethod + def completed( + cls, index: int, result: ResultType | None + ) -> BranchEvent[ResultType]: + return cls(index=index, kind=BranchEventKind.COMPLETED, result=result) - def should_complete(self) -> bool: - """ - Check if execution should complete. - Combines TypeScript shouldContinue() and isComplete() logic. - """ - return self.is_complete() or not self.should_continue() - - def is_all_completed(self) -> bool: - """True if all tasks completed successfully.""" - with self._lock: - return self.success_count == self.total_tasks - - def is_min_successful_reached(self) -> bool: - """True if minimum successful tasks reached.""" - with self._lock: - return self.success_count >= self.min_successful - - def is_failure_tolerance_exceeded(self) -> bool: - """True if failure tolerance was exceeded.""" - with self._lock: - return self._is_failure_condition_reached( - tolerated_count=self.tolerated_failure_count, - tolerated_percentage=self.tolerated_failure_percentage, - failure_count=self.failure_count, - ) + @classmethod + def failed(cls, index: int, error: Exception) -> BranchEvent[ResultType]: + return cls(index=index, kind=BranchEventKind.FAILED, error=error) - def _is_failure_condition_reached( - self, - tolerated_count: int | None, - tolerated_percentage: float | None, - failure_count: int, - ) -> bool: - """True if failure conditions are reached (no locking - caller must lock).""" - # Failure count condition - if tolerated_count is not None and failure_count > tolerated_count: - return True + @classmethod + def suspended(cls, index: int) -> BranchEvent[ResultType]: + return cls(index=index, kind=BranchEventKind.SUSPENDED) - # Failure percentage condition - if tolerated_percentage is not None and self.total_tasks > 0: - failure_percentage = (failure_count / self.total_tasks) * 100 - if failure_percentage > tolerated_percentage: - return True + @classmethod + def suspended_until(cls, index: int, resume_at: float) -> BranchEvent[ResultType]: + return cls( + index=index, kind=BranchEventKind.SUSPENDED_UNTIL, resume_at=resume_at + ) - return False + @classmethod + def orphaned(cls, index: int) -> BranchEvent[ResultType]: + return cls(index=index, kind=BranchEventKind.ORPHANED) -# endegion concurrency models +# endregion concurrency models diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/config.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/config.py index 71eb01ea..ac4b4e94 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/config.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/config.py @@ -152,6 +152,28 @@ class CompletionConfig: tolerated_failure_count: int | None = None tolerated_failure_percentage: int | float | None = None + def __post_init__(self): + if self.min_successful is not None and self.min_successful < 1: + msg = f"min_successful must be at least 1, got: {self.min_successful}" + raise ValidationError(msg) + if ( + self.tolerated_failure_count is not None + and self.tolerated_failure_count < 0 + ): + msg = ( + "tolerated_failure_count must be non-negative, got: " + f"{self.tolerated_failure_count}" + ) + raise ValidationError(msg) + if self.tolerated_failure_percentage is not None and not ( + 0 <= self.tolerated_failure_percentage <= 100 # noqa: PLR2004 + ): + msg = ( + "tolerated_failure_percentage must be between 0 and 100, got: " + f"{self.tolerated_failure_percentage}" + ) + raise ValidationError(msg) + # TODO: reevaluate this # @staticmethod # def first_completed(): @@ -246,6 +268,11 @@ class ParallelConfig: summary_generator: SummaryGenerator | None = None nesting_type: NestingType = NestingType.NESTED + def __post_init__(self): + if self.max_concurrency is not None and self.max_concurrency < 1: + msg = f"max_concurrency must be at least 1, got: {self.max_concurrency}" + raise ValidationError(msg) + @dataclass(frozen=True) class ParallelBranch(Generic[T]): @@ -472,6 +499,11 @@ class MapConfig(Generic[T]): nesting_type: NestingType = NestingType.NESTED item_namer: Callable[[T, int], str] | None = None + def __post_init__(self): + if self.max_concurrency is not None and self.max_concurrency < 1: + msg = f"max_concurrency must be at least 1, got: {self.max_concurrency}" + raise ValidationError(msg) + @dataclass(frozen=True) class InvokeConfig(Generic[P, R]): diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/map.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/map.py index f201efce..9e381656 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/map.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/map.py @@ -10,6 +10,7 @@ from aws_durable_execution_sdk_python.concurrency.executor import ConcurrentExecutor from aws_durable_execution_sdk_python.concurrency.models import ( BatchResult, + CompletionRecord, Executable, ) from aws_durable_execution_sdk_python.config import MapConfig, NestingType @@ -132,7 +133,7 @@ def map_handler( ) if checkpoint.is_succeeded(): # if we've reached this point, then not only is the step succeeded, but it is also `replay_children`. - return executor.replay(execution_state, map_context) + return executor.replay(execution_state, map_context, checkpoint) # we are making it explicit that we are now executing within the map_context return executor.execute(execution_state, executor_context=map_context) @@ -146,5 +147,6 @@ def __call__(self, result: BatchResult) -> str: "completionReason": result.completion_reason.value, "status": result.status.value, "type": "MapResult", + **CompletionRecord.summary_index_fields(result.all), } return json.dumps(fields) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/parallel.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/parallel.py index 76fc16f0..0593dd5e 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/parallel.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/parallel.py @@ -8,7 +8,10 @@ from typing import TYPE_CHECKING, TypeVar from aws_durable_execution_sdk_python.concurrency.executor import ConcurrentExecutor -from aws_durable_execution_sdk_python.concurrency.models import Executable +from aws_durable_execution_sdk_python.concurrency.models import ( + CompletionRecord, + Executable, +) from aws_durable_execution_sdk_python.config import ( NestingType, ParallelBranch, @@ -124,7 +127,7 @@ def parallel_handler( operation_identifier.operation_id ) if checkpoint.is_succeeded(): - return executor.replay(execution_state, parallel_context) + return executor.replay(execution_state, parallel_context, checkpoint) return executor.execute(execution_state, executor_context=parallel_context) @@ -138,6 +141,7 @@ def __call__(self, result: BatchResult) -> str: "status": result.status.value, "startedCount": result.started_count, "type": "ParallelResult", + **CompletionRecord.summary_index_fields(result.all), } return json.dumps(fields) diff --git a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py index 887e91c3..4b23c38f 100644 --- a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py @@ -1,10 +1,11 @@ """Tests for the concurrency module.""" import json +import queue import random import threading import time -from concurrent.futures import Future +from collections.abc import Callable from functools import partial from itertools import combinations from unittest.mock import Mock, patch @@ -13,17 +14,18 @@ from aws_durable_execution_sdk_python.concurrency.executor import ( ConcurrentExecutor, - TimerScheduler, ) from aws_durable_execution_sdk_python.concurrency.models import ( BatchItem, BatchItemStatus, BatchResult, + Branch, + BranchEventKind, BranchStatus, + CompletionPolicy, + CompletionRecord, CompletionReason, Executable, - ExecutableWithState, - ExecutionCounters, ) from aws_durable_execution_sdk_python.config import ( ChildConfig, @@ -38,6 +40,7 @@ from aws_durable_execution_sdk_python.exceptions import ( CallableRuntimeError, InvalidStateError, + OrphanedChildException, SuspendExecution, TimedSuspendExecution, ) @@ -48,7 +51,10 @@ OperationSubType, OperationType, ) -from aws_durable_execution_sdk_python.operation.map import MapExecutor +from aws_durable_execution_sdk_python.operation.map import ( + MapExecutor, + MapSummaryGenerator, +) from aws_durable_execution_sdk_python.state import CheckpointedResult @@ -453,14 +459,14 @@ def test_batch_result_from_dict_with_explicit_completion_reason(): def test_batch_result_infer_completion_reason_edge_cases(): """Test _infer_completion_reason method with various edge cases.""" - # Test with only started items and min_successful=0 + # Test with min_successful reached while other items are still started started_items = [ - BatchItem(0, BatchItemStatus.STARTED).to_dict(), + BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok").to_dict(), BatchItem(1, BatchItemStatus.STARTED).to_dict(), ] items = {"all": started_items} - batch = BatchResult.from_dict(items, CompletionConfig(0)) # SLF001 - # With min_successful=0 and no failures, should be MIN_SUCCESSFUL_REACHED + batch = BatchResult.from_dict(items, CompletionConfig(1)) # SLF001 + # With min_successful=1 and one success, should be MIN_SUCCESSFUL_REACHED assert batch.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED # Test with only started items and no config @@ -547,311 +553,6 @@ def test_func(): assert executable.func == test_func -def test_executable_with_state_creation(): - """Test ExecutableWithState creation.""" - executable = Executable(index=1, func=lambda: "test") - exe_state = ExecutableWithState(executable) - - assert exe_state.executable == executable - assert exe_state.status == BranchStatus.PENDING - assert exe_state.index == 1 - assert exe_state.callable == executable.func - - -def test_executable_with_state_properties(): - """Test ExecutableWithState property access.""" - - def test_callable(): - return "test" - - executable = Executable(index=42, func=test_callable) - exe_state = ExecutableWithState(executable) - - assert exe_state.index == 42 - assert exe_state.callable == test_callable - assert exe_state.suspend_until is None - - -def test_executable_with_state_future_not_available(): - """Test ExecutableWithState future property when not started.""" - executable = Executable(index=1, func=lambda: "test") - exe_state = ExecutableWithState(executable) - - with pytest.raises(InvalidStateError): - _ = exe_state.future - - -def test_executable_with_state_result_not_available(): - """Test ExecutableWithState result property when not completed.""" - executable = Executable(index=1, func=lambda: "test") - exe_state = ExecutableWithState(executable) - - with pytest.raises(InvalidStateError): - _ = exe_state.result - - -def test_executable_with_state_error_not_available(): - """Test ExecutableWithState error property when not failed.""" - executable = Executable(index=1, func=lambda: "test") - exe_state = ExecutableWithState(executable) - - with pytest.raises(InvalidStateError): - _ = exe_state.error - - -def test_executable_with_state_is_running(): - """Test ExecutableWithState is_running property.""" - executable = Executable(index=1, func=lambda: "test") - exe_state = ExecutableWithState(executable) - - assert not exe_state.is_running - - future = Future() - exe_state.run(future) - assert exe_state.is_running - - -def test_executable_with_state_can_resume(): - """Test ExecutableWithState can_resume property.""" - executable = Executable(index=1, func=lambda: "test") - exe_state = ExecutableWithState(executable) - - # Not suspended - assert not exe_state.can_resume - - # Suspended indefinitely - exe_state.suspend() - assert exe_state.can_resume - - # Suspended with timeout in future - future_time = time.time() + 10 - exe_state.suspend_with_timeout(future_time) - assert not exe_state.can_resume - - # Suspended with timeout in past - past_time = time.time() - 10 - exe_state.suspend_with_timeout(past_time) - assert exe_state.can_resume - - -def test_executable_with_state_run(): - """Test ExecutableWithState run method.""" - executable = Executable(index=1, func=lambda: "test") - exe_state = ExecutableWithState(executable) - future = Future() - - exe_state.run(future) - assert exe_state.status == BranchStatus.RUNNING - assert exe_state.future == future - - -def test_executable_with_state_run_invalid_state(): - """Test ExecutableWithState run method from invalid state.""" - executable = Executable(index=1, func=lambda: "test") - exe_state = ExecutableWithState(executable) - future1 = Future() - future2 = Future() - - exe_state.run(future1) - - with pytest.raises(InvalidStateError): - exe_state.run(future2) - - -def test_executable_with_state_suspend(): - """Test ExecutableWithState suspend method.""" - executable = Executable(index=1, func=lambda: "test") - exe_state = ExecutableWithState(executable) - - exe_state.suspend() - assert exe_state.status == BranchStatus.SUSPENDED - assert exe_state.suspend_until is None - - -def test_executable_with_state_suspend_with_timeout(): - """Test ExecutableWithState suspend_with_timeout method.""" - executable = Executable(index=1, func=lambda: "test") - exe_state = ExecutableWithState(executable) - timestamp = time.time() + 5 - - exe_state.suspend_with_timeout(timestamp) - assert exe_state.status == BranchStatus.SUSPENDED_WITH_TIMEOUT - assert exe_state.suspend_until == timestamp - - -def test_executable_with_state_complete(): - """Test ExecutableWithState complete method.""" - executable = Executable(index=1, func=lambda: "test") - exe_state = ExecutableWithState(executable) - - exe_state.complete("test_result") - assert exe_state.status == BranchStatus.COMPLETED - assert exe_state.result == "test_result" - - -def test_executable_with_state_fail(): - """Test ExecutableWithState fail method.""" - executable = Executable(index=1, func=lambda: "test") - exe_state = ExecutableWithState(executable) - error = Exception("test error") - - exe_state.fail(error) - assert exe_state.status == BranchStatus.FAILED - assert exe_state.error == error - - -def test_execution_counters_creation(): - """Test ExecutionCounters creation.""" - counters = ExecutionCounters( - total_tasks=10, - min_successful=8, - tolerated_failure_count=2, - tolerated_failure_percentage=20.0, - ) - - assert counters.total_tasks == 10 - assert counters.min_successful == 8 - assert counters.tolerated_failure_count == 2 - assert counters.tolerated_failure_percentage == 20.0 - assert counters.success_count == 0 - assert counters.failure_count == 0 - - -def test_execution_counters_complete_task(): - """Test ExecutionCounters complete_task method.""" - counters = ExecutionCounters(5, 3, None, None) - - counters.complete_task() - assert counters.success_count == 1 - - -def test_execution_counters_fail_task(): - """Test ExecutionCounters fail_task method.""" - counters = ExecutionCounters(5, 3, None, None) - - counters.fail_task() - assert counters.failure_count == 1 - - -def test_execution_counters_should_complete_min_successful(): - """Test ExecutionCounters should_complete with min successful reached.""" - counters = ExecutionCounters(5, 3, None, None) - - assert not counters.should_complete() - - counters.complete_task() - counters.complete_task() - counters.complete_task() - - assert counters.should_complete() - - -def test_execution_counters_should_complete_failure_count(): - """Test ExecutionCounters should_complete with failure count exceeded.""" - counters = ExecutionCounters(5, 3, 1, None) - - assert not counters.should_complete() - - counters.fail_task() - assert not counters.should_complete() - - counters.fail_task() - assert counters.should_complete() - - -def test_execution_counters_should_complete_failure_percentage(): - """Test ExecutionCounters should_complete with failure percentage exceeded.""" - counters = ExecutionCounters(10, 8, None, 15.0) - - assert not counters.should_complete() - - counters.fail_task() - assert not counters.should_complete() - - counters.fail_task() - assert counters.should_complete() # 20% > 15% - - -def test_execution_counters_is_all_completed(): - """Test ExecutionCounters is_all_completed method.""" - counters = ExecutionCounters(3, 2, None, None) - - assert not counters.is_all_completed() - - counters.complete_task() - counters.complete_task() - assert not counters.is_all_completed() - - counters.complete_task() - assert counters.is_all_completed() - - -def test_execution_counters_is_min_successful_reached(): - """Test ExecutionCounters is_min_successful_reached method.""" - counters = ExecutionCounters(5, 3, None, None) - - assert not counters.is_min_successful_reached() - - counters.complete_task() - counters.complete_task() - assert not counters.is_min_successful_reached() - - counters.complete_task() - assert counters.is_min_successful_reached() - - -def test_execution_counters_is_failure_tolerance_exceeded(): - """Test ExecutionCounters is_failure_tolerance_exceeded method.""" - counters = ExecutionCounters(10, 8, 2, None) - - assert not counters.is_failure_tolerance_exceeded() - - counters.fail_task() - counters.fail_task() - assert not counters.is_failure_tolerance_exceeded() - - counters.fail_task() - assert counters.is_failure_tolerance_exceeded() - - -def test_execution_counters_zero_total_tasks(): - """Test ExecutionCounters with zero total tasks.""" - counters = ExecutionCounters(0, 0, None, 50.0) - - # Should not fail with division by zero - assert not counters.is_failure_tolerance_exceeded() - - -def test_execution_counters_failure_percentage_edge_case(): - """Test ExecutionCounters failure percentage at exact threshold.""" - counters = ExecutionCounters(10, 5, None, 20.0) - - # Exactly at threshold (20%) - counters.failure_count = 2 - assert not counters.is_failure_tolerance_exceeded() - - # Just over threshold - counters.failure_count = 3 - assert counters.is_failure_tolerance_exceeded() - - -def test_execution_counters_thread_safety(): - """Test ExecutionCounters thread safety.""" - counters = ExecutionCounters(100, 50, None, None) - - def worker(): - for _ in range(10): - counters.complete_task() - - threads = [threading.Thread(target=worker) for _ in range(5)] - for t in threads: - t.start() - for t in threads: - t.join() - - assert counters.success_count == 50 - - def test_batch_result_failed_with_none_error(): """Test BatchResult failed method filters out None errors.""" items = [ @@ -965,137 +666,6 @@ def mock_run_in_child_context(func, name, config): assert len(result.all) >= 1 -def test_timer_scheduler_double_check_resume_queue(): - """Test TimerScheduler double-check logic in scheduler loop.""" - callback = Mock() - - with TimerScheduler(callback) as scheduler: - exe_state1 = ExecutableWithState(Executable(0, lambda: "test")) - exe_state2 = ExecutableWithState(Executable(1, lambda: "test")) - - # Schedule two tasks with different times to avoid comparison issues - past_time1 = time.time() - 2 - past_time2 = time.time() - 1 - scheduler.schedule_resume(exe_state1, past_time1) - scheduler.schedule_resume(exe_state2, past_time2) - - # Give scheduler time to process - time.sleep(0.1) - - # At least one callback should have been made - assert callback.call_count >= 0 - - -def test_concurrent_executor_on_task_complete_timed_suspend(): - """Test ConcurrentExecutor _on_task_complete with TimedSuspendExecution.""" - - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" - - executables = [Executable(0, lambda: "test")] - completion_config = CompletionConfig( - min_successful=1, - tolerated_failure_count=None, - tolerated_failure_percentage=None, - ) - - executor = TestExecutor( - executables=executables, - max_concurrency=1, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) - - exe_state = ExecutableWithState(executables[0]) - future = Mock() - future.result.side_effect = TimedSuspendExecution("test message", time.time() + 1) - future.cancelled.return_value = False - - scheduler = Mock() - scheduler.schedule_resume = Mock() - - executor._on_task_complete(exe_state, future, scheduler) # noqa: SLF001 - - assert exe_state.status == BranchStatus.SUSPENDED_WITH_TIMEOUT - scheduler.schedule_resume.assert_called_once() - - -def test_concurrent_executor_on_task_complete_suspend(): - """Test ConcurrentExecutor _on_task_complete with SuspendExecution.""" - - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" - - executables = [Executable(0, lambda: "test")] - completion_config = CompletionConfig( - min_successful=1, - tolerated_failure_count=None, - tolerated_failure_percentage=None, - ) - - executor = TestExecutor( - executables=executables, - max_concurrency=1, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) - - exe_state = ExecutableWithState(executables[0]) - future = Mock() - future.result.side_effect = SuspendExecution("test message") - - scheduler = Mock() - - executor._on_task_complete(exe_state, future, scheduler) # noqa: SLF001 - - assert exe_state.status == BranchStatus.SUSPENDED - - -def test_concurrent_executor_on_task_complete_exception(): - """Test ConcurrentExecutor _on_task_complete with general exception.""" - - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" - - executables = [Executable(0, lambda: "test")] - completion_config = CompletionConfig( - min_successful=1, - tolerated_failure_count=None, - tolerated_failure_percentage=None, - ) - - executor = TestExecutor( - executables=executables, - max_concurrency=1, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) - - exe_state = ExecutableWithState(executables[0]) - future = Mock() - future.result.side_effect = ValueError("Test error") - future.cancelled.return_value = False - - scheduler = Mock() - - executor._on_task_complete(exe_state, future, scheduler) # noqa: SLF001 - - assert exe_state.status == BranchStatus.FAILED - assert isinstance(exe_state.error, ValueError) - - def test_concurrent_executor_create_result_with_early_exit(): """Test ConcurrentExecutor with failed branches using public execute method.""" @@ -1320,18 +890,6 @@ def execute_item(self, child_context, executable): child_context._set_replay_status_new.assert_not_called() # noqa: SLF001 -def test_execution_counters_impossible_to_succeed(): - """Test ExecutionCounters should_complete when impossible to succeed.""" - counters = ExecutionCounters(5, 4, None, None) - - # Fail 3 tasks, leaving only 2 remaining (can't reach min_successful of 4) - counters.fail_task() - counters.fail_task() - counters.fail_task() - - assert counters.should_complete() - - def test_concurrent_executor_create_result_failure_tolerance_exceeded(): """Test ConcurrentExecutor with failure tolerance exceeded using public execute method.""" @@ -1658,36 +1216,21 @@ def execute_item(self, child_context, executable): assert executor.call_counts[0] == 1 -def test_timer_scheduler_double_check_condition(): - """Test TimerScheduler double-check condition in _timer_loop (line 434).""" - callback = Mock() - - with TimerScheduler(callback) as scheduler: - exe_state = ExecutableWithState(Executable(0, lambda: "test")) - exe_state.suspend() # Make it resumable - - # Schedule a task with past time - past_time = time.time() - 1 - scheduler.schedule_resume(exe_state, past_time) - - # Give scheduler time to process and hit the double-check condition - time.sleep(0.2) - - # The callback should be called - assert callback.call_count >= 1 - - -def test_concurrent_executor_should_execution_suspend_with_timeout(): - """Test should_execution_suspend with SUSPENDED_WITH_TIMEOUT state.""" +def test_concurrent_executor_create_result_with_failed_status(): + """Test with failed executable status using public execute method.""" class TestExecutor(ConcurrentExecutor): def execute_item(self, child_context, executable): - return f"result_{executable.index}" + msg = "Test error" + raise ValueError(msg) - executables = [Executable(0, lambda: "test")] + def failure_callable(): + return "test" + + executables = [Executable(0, failure_callable)] completion_config = CompletionConfig( min_successful=1, - tolerated_failure_count=None, + tolerated_failure_count=0, tolerated_failure_percentage=None, ) @@ -1701,32 +1244,37 @@ def execute_item(self, child_context, executable): serdes=None, ) - # Create executable with state in SUSPENDED_WITH_TIMEOUT - exe_state = ExecutableWithState(executables[0]) - future_time = time.time() + 10 - exe_state.suspend_with_timeout(future_time) + execution_state = Mock() + execution_state.create_checkpoint = Mock() - executor.executables_with_state = [exe_state] + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + child_context = Mock() + child_context.state.wrap_user_function = lambda func, *args, **kwargs: func + executor_context.create_child_context = lambda *args, **kwargs: child_context - result = executor.should_execution_suspend() + result = executor.execute(execution_state, executor_context) - assert result.should_suspend - assert isinstance(result.exception, TimedSuspendExecution) - assert result.exception.scheduled_timestamp == future_time + assert len(result.all) == 1 + assert result.all[0].status == BatchItemStatus.FAILED + assert result.all[0].error is not None + assert result.all[0].error.message == "Test error" -def test_concurrent_executor_should_execution_suspend_indefinite(): - """Test should_execution_suspend with indefinite SUSPENDED state.""" +def test_concurrent_executor_execute_with_failing_task(): + """Test execute() with a task that fails using public execute method.""" class TestExecutor(ConcurrentExecutor): def execute_item(self, child_context, executable): - return f"result_{executable.index}" + msg = "Task failed" + raise ValueError(msg) - executables = [Executable(0, lambda: "test")] + def failure_callable(): + return "test" + + executables = [Executable(0, failure_callable)] completion_config = CompletionConfig( - min_successful=1, - tolerated_failure_count=None, - tolerated_failure_percentage=None, + min_successful=1, tolerated_failure_count=0, tolerated_failure_percentage=None ) executor = TestExecutor( @@ -1739,34 +1287,36 @@ def execute_item(self, child_context, executable): serdes=None, ) - # Create executable with state in SUSPENDED (indefinite) - exe_state = ExecutableWithState(executables[0]) - exe_state.suspend() + execution_state = Mock() + execution_state.create_checkpoint = Mock() - executor.executables_with_state = [exe_state] + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + child_context = Mock() + child_context.state.wrap_user_function = lambda func, *args, **kwargs: func + executor_context.create_child_context = lambda *args, **kwargs: child_context - result = executor.should_execution_suspend() + result = executor.execute(execution_state, executor_context) - assert result.should_suspend - assert isinstance(result.exception, SuspendExecution) - assert "pending external callback" in str(result.exception) + assert len(result.all) == 1 + assert result.all[0].status == BatchItemStatus.FAILED + assert result.all[0].error.message == "Task failed" -def test_concurrent_executor_create_result_with_failed_status(): - """Test with failed executable status using public execute method.""" +def test_create_result_no_failed_executables(): + """Test when no executables are failed using public execute method.""" class TestExecutor(ConcurrentExecutor): def execute_item(self, child_context, executable): - msg = "Test error" - raise ValueError(msg) + return f"result_{executable.index}" - def failure_callable(): + def success_callable(): return "test" - executables = [Executable(0, failure_callable)] + executables = [Executable(0, success_callable)] completion_config = CompletionConfig( min_successful=1, - tolerated_failure_count=0, + tolerated_failure_count=None, tolerated_failure_percentage=None, ) @@ -1785,47 +1335,27 @@ def failure_callable(): executor_context = Mock() executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 - child_context = Mock() - child_context.state.wrap_user_function = lambda func, *args, **kwargs: func - executor_context.create_child_context = lambda *args, **kwargs: child_context + executor_context.create_child_context = lambda *args, **kwargs: Mock() result = executor.execute(execution_state, executor_context) assert len(result.all) == 1 - assert result.all[0].status == BatchItemStatus.FAILED - assert result.all[0].error is not None - assert result.all[0].error.message == "Test error" - - -def test_timer_scheduler_can_resume_false(): - """Test TimerScheduler when exe_state.can_resume is False.""" - callback = Mock() - - with TimerScheduler(callback) as scheduler: - exe_state = ExecutableWithState(Executable(0, lambda: "test")) - - # Set state to something that can't resume - exe_state.complete("done") - - # Schedule with past time - past_time = time.time() - 1 - scheduler.schedule_resume(exe_state, past_time) - - # Give scheduler time to process - time.sleep(0.15) - - # Callback should not be called since can_resume is False - callback.assert_not_called() + assert result.all[0].status == BatchItemStatus.SUCCEEDED + assert result.completion_reason == CompletionReason.ALL_COMPLETED -def test_concurrent_executor_mixed_suspend_states(): - """Test should_execution_suspend with mixed suspend states.""" +def test_create_result_with_suspended_executable(): + """Test with suspended executable using public execute method.""" class TestExecutor(ConcurrentExecutor): def execute_item(self, child_context, executable): - return f"result_{executable.index}" + msg = "Test suspend" + raise SuspendExecution(msg) - executables = [Executable(0, lambda: "test"), Executable(1, lambda: "test2")] + def suspend_callable(): + return "test" + + executables = [Executable(0, suspend_callable)] completion_config = CompletionConfig( min_successful=1, tolerated_failure_count=None, @@ -1834,7 +1364,7 @@ def execute_item(self, child_context, executable): executor = TestExecutor( executables=executables, - max_concurrency=2, + max_concurrency=1, completion_config=completion_config, sub_type_top="TOP", sub_type_iteration="ITER", @@ -1842,270 +1372,452 @@ def execute_item(self, child_context, executable): serdes=None, ) - # Create one with timed suspend and one with indefinite suspend - exe_state1 = ExecutableWithState(executables[0]) - exe_state2 = ExecutableWithState(executables[1]) - - future_time = time.time() + 5 - exe_state1.suspend_with_timeout(future_time) - exe_state2.suspend() # Indefinite + execution_state = Mock() + execution_state.create_checkpoint = Mock() - executor.executables_with_state = [exe_state1, exe_state2] + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + child_context = Mock() + child_context.state.wrap_user_function = lambda func, *args, **kwargs: func + executor_context.create_child_context = lambda *args, **kwargs: child_context - result = executor.should_execution_suspend() + # Should raise SuspendExecution since single task suspends + with pytest.raises(SuspendExecution): + executor.execute(execution_state, executor_context) - # Should return timed suspend (earliest timestamp takes precedence) - assert result.should_suspend - assert isinstance(result.exception, TimedSuspendExecution) +# Tests for _create_result method match statement branches +def test_batch_result_from_dict_with_completion_config(): + """Test BatchResult from_dict with completion config parameter.""" + data = { + "all": [ + {"index": 0, "status": "SUCCEEDED", "result": "result1", "error": None}, + {"index": 1, "status": "STARTED", "result": None, "error": None}, + ], + # No completionReason provided + } -def test_concurrent_executor_multiple_timed_suspends(): - """Test should_execution_suspend with multiple timed suspends to find earliest.""" + # With started items, should infer MIN_SUCCESSFUL_REACHED + completion_config = CompletionConfig(min_successful=1) - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + with patch( + "aws_durable_execution_sdk_python.concurrency.models.logger" + ) as mock_logger: + result = BatchResult.from_dict(data, completion_config) + assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED + mock_logger.warning.assert_called_once() - executables = [Executable(0, lambda: "test"), Executable(1, lambda: "test2")] - completion_config = CompletionConfig( - min_successful=1, - tolerated_failure_count=None, - tolerated_failure_percentage=None, - ) - executor = TestExecutor( - executables=executables, - max_concurrency=2, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) +def test_batch_result_from_dict_all_completed(): + """Test BatchResult from_dict infers completion reason when all items are completed.""" + data = { + "all": [ + {"index": 0, "status": "SUCCEEDED", "result": "result1", "error": None}, + { + "index": 1, + "status": "FAILED", + "result": None, + "error": { + "message": "error", + "type": "Error", + "data": None, + "stackTrace": None, + }, + }, + ], + # No completionReason provided + } - # Create two with different timed suspends - exe_state1 = ExecutableWithState(executables[0]) - exe_state2 = ExecutableWithState(executables[1]) + # With no config and failures, fail-fast + with patch( + "aws_durable_execution_sdk_python.concurrency.models.logger" + ) as mock_logger: + result = BatchResult.from_dict(data) + assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED + mock_logger.warning.assert_called_once() - later_time = time.time() + 10 - earlier_time = time.time() + 5 - exe_state1.suspend_with_timeout(later_time) - exe_state2.suspend_with_timeout(earlier_time) +def test_batch_result_from_dict_backward_compatibility(): + """Test BatchResult from_dict maintains backward compatibility when no completion_config provided.""" + data = { + "all": [ + {"index": 0, "status": "SUCCEEDED", "result": "result1", "error": None} + ], + "completionReason": "MIN_SUCCESSFUL_REACHED", + } - executor.executables_with_state = [exe_state1, exe_state2] + # Should work without completion_config parameter + result = BatchResult.from_dict(data) + assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - result = executor.should_execution_suspend() + # Should also work with None completion_config + result2 = BatchResult.from_dict(data, None) + assert result2.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - # Should return the earlier timestamp - assert result.should_suspend - assert isinstance(result.exception, TimedSuspendExecution) - assert result.exception.scheduled_timestamp == earlier_time +def test_batch_result_infer_completion_reason_basic_cases(): + """Test _infer_completion_reason method with basic scenarios.""" + # Test with started items - should be MIN_SUCCESSFUL_REACHED + items = { + "all": [ + BatchItem(0, BatchItemStatus.SUCCEEDED, "result1").to_dict(), + BatchItem(1, BatchItemStatus.STARTED).to_dict(), + ] + } + batch = BatchResult.from_dict(items, CompletionConfig(1)) + assert batch.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED -def test_timer_scheduler_double_check_condition_race(): - """Test TimerScheduler double-check condition when heap changes between checks.""" - callback = Mock() + # Test with all completed items - should be ALL_COMPLETED + completed_items = [ + BatchItem(0, BatchItemStatus.SUCCEEDED, "result1").to_dict(), + BatchItem( + 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ).to_dict(), + ] + completed_items = {"all": completed_items} + batch = BatchResult.from_dict(completed_items, CompletionConfig(1)) + assert batch.completion_reason == CompletionReason.ALL_COMPLETED - with TimerScheduler(callback) as scheduler: - exe_state1 = ExecutableWithState(Executable(0, lambda: "test")) - exe_state2 = ExecutableWithState(Executable(1, lambda: "test")) + # Test empty items - should be ALL_COMPLETED + batch = BatchResult.from_dict({"all": []}, CompletionConfig(1)) + assert batch.completion_reason == CompletionReason.ALL_COMPLETED - exe_state1.suspend() - exe_state2.suspend() - # Schedule first task with past time - past_time = time.time() - 1 - scheduler.schedule_resume(exe_state1, past_time) +def test_operation_id_determinism_across_shuffles(): + """Test that operation_id depends on Executable.index, not execution order.""" - # Brief delay to let timer thread see the first task - time.sleep(0.05) + def index_based_function(index, ctx): + """Function that returns a result based on the executable index.""" + return f"result_for_index_{index}" - # Schedule second task with even more past time (will be heap[0]) - very_past_time = time.time() - 2 - scheduler.schedule_resume(exe_state2, very_past_time) + class TestExecutor(ConcurrentExecutor): + """Custom executor for testing operation_id determinism.""" - # Wait for processing - time.sleep(0.2) + def execute_item(self, child_context, executable): + return executable.func(child_context) - assert callback.call_count >= 1 + # Create executables with specific indices using partial + num_executables = 50 + funcs = [partial(index_based_function, i) for i in range(num_executables)] + # Track operation_id -> result associations + captured_associations = [] -def test_should_execution_suspend_earliest_timestamp_comparison(): - """Test should_execution_suspend timestamp comparison logic (line 554).""" + def patched_child_handler( + func, + execution_state, + operation_identifier, + config: ChildConfig, + ): + """Patched child handler that captures operation_id -> result mapping.""" + assert config.is_virtual + assert config.sub_type == "TEST_ITER" + result = func() # Execute the function + captured_associations.append((operation_identifier.operation_id, result)) + return result - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + execution_state = Mock() + execution_state.create_checkpoint = Mock() - executables = [ - Executable(0, lambda: "test"), - Executable(1, lambda: "test2"), - Executable(2, lambda: "test3"), - ] - completion_config = CompletionConfig( - min_successful=1, - tolerated_failure_count=None, - tolerated_failure_percentage=None, - ) + completion_config = CompletionConfig(min_successful=num_executables) - executor = TestExecutor( - executables=executables, - max_concurrency=3, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) + # Run multiple times with different shuffle orders + associations_per_run = [] - # Create three executables with different suspend times - exe_state1 = ExecutableWithState(executables[0]) - exe_state2 = ExecutableWithState(executables[1]) - exe_state3 = ExecutableWithState(executables[2]) + for run in range(10): # Test 10 different shuffle orders + captured_associations.clear() - time1 = time.time() + 10 - time2 = time.time() + 5 # Earliest - time3 = time.time() + 15 + # Create executables from shuffled functions + executables = [Executable(index=i, func=func) for i, func in enumerate(funcs)] + random.seed(run) # Different seed for each run + random.shuffle(executables) - exe_state1.suspend_with_timeout(time1) - exe_state2.suspend_with_timeout(time2) - exe_state3.suspend_with_timeout(time3) + executor = TestExecutor( + executables=executables, + max_concurrency=2, + completion_config=completion_config, + sub_type_top="TEST", + sub_type_iteration="TEST_ITER", + name_prefix="test_", + serdes=None, + nesting_type=NestingType.FLAT, + ) - executor.executables_with_state = [exe_state1, exe_state2, exe_state3] + # Create executor context mock + executor_context = Mock() + executor_context._parent_id = "parent_123" # noqa SLF001 - result = executor.should_execution_suspend() + def create_step_id(index): + return f"step_{index}" - assert result.should_suspend - assert isinstance(result.exception, TimedSuspendExecution) - assert result.exception.scheduled_timestamp == time2 + executor_context._create_step_id_for_logical_step = create_step_id # noqa SLF001 + def create_child_context(operation_id, *, is_virtual=False): + child_ctx = Mock() + child_ctx.state = execution_state + return child_ctx -def test_concurrent_executor_execute_with_failing_task(): - """Test execute() with a task that fails using public execute method.""" + executor_context.create_child_context = create_child_context - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - msg = "Task failed" - raise ValueError(msg) + with patch( + "aws_durable_execution_sdk_python.concurrency.executor.child_handler", + patched_child_handler, + ): + executor.execute(execution_state, executor_context) - def failure_callable(): - return "test" + associations_per_run.append(captured_associations.copy()) - executables = [Executable(0, failure_callable)] - completion_config = CompletionConfig( - min_successful=1, tolerated_failure_count=0, tolerated_failure_percentage=None + # first we will verify the validity of the test by ensuring that there exist at least 2 runs with different ordering + assert any( + assoc1 != assoc2 for assoc1, assoc2 in combinations(associations_per_run, 2) ) - - executor = TestExecutor( - executables=executables, - max_concurrency=1, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, + # then we will verify the invariant of association between step_id and result + associations_per_run = [dict(assoc) for assoc in associations_per_run] + assert all( + assoc1 == assoc2 for assoc1, assoc2 in combinations(associations_per_run, 2) ) - execution_state = Mock() - execution_state.create_checkpoint = Mock() - executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 - child_context = Mock() - child_context.state.wrap_user_function = lambda func, *args, **kwargs: func - executor_context.create_child_context = lambda *args, **kwargs: child_context +def test_concurrent_executor_replay_with_succeeded_operations(): + """Test ConcurrentExecutor replay method with succeeded operations.""" - result = executor.execute(execution_state, executor_context) + def func1(ctx, item, idx, items): + return f"result_{item}" + + items = ["a", "b"] + config = MapConfig() + + executor = MapExecutor.from_items( + items=items, + func=func1, + config=config, + ) + + # Mock execution state with succeeded operations + mock_execution_state = Mock() + mock_execution_state.durable_execution_arn = ( + "arn:aws:durable:us-east-1:123456789012:execution/test" + ) + + def mock_get_checkpoint_result(operation_id): + mock_result = Mock() + mock_result.is_succeeded.return_value = True + mock_result.is_failed.return_value = False + mock_result.is_replay_children.return_value = False + mock_result.is_existent.return_value = True + # Provide properly serialized JSON data + mock_result.result = f'"cached_result_{operation_id}"' # JSON string + return mock_result + + mock_execution_state.get_checkpoint_result = mock_get_checkpoint_result + + def mock_create_step_id_for_logical_step(step): + return f"op_{step}" + + # Mock executor context + mock_executor_context = Mock() + mock_executor_context._create_step_id_for_logical_step = ( # noqa + mock_create_step_id_for_logical_step + ) + + # Mock child context that has the same execution state + mock_child_context = Mock() + mock_child_context.state = mock_execution_state + mock_executor_context.create_child_context = Mock(return_value=mock_child_context) + mock_executor_context._parent_id = "parent_id" # noqa + + result = executor.replay(mock_execution_state, mock_executor_context) + + assert isinstance(result, BatchResult) + assert len(result.all) == 2 + assert result.all[0].status == BatchItemStatus.SUCCEEDED + assert result.all[0].result == "cached_result_op_0" + assert result.all[1].status == BatchItemStatus.SUCCEEDED + assert result.all[1].result == "cached_result_op_1" + + +def test_concurrent_executor_replay_with_failed_operations(): + """Test ConcurrentExecutor replay method with failed operations.""" + def func1(ctx, item, idx, items): + return f"result_{item}" + + items = ["a"] + config = MapConfig() + + executor = MapExecutor.from_items( + items=items, + func=func1, + config=config, + ) + + # Mock execution state with failed operation + mock_execution_state = Mock() + + def mock_get_checkpoint_result(operation_id): + mock_result = Mock() + mock_result.is_succeeded.return_value = False + mock_result.is_failed.return_value = True + mock_result.error = Exception("Test error") + return mock_result + + mock_execution_state.get_checkpoint_result = mock_get_checkpoint_result + + # Mock executor context + mock_executor_context = Mock() + mock_executor_context._create_step_id_for_logical_step = Mock(return_value="op_1") # noqa: SLF001 + + result = executor.replay(mock_execution_state, mock_executor_context) + + assert isinstance(result, BatchResult) assert len(result.all) == 1 assert result.all[0].status == BatchItemStatus.FAILED - assert result.all[0].error.message == "Task failed" + assert result.all[0].error is not None -def test_timer_scheduler_cannot_resume_branch(): - """Test TimerScheduler when exe_state cannot resume (434->433 branch).""" - callback = Mock() +def test_concurrent_executor_replay_with_replay_children(): + """Test ConcurrentExecutor replay method when children need re-execution.""" - with TimerScheduler(callback) as scheduler: - exe_state = ExecutableWithState(Executable(0, lambda: "test")) + def func1(ctx, item, idx, items): + return f"result_{item}" - # Set to completed state so can_resume returns False - exe_state.complete("done") + items = ["a"] + config = MapConfig() - # Schedule with past time - past_time = time.time() - 1 - scheduler.schedule_resume(exe_state, past_time) + executor = MapExecutor.from_items( + items=items, + func=func1, + config=config, + ) - # Wait for processing - time.sleep(0.2) + # Mock execution state with succeeded operation that needs replay + mock_execution_state = Mock() - # Callback should not be called since can_resume is False - callback.assert_not_called() + def mock_get_checkpoint_result(operation_id): + mock_result = Mock() + mock_result.is_succeeded.return_value = True + mock_result.is_failed.return_value = False + mock_result.is_replay_children.return_value = True + return mock_result + mock_execution_state.get_checkpoint_result = mock_get_checkpoint_result -def test_create_result_no_failed_executables(): - """Test when no executables are failed using public execute method.""" + # Mock executor context + mock_executor_context = Mock() + mock_executor_context._create_step_id_for_logical_step = Mock(return_value="op_1") # noqa: SLF001 - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + # Mock _execute_item_in_child_context to return a result + with patch.object( + executor, "_execute_item_in_child_context", return_value="re_executed_result" + ): + result = executor.replay(mock_execution_state, mock_executor_context) - def success_callable(): - return "test" + assert isinstance(result, BatchResult) + assert len(result.all) == 1 + assert result.all[0].status == BatchItemStatus.SUCCEEDED + assert result.all[0].result == "re_executed_result" - executables = [Executable(0, success_callable)] - completion_config = CompletionConfig( - min_successful=1, - tolerated_failure_count=None, - tolerated_failure_percentage=None, - ) - executor = TestExecutor( - executables=executables, - max_concurrency=1, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, +def test_batch_item_from_dict_with_error(): + """Test BatchItem.from_dict() with error.""" + data = { + "index": 3, + "status": "FAILED", + "result": None, + "error": { + "ErrorType": "ValueError", + "ErrorMessage": "bad value", + "StackTrace": [], + }, + } + + item = BatchItem.from_dict(data) + + assert item.index == 3 + assert item.status == BatchItemStatus.FAILED + assert item.error.type == "ValueError" + assert item.error.message == "bad value" + + +def test_batch_result_with_mixed_statuses(): + """Test BatchResult serialization with mixed item statuses.""" + result = BatchResult( + all=[ + BatchItem(0, BatchItemStatus.SUCCEEDED, result="success"), + BatchItem( + 1, + BatchItemStatus.FAILED, + error=ErrorObject(message="msg", type="E", data=None, stack_trace=[]), + ), + BatchItem(2, BatchItemStatus.STARTED), + ], + completion_reason=CompletionReason.FAILURE_TOLERANCE_EXCEEDED, ) - execution_state = Mock() - execution_state.create_checkpoint = Mock() + serialized = json.dumps(result.to_dict()) + deserialized = BatchResult.from_dict(json.loads(serialized)) - executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 - executor_context.create_child_context = lambda *args, **kwargs: Mock() + assert len(deserialized.all) == 3 + assert deserialized.all[0].status == BatchItemStatus.SUCCEEDED + assert deserialized.all[1].status == BatchItemStatus.FAILED + assert deserialized.all[2].status == BatchItemStatus.STARTED + assert deserialized.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED - result = executor.execute(execution_state, executor_context) - assert len(result.all) == 1 - assert result.all[0].status == BatchItemStatus.SUCCEEDED - assert result.completion_reason == CompletionReason.ALL_COMPLETED +def test_batch_result_empty_list(): + """Test BatchResult serialization with empty items list.""" + result = BatchResult(all=[], completion_reason=CompletionReason.ALL_COMPLETED) + serialized = json.dumps(result.to_dict()) + deserialized = BatchResult.from_dict(json.loads(serialized)) -def test_create_result_with_suspended_executable(): - """Test with suspended executable using public execute method.""" + assert len(deserialized.all) == 0 + assert deserialized.completion_reason == CompletionReason.ALL_COMPLETED + + +def test_batch_result_complex_nested_data(): + """Test BatchResult with complex nested data structures.""" + complex_result = { + "users": [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}], + "metadata": {"count": 2, "timestamp": "2025-10-31"}, + } + + result = BatchResult( + all=[BatchItem(0, BatchItemStatus.SUCCEEDED, result=complex_result)], + completion_reason=CompletionReason.ALL_COMPLETED, + ) + + serialized = json.dumps(result.to_dict()) + deserialized = BatchResult.from_dict(json.loads(serialized)) + + assert deserialized.all[0].result == complex_result + assert deserialized.all[0].result["users"][0]["name"] == "Alice" + + +def test_executor_does_not_deadlock_when_all_tasks_terminal_but_completion_config_allows_failures(): + """Ensure executor returns when all tasks are terminal even if completion rules are confusing.""" class TestExecutor(ConcurrentExecutor): def execute_item(self, child_context, executable): - msg = "Test suspend" - raise SuspendExecution(msg) - - def suspend_callable(): - return "test" + if executable.index == 0: + # fail one task + raise Exception("boom") # noqa EM101 TRY002 + return f"ok_{executable.index}" - executables = [Executable(0, suspend_callable)] + # Two tasks, min_successful=2 but tolerated failure_count set to 1. + # After one fail + one success, counters.is_complete() should return true, + # should_continue() should return false. counters.is_complete was failing to + # stop early, which caused map to hang. + executables = [Executable(0, lambda: "a"), Executable(1, lambda: "b")] completion_config = CompletionConfig( - min_successful=1, - tolerated_failure_count=None, + min_successful=2, + tolerated_failure_count=1, tolerated_failure_percentage=None, ) executor = TestExecutor( executables=executables, - max_concurrency=1, + max_concurrency=2, completion_config=completion_config, sub_type_top="TOP", sub_type_iteration="ITER", @@ -2115,103 +1827,95 @@ def suspend_callable(): execution_state = Mock() execution_state.create_checkpoint = Mock() - executor_context = Mock() executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context - # Should raise SuspendExecution since single task suspends - with pytest.raises(SuspendExecution): - executor.execute(execution_state, executor_context) + # Should return (not hang) and batch should reflect one FAILED and one SUCCEEDED + result = executor.execute(execution_state, executor_context) + statuses = {item.index: item.status for item in result.all} + assert statuses[0] == BatchItemStatus.FAILED + assert statuses[1] == BatchItemStatus.SUCCEEDED -# Tests for _create_result method match statement branches -def test_create_result_completed_branch(): - """Test _create_result with COMPLETED status branch.""" +def test_executor_terminates_quickly_when_impossible_to_succeed(): + """Test that executor terminates when min_successful becomes impossible.""" + executed_count = {"value": 0} - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + def task_func(ctx, item, idx, items): + executed_count["value"] += 1 + if idx < 2: + raise Exception(f"fail_{idx}") # noqa EM102 TRY002 + time.sleep(0.05) + return f"ok_{idx}" - executables = [Executable(0, lambda: "test")] - completion_config = CompletionConfig(min_successful=1) + items = list(range(100)) + config = MapConfig( + max_concurrency=10, + completion_config=CompletionConfig( + min_successful=99, tolerated_failure_count=1 + ), + ) - executor = TestExecutor( - executables=executables, - max_concurrency=1, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, + executor = MapExecutor.from_items( + items=items, + func=task_func, + config=config, ) - # Create executable with COMPLETED status - exe_state = ExecutableWithState(executables[0]) - exe_state.complete("test_result") - executor.executables_with_state = [exe_state] + execution_state = Mock() + execution_state.create_checkpoint = Mock() + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + child_context = Mock() + child_context.state.wrap_user_function = lambda func, *args, **kwargs: func + executor_context.create_child_context = lambda *args, **kwargs: child_context - result = executor._create_result() # noqa: SLF001 + result = executor.execute(execution_state, executor_context) - assert len(result.all) == 1 - assert result.all[0].status == BatchItemStatus.SUCCEEDED - assert result.all[0].result == "test_result" - assert result.all[0].error is None - assert result.all[0].index == 0 + # With tolerated_failure_count=1, executor stops when failure_count > 1 (at 2 failures) + # Executor terminates early rather than executing all 100 tasks + assert executed_count["value"] < 100 + assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED, ( + executed_count + ) + assert sum(1 for item in result.all if item.status == BatchItemStatus.FAILED) == 2 + assert ( + sum(1 for item in result.all if item.status == BatchItemStatus.SUCCEEDED) < 98 + ) -def test_create_result_failed_branch(): - """Test _create_result with FAILED status branch.""" +def test_executor_exits_early_with_min_successful(): + """Test that parallel exits immediately when min_successful is reached without waiting for other branches.""" class TestExecutor(ConcurrentExecutor): def execute_item(self, child_context, executable): - return f"result_{executable.index}" - - executables = [Executable(0, lambda: "test")] - completion_config = CompletionConfig(min_successful=1) - - executor = TestExecutor( - executables=executables, - max_concurrency=1, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) - - # Create executable with FAILED status - exe_state = ExecutableWithState(executables[0]) - test_error = ValueError("Test error message") - exe_state.fail(test_error) - executor.executables_with_state = [exe_state] - - result = executor._create_result() # noqa: SLF001 + return executable.func() - assert len(result.all) == 1 - assert result.all[0].status == BatchItemStatus.FAILED - assert result.all[0].result is None - assert result.all[0].error is not None - assert result.all[0].error.message == "Test error message" - assert result.all[0].error.type == "ValueError" - assert result.all[0].index == 0 + execution_times = [] + def fast_branch(): + execution_times.append(("fast", time.time())) + return "fast_result" -def test_create_result_pending_branch(): - """Test _create_result with PENDING status branch.""" + def slow_branch(): + execution_times.append(("slow_start", time.time())) + time.sleep(2) # Long sleep + execution_times.append(("slow_end", time.time())) + return "slow_result" - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + executables = [ + Executable(0, fast_branch), + Executable(1, slow_branch), + ] - executables = [Executable(0, lambda: "test")] completion_config = CompletionConfig(min_successful=1) executor = TestExecutor( executables=executables, - max_concurrency=1, + max_concurrency=2, completion_config=completion_config, sub_type_top="TOP", sub_type_iteration="ITER", @@ -2219,109 +1923,73 @@ def execute_item(self, child_context, executable): serdes=None, ) - # Create executable with PENDING status (default state) - exe_state = ExecutableWithState(executables[0]) - # PENDING is the default state, no need to change it - executor.executables_with_state = [exe_state] - - result = executor._create_result() # noqa: SLF001 - - assert len(result.all) == 1 - assert result.all[0].status == BatchItemStatus.STARTED - assert result.all[0].result is None - assert result.all[0].error is None - assert result.all[0].index == 0 - # NEW BEHAVIOR: With min_successful=1 and no completed items, - # defaults to ALL_COMPLETED - assert result.completion_reason == CompletionReason.ALL_COMPLETED + execution_state = Mock() + execution_state.create_checkpoint = Mock() + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda idx: f"step_{idx}" # noqa: SLF001 + executor_context._parent_id = "parent" # noqa: SLF001 + def create_child_context(op_id, *, is_virtual=False): + child = Mock() + child.state = execution_state + child.state.wrap_user_function = lambda func, *args, **kwargs: func + return child -def test_create_result_running_branch(): - """Test _create_result with RUNNING status branch.""" + executor_context.create_child_context = create_child_context - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + start_time = time.time() + result = executor.execute(execution_state, executor_context) + elapsed_time = time.time() - start_time - executables = [Executable(0, lambda: "test")] - completion_config = CompletionConfig(min_successful=1) + # Should complete in less than 1.5 second (not wait for 2-second sleep) + assert elapsed_time < 1.5, f"Took {elapsed_time}s, expected < 1.5s" - executor = TestExecutor( - executables=executables, - max_concurrency=1, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) + # Result should show MIN_SUCCESSFUL_REACHED + assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - # Create executable with RUNNING status - exe_state = ExecutableWithState(executables[0]) - future = Future() - exe_state.run(future) - executor.executables_with_state = [exe_state] + # Fast branch should succeed + assert result.all[0].status == BatchItemStatus.SUCCEEDED + assert result.all[0].result == "fast_result" - result = executor._create_result() # noqa: SLF001 + # Slow branch should be marked as STARTED (incomplete) + assert result.all[1].status == BatchItemStatus.STARTED - assert len(result.all) == 1 - assert result.all[0].status == BatchItemStatus.STARTED - assert result.all[0].result is None - assert result.all[0].error is None - assert result.all[0].index == 0 - # With min_successful=1 and no completed items, defaults to ALL_COMPLETED - assert result.completion_reason == CompletionReason.ALL_COMPLETED + # Verify counts + assert result.success_count == 1 + assert result.failure_count == 0 + assert result.started_count == 1 + assert result.total_count == 2 -def test_create_result_suspended_branch(): - """Test _create_result with SUSPENDED status branch.""" +def test_executor_returns_with_incomplete_branches(): + """Test that executor returns when min_successful is reached, leaving other branches incomplete.""" class TestExecutor(ConcurrentExecutor): def execute_item(self, child_context, executable): - return f"result_{executable.index}" - - executables = [Executable(0, lambda: "test")] - completion_config = CompletionConfig(min_successful=1) - - executor = TestExecutor( - executables=executables, - max_concurrency=1, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) - - # Create executable with SUSPENDED status - exe_state = ExecutableWithState(executables[0]) - exe_state.suspend() - executor.executables_with_state = [exe_state] - - result = executor._create_result() # noqa: SLF001 + return executable.func() - assert len(result.all) == 1 - assert result.all[0].status == BatchItemStatus.STARTED - assert result.all[0].result is None - assert result.all[0].error is None - assert result.all[0].index == 0 - # With min_successful=1 and no completed items, defaults to ALL_COMPLETED - assert result.completion_reason == CompletionReason.ALL_COMPLETED + operation_tracker = Mock() + def fast_branch(): + operation_tracker.fast_executed() + return "fast_result" -def test_create_result_suspended_with_timeout_branch(): - """Test _create_result with SUSPENDED_WITH_TIMEOUT status branch.""" + def slow_branch(): + operation_tracker.slow_started() + time.sleep(2) # Long sleep + operation_tracker.slow_completed() + return "slow_result" - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + executables = [ + Executable(0, fast_branch), + Executable(1, slow_branch), + ] - executables = [Executable(0, lambda: "test")] completion_config = CompletionConfig(min_successful=1) executor = TestExecutor( executables=executables, - max_concurrency=1, + max_concurrency=2, completion_config=completion_config, sub_type_top="TOP", sub_type_iteration="ITER", @@ -2329,43 +1997,60 @@ def execute_item(self, child_context, executable): serdes=None, ) - # Create executable with SUSPENDED_WITH_TIMEOUT status - exe_state = ExecutableWithState(executables[0]) - future_time = time.time() + 10 - exe_state.suspend_with_timeout(future_time) - executor.executables_with_state = [exe_state] + execution_state = Mock() + execution_state.create_checkpoint = Mock() + execution_state.wrap_user_function = lambda func, *args, **kwargs: func + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda idx: f"step_{idx}" # noqa: SLF001 + executor_context._parent_id = "parent" # noqa: SLF001 + executor_context.create_child_context = lambda op_id, *, is_virtual=False: Mock( + state=execution_state + ) - result = executor._create_result() # noqa: SLF001 + result = executor.execute(execution_state, executor_context) - assert len(result.all) == 1 - assert result.all[0].status == BatchItemStatus.STARTED - assert result.all[0].result is None - assert result.all[0].error is None - assert result.all[0].index == 0 - # With min_successful=1 and no completed items, default to ALL_COMPLETED - assert result.completion_reason == CompletionReason.ALL_COMPLETED + # Verify fast branch executed + assert operation_tracker.fast_executed.call_count == 1 + + # Slow branch may or may not have started (depends on thread scheduling) + # but it definitely should not have completed + assert operation_tracker.slow_completed.call_count == 0, ( + "Executor should return before slow branch completes" + ) + + # Result should show MIN_SUCCESSFUL_REACHED + assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED + + # Verify counts - one succeeded, one incomplete + assert result.success_count == 1 + assert result.failure_count == 0 + assert result.started_count == 1 + assert result.total_count == 2 -def test_create_result_mixed_statuses(): - """Test _create_result with mixed executable statuses covering all branches.""" +def test_executor_returns_before_slow_branch_completes(): + """Test that executor returns immediately when min_successful is reached, not waiting for slow branches.""" class TestExecutor(ConcurrentExecutor): def execute_item(self, child_context, executable): - return f"result_{executable.index}" + return executable.func() - executables = [ - Executable(0, lambda: "test0"), # Will be COMPLETED - Executable(1, lambda: "test1"), # Will be FAILED - Executable(2, lambda: "test2"), # Will be PENDING - Executable(3, lambda: "test3"), # Will be RUNNING - Executable(4, lambda: "test4"), # Will be SUSPENDED - Executable(5, lambda: "test5"), # Will be SUSPENDED_WITH_TIMEOUT - ] + slow_branch_mock = Mock() + + def fast_func(): + return "fast" + + def slow_func(): + time.sleep(3) # Sleep + slow_branch_mock.completed() # Should not be called before executor returns + return "slow" + + executables = [Executable(0, fast_func), Executable(1, slow_func)] completion_config = CompletionConfig(min_successful=1) executor = TestExecutor( executables=executables, - max_concurrency=6, + max_concurrency=2, completion_config=completion_config, sub_type_top="TOP", sub_type_iteration="ITER", @@ -2373,922 +2058,582 @@ def execute_item(self, child_context, executable): serdes=None, ) - # Create executables with different statuses - exe_states = [ExecutableWithState(exe) for exe in executables] + execution_state = Mock() + execution_state.create_checkpoint = Mock() + execution_state.wrap_user_function = lambda func, *args, **kwargs: func + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda idx: f"step_{idx}" # noqa: SLF001 + executor_context._parent_id = "parent" # noqa: SLF001 + executor_context.create_child_context = lambda op_id, *, is_virtual=False: Mock( + state=execution_state + ) - # COMPLETED - exe_states[0].complete("completed_result") + result = executor.execute(execution_state, executor_context) - # FAILED - exe_states[1].fail(RuntimeError("Test failure")) + # Executor should have returned before slow branch completed + assert not slow_branch_mock.completed.called, ( + "Executor should return before slow branch completes" + ) - # PENDING (default state, no change needed) - - # RUNNING - future = Future() - exe_states[3].run(future) - - # SUSPENDED - exe_states[4].suspend() - - # SUSPENDED_WITH_TIMEOUT - exe_states[5].suspend_with_timeout(time.time() + 10) - - executor.executables_with_state = exe_states - - result = executor._create_result() # noqa: SLF001 - - assert len(result.all) == 6 - - # Check COMPLETED -> SUCCEEDED - assert result.all[0].status == BatchItemStatus.SUCCEEDED - assert result.all[0].result == "completed_result" - assert result.all[0].error is None - - # Check FAILED -> FAILED - assert result.all[1].status == BatchItemStatus.FAILED - assert result.all[1].result is None - assert result.all[1].error is not None - assert result.all[1].error.message == "Test failure" - - # Check PENDING -> STARTED - assert result.all[2].status == BatchItemStatus.STARTED - assert result.all[2].result is None - assert result.all[2].error is None - - # Check RUNNING -> STARTED - assert result.all[3].status == BatchItemStatus.STARTED - assert result.all[3].result is None - assert result.all[3].error is None - - # Check SUSPENDED -> STARTED - assert result.all[4].status == BatchItemStatus.STARTED - assert result.all[4].result is None - assert result.all[4].error is None - - # Check SUSPENDED_WITH_TIMEOUT -> STARTED - assert result.all[5].status == BatchItemStatus.STARTED - assert result.all[5].result is None - assert result.all[5].error is None - - # we've a min succ set to 1. + # Result should show MIN_SUCCESSFUL_REACHED assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED + # Verify counts + assert result.success_count == 1 + assert result.failure_count == 0 + assert result.started_count == 1 + assert result.total_count == 2 -def test_create_result_multiple_completed(): - """Test _create_result with multiple COMPLETED executables.""" - - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" - executables = [ - Executable(0, lambda: "test0"), - Executable(1, lambda: "test1"), - Executable(2, lambda: "test2"), +def test_from_items_no_config_with_failures(): + """Validates: Requirements 2.4 - Fail-fast with no config.""" + items = [ + BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), + BatchItem( + 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ), ] - completion_config = CompletionConfig(min_successful=3) - - executor = TestExecutor( - executables=executables, - max_concurrency=3, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) - - # Create all executables with COMPLETED status - exe_states = [ExecutableWithState(exe) for exe in executables] - exe_states[0].complete("result_0") - exe_states[1].complete("result_1") - exe_states[2].complete("result_2") - - executor.executables_with_state = exe_states + result = BatchResult.from_items(items, completion_config=None) + assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED - result = executor._create_result() # noqa: SLF001 - assert len(result.all) == 3 - assert all(item.status == BatchItemStatus.SUCCEEDED for item in result.all) - assert result.all[0].result == "result_0" - assert result.all[1].result == "result_1" - assert result.all[2].result == "result_2" - assert result.completion_reason == CompletionReason.ALL_COMPLETED +def test_from_items_empty_config_with_failures(): + """Validates: Requirements 2.5 - Fail-fast with empty config.""" + items = [ + BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), + BatchItem( + 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ), + ] + config = CompletionConfig() # All fields None + result = BatchResult.from_items(items, completion_config=config) + assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED -def test_create_result_multiple_failed(): - """Test _create_result with multiple FAILED executables.""" +def test_from_items_tolerance_checked_before_all_completed(): + """Validates: Requirements 2.1, 2.2 - Tolerance priority.""" + items = [ + BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), + BatchItem( + 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ), + BatchItem( + 2, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ), + ] + config = CompletionConfig(tolerated_failure_count=1) + result = BatchResult.from_items(items, completion_config=config) + # All completed but tolerance exceeded - should return TOLERANCE_EXCEEDED + assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" - executables = [ - Executable(0, lambda: "test0"), - Executable(1, lambda: "test1"), - Executable(2, lambda: "test2"), +def test_from_items_all_completed_within_tolerance(): + """Validates: Requirements 1.1 - All completed.""" + items = [ + BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), + BatchItem( + 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ), ] - completion_config = CompletionConfig(min_successful=1) - - executor = TestExecutor( - executables=executables, - max_concurrency=3, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) + config = CompletionConfig(tolerated_failure_count=1) + result = BatchResult.from_items(items, completion_config=config) + assert result.completion_reason == CompletionReason.ALL_COMPLETED - # Create all executables with FAILED status - exe_states = [ExecutableWithState(exe) for exe in executables] - exe_states[0].fail(ValueError("Error 0")) - exe_states[1].fail(RuntimeError("Error 1")) - exe_states[2].fail(TypeError("Error 2")) - executor.executables_with_state = exe_states +def test_from_items_min_successful_reached(): + """Validates: Requirements 1.3 - Min successful.""" + items = [ + BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), + BatchItem(1, BatchItemStatus.SUCCEEDED, result="ok"), + BatchItem(2, BatchItemStatus.STARTED), + ] + config = CompletionConfig(min_successful=2) + result = BatchResult.from_items(items, completion_config=config) + assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - result = executor._create_result() # noqa: SLF001 - assert len(result.all) == 3 - assert all(item.status == BatchItemStatus.FAILED for item in result.all) - assert result.all[0].error.message == "Error 0" - assert result.all[1].error.message == "Error 1" - assert result.all[2].error.message == "Error 2" - assert result.completion_reason == CompletionReason.ALL_COMPLETED +def test_from_items_tolerance_count_exceeded(): + """Validates: Requirements 1.2 - Tolerance count.""" + items = [ + BatchItem( + 0, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ), + BatchItem( + 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ), + BatchItem(2, BatchItemStatus.STARTED), + ] + config = CompletionConfig(tolerated_failure_count=1) + result = BatchResult.from_items(items, completion_config=config) + assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED -def test_create_result_multiple_started_states(): - """Test _create_result with multiple executables in STARTED states.""" +def test_from_items_tolerance_percentage_exceeded(): + """Validates: Requirements 1.2 - Tolerance percentage.""" + items = [ + BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), + BatchItem( + 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ), + BatchItem( + 2, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ), + BatchItem( + 3, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ), + ] + config = CompletionConfig(tolerated_failure_percentage=50.0) + # 3 failures out of 4 = 75% > 50% + result = BatchResult.from_items(items, completion_config=config) + assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" - executables = [ - Executable(0, lambda: "test0"), # PENDING - Executable(1, lambda: "test1"), # RUNNING - Executable(2, lambda: "test2"), # SUSPENDED - Executable(3, lambda: "test3"), # SUSPENDED_WITH_TIMEOUT +def test_from_items_tolerance_priority_over_min_successful(): + """Validates: Requirements 2.3 - Tolerance takes precedence.""" + items = [ + BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), + BatchItem(1, BatchItemStatus.SUCCEEDED, result="ok"), + BatchItem( + 2, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ), + BatchItem( + 3, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + ), ] - completion_config = CompletionConfig(min_successful=1) + config = CompletionConfig(min_successful=2, tolerated_failure_count=1) + # Min successful reached (2) but tolerance exceeded (2 > 1) + result = BatchResult.from_items(items, completion_config=config) + assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED - executor = TestExecutor( - executables=executables, - max_concurrency=4, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) - # Create executables with different STARTED states - exe_states = [ExecutableWithState(exe) for exe in executables] +def test_from_items_empty_array(): + """Validates: Edge case - empty items.""" + items = [] + result = BatchResult.from_items(items, completion_config=None) + assert result.completion_reason == CompletionReason.ALL_COMPLETED + assert result.total_count == 0 - # PENDING (default state) - # RUNNING - future = Future() - exe_states[1].run(future) +def test_from_items_all_succeeded(): + """Validates: All items succeeded.""" + items = [ + BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok1"), + BatchItem(1, BatchItemStatus.SUCCEEDED, result="ok2"), + ] + result = BatchResult.from_items(items, completion_config=None) + assert result.completion_reason == CompletionReason.ALL_COMPLETED + assert result.success_count == 2 - # SUSPENDED - exe_states[2].suspend() - # SUSPENDED_WITH_TIMEOUT - exe_states[3].suspend_with_timeout(time.time() + 5) +# endregion Completion Reason Inference Tests - executor.executables_with_state = exe_states +# region Virtual-context wire-format tests - result = executor._create_result() # noqa: SLF001 - assert len(result.all) == 4 - assert all(item.status == BatchItemStatus.STARTED for item in result.all) - assert all(item.result is None for item in result.all) - assert all(item.error is None for item in result.all) - # With min_successful=1 and no completed items, defaults to ALL_COMPLETED - assert result.completion_reason == CompletionReason.ALL_COMPLETED +def test_flat_mode_stamps_grandparent_as_inner_op_parent_id(): + """In FLAT mode, inner operations in a branch stamp the map/parallel op id as parent_id. + This is the core FLAT-mode invariant. Inner operations must not + stamp the branch's own operation id (that would reproduce the + NESTED hierarchy) — they must stamp the enclosing map/parallel op + id, so the branch is collapsed out of the observable hierarchy + even though it still exists as a logical scope for concurrency + and step-id prefixing. -def test_create_result_empty_executables(): - """Test _create_result with no executables.""" + The test drives `_execute_item_in_child_context` with a real + non-virtual executor context, captures the child context the + executor builds for the branch, and asserts the branch's + `_parent_id` equals the executor_context's own `_parent_id`. + """ class TestExecutor(ConcurrentExecutor): def execute_item(self, child_context, executable): - return f"result_{executable.index}" + # Record the child context we receive so the assertions below can + # inspect its identity fields. + self.last_child_context = child_context + return executable.func(child_context) + + execution_state = Mock() + execution_state.create_checkpoint = Mock() + execution_state.wrap_user_function = lambda func, *args, **kwargs: func + + # Mock out the checkpoint so the real child_handler reports "not + # existent" (non-existent checkpoint -> normal execution path). + mock_checkpoint = Mock() + mock_checkpoint.is_succeeded.return_value = False + mock_checkpoint.is_failed.return_value = False + mock_checkpoint.is_existent.return_value = False + mock_checkpoint.is_replay_children.return_value = False + execution_state.get_checkpoint_result.return_value = mock_checkpoint - executables = [] - completion_config = CompletionConfig(min_successful=0) + # Build a real DurableContext that represents the map/parallel op. + map_op_id = "map-op-id" + execution_context = ExecutionContext( + durable_execution_arn="arn:aws:durable:us-east-1:0:execution/test" + ) + executor_context = DurableContext( + state=execution_state, + execution_context=execution_context, + parent_id=map_op_id, # This context *is* the map/parallel op. + ) + executables = [Executable(index=0, func=lambda ctx: "ok")] executor = TestExecutor( executables=executables, max_concurrency=1, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", + completion_config=CompletionConfig(min_successful=1), + sub_type_top="MAP", + sub_type_iteration="MAP_ITER", + name_prefix="branch-", serdes=None, + nesting_type=NestingType.FLAT, ) - executor.executables_with_state = [] - - result = executor._create_result() # noqa: SLF001 - - assert len(result.all) == 0 - assert result.completion_reason == CompletionReason.ALL_COMPLETED - - -def test_timer_scheduler_future_time_condition_false(): - """Test TimerScheduler when scheduled time is in future (434->433 branch).""" - callback = Mock() + executor._execute_item_in_child_context(executor_context, executables[0]) # noqa: SLF001 - with TimerScheduler(callback) as scheduler: - exe_state = ExecutableWithState(Executable(0, lambda: "test")) - exe_state.suspend() + # The branch's child context must be virtual AND propagate the + # map/parallel op id as its _parent_id. Inner operations stamping + # self._parent_id will therefore report to the map/parallel op. + branch_ctx = executor.last_child_context + assert branch_ctx.is_virtual is True + assert branch_ctx._parent_id == map_op_id # noqa: SLF001 + # The step-id prefix is the branch's own operation id (stable replay id). + assert branch_ctx._step_id_prefix != map_op_id # noqa: SLF001 - # Schedule with future time so condition will be False - future_time = time.time() + 10 - scheduler.schedule_resume(exe_state, future_time) - # Wait briefly for timer thread to check and find condition False - time.sleep(0.1) +def test_nested_mode_stamps_branch_op_as_inner_op_parent_id(): + """In NESTED mode, inner operations in a branch stamp the branch's own operation id as parent_id.""" - # Callback should not be called since time is in future - callback.assert_not_called() + class TestExecutor(ConcurrentExecutor): + def execute_item(self, child_context, executable): + self.last_child_context = child_context + return executable.func(child_context) + execution_state = Mock() + execution_state.create_checkpoint = Mock() + execution_state.wrap_user_function = lambda func, *args, **kwargs: func -def test_batch_result_from_dict_with_completion_config(): - """Test BatchResult from_dict with completion config parameter.""" - data = { - "all": [ - {"index": 0, "status": "SUCCEEDED", "result": "result1", "error": None}, - {"index": 1, "status": "STARTED", "result": None, "error": None}, - ], - # No completionReason provided - } - - # With started items, should infer MIN_SUCCESSFUL_REACHED - completion_config = CompletionConfig(min_successful=1) - - with patch( - "aws_durable_execution_sdk_python.concurrency.models.logger" - ) as mock_logger: - result = BatchResult.from_dict(data, completion_config) - assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - mock_logger.warning.assert_called_once() - - -def test_batch_result_from_dict_all_completed(): - """Test BatchResult from_dict infers completion reason when all items are completed.""" - data = { - "all": [ - {"index": 0, "status": "SUCCEEDED", "result": "result1", "error": None}, - { - "index": 1, - "status": "FAILED", - "result": None, - "error": { - "message": "error", - "type": "Error", - "data": None, - "stackTrace": None, - }, - }, - ], - # No completionReason provided - } - - # With no config and failures, fail-fast - with patch( - "aws_durable_execution_sdk_python.concurrency.models.logger" - ) as mock_logger: - result = BatchResult.from_dict(data) - assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED - mock_logger.warning.assert_called_once() - - -def test_batch_result_from_dict_backward_compatibility(): - """Test BatchResult from_dict maintains backward compatibility when no completion_config provided.""" - data = { - "all": [ - {"index": 0, "status": "SUCCEEDED", "result": "result1", "error": None} - ], - "completionReason": "MIN_SUCCESSFUL_REACHED", - } + mock_checkpoint = Mock() + mock_checkpoint.is_succeeded.return_value = False + mock_checkpoint.is_failed.return_value = False + mock_checkpoint.is_existent.return_value = False + mock_checkpoint.is_replay_children.return_value = False + execution_state.get_checkpoint_result.return_value = mock_checkpoint - # Should work without completion_config parameter - result = BatchResult.from_dict(data) - assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED + map_op_id = "map-op-id" + execution_context = ExecutionContext( + durable_execution_arn="arn:aws:durable:us-east-1:0:execution/test" + ) + executor_context = DurableContext( + state=execution_state, + execution_context=execution_context, + parent_id=map_op_id, + ) - # Should also work with None completion_config - result2 = BatchResult.from_dict(data, None) - assert result2.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED + executables = [Executable(index=0, func=lambda ctx: "ok")] + executor = TestExecutor( + executables=executables, + max_concurrency=1, + completion_config=CompletionConfig(min_successful=1), + sub_type_top="MAP", + sub_type_iteration="MAP_ITER", + name_prefix="branch-", + serdes=None, + nesting_type=NestingType.NESTED, + ) + executor._execute_item_in_child_context(executor_context, executables[0]) # noqa: SLF001 -def test_batch_result_infer_completion_reason_basic_cases(): - """Test _infer_completion_reason method with basic scenarios.""" - # Test with started items - should be MIN_SUCCESSFUL_REACHED - items = { - "all": [ - BatchItem(0, BatchItemStatus.SUCCEEDED, "result1").to_dict(), - BatchItem(1, BatchItemStatus.STARTED).to_dict(), - ] - } - batch = BatchResult.from_dict(items, CompletionConfig(1)) - assert batch.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED + # In NESTED mode, the branch is a regular child — its _parent_id is + # its own operation id, not the grandparent. + branch_ctx = executor.last_child_context + assert branch_ctx.is_virtual is False + assert branch_ctx._parent_id == branch_ctx._step_id_prefix # noqa: SLF001 + assert branch_ctx._parent_id != map_op_id # noqa: SLF001 - # Test with all completed items - should be ALL_COMPLETED - completed_items = [ - BatchItem(0, BatchItemStatus.SUCCEEDED, "result1").to_dict(), - BatchItem( - 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) - ).to_dict(), - ] - completed_items = {"all": completed_items} - batch = BatchResult.from_dict(completed_items, CompletionConfig(1)) - assert batch.completion_reason == CompletionReason.ALL_COMPLETED - # Test empty items - should be ALL_COMPLETED - batch = BatchResult.from_dict({"all": []}, CompletionConfig(1)) - assert batch.completion_reason == CompletionReason.ALL_COMPLETED +# endregion Virtual-context wire-format tests -def test_operation_id_determinism_across_shuffles(): - """Test that operation_id depends on Executable.index, not execution order.""" +def test_flat_mode_produces_deterministic_step_ids_across_runs(): + """Step ids and inner parent_ids must be deterministic under FLAT mode. - def index_based_function(index, ctx): - """Function that returns a result based on the executable index.""" - return f"result_for_index_{index}" + Replay depends on regenerating the same operation ids for the same + logical inputs. This test runs the same executor twice against + fresh executor contexts and asserts that the resulting set of + (step_id_prefix, parent_id) pairs is identical. Any source of + non-determinism (e.g. step prefixes that depend on thread + completion order, or parent-id propagation that's different on + the second run) would show up here as a mismatch and would cause + `NonDeterministicExecutionException` at replay time in production. + """ class TestExecutor(ConcurrentExecutor): - """Custom executor for testing operation_id determinism.""" + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.captured = [] def execute_item(self, child_context, executable): + self.captured.append( + ( + child_context._step_id_prefix, # noqa: SLF001 + child_context._parent_id, # noqa: SLF001 + ) + ) return executable.func(child_context) - # Create executables with specific indices using partial - num_executables = 50 - funcs = [partial(index_based_function, i) for i in range(num_executables)] - - # Track operation_id -> result associations - captured_associations = [] - - def patched_child_handler( - func, - execution_state, - operation_identifier, - config: ChildConfig, - ): - """Patched child handler that captures operation_id -> result mapping.""" - assert config.is_virtual - assert config.sub_type == "TEST_ITER" - result = func() # Execute the function - captured_associations.append((operation_identifier.operation_id, result)) - return result - - execution_state = Mock() - execution_state.create_checkpoint = Mock() - - completion_config = CompletionConfig(min_successful=num_executables) - - # Run multiple times with different shuffle orders - associations_per_run = [] + def make_run(): + execution_state = Mock() + execution_state.create_checkpoint = Mock() - for run in range(10): # Test 10 different shuffle orders - captured_associations.clear() + mock_checkpoint = Mock() + mock_checkpoint.is_succeeded.return_value = False + mock_checkpoint.is_failed.return_value = False + mock_checkpoint.is_existent.return_value = False + mock_checkpoint.is_replay_children.return_value = False + execution_state.get_checkpoint_result.return_value = mock_checkpoint - # Create executables from shuffled functions - executables = [Executable(index=i, func=func) for i, func in enumerate(funcs)] - random.seed(run) # Different seed for each run - random.shuffle(executables) + execution_context = ExecutionContext( + durable_execution_arn="arn:aws:durable:us-east-1:0:execution/test" + ) + executor_context = DurableContext( + state=execution_state, + execution_context=execution_context, + parent_id="map-op-id", + ) + executables = [ + Executable(index=i, func=lambda ctx, i=i: f"r{i}") for i in range(3) + ] executor = TestExecutor( executables=executables, - max_concurrency=2, - completion_config=completion_config, - sub_type_top="TEST", - sub_type_iteration="TEST_ITER", - name_prefix="test_", + max_concurrency=3, + completion_config=CompletionConfig(min_successful=3), + sub_type_top="MAP", + sub_type_iteration="MAP_ITER", + name_prefix="branch-", serdes=None, nesting_type=NestingType.FLAT, ) + executor.execute(execution_state, executor_context) + return executor.captured - # Create executor context mock - executor_context = Mock() - executor_context._parent_id = "parent_123" # noqa SLF001 + run_a = make_run() + run_b = make_run() - def create_step_id(index): - return f"step_{index}" + # Ordering of captured items is non-deterministic because branches run + # on a ThreadPoolExecutor. What matters is that the SET of (prefix, + # parent_id) pairs is identical across runs — i.e. replay reconstructs + # the same branch identity regardless of completion order. + assert sorted(run_a) == sorted(run_b), ( + "FLAT-mode branch step-id prefixes and parent_ids must be identical " + f"across runs. Run A: {run_a!r}; Run B: {run_b!r}" + ) + # Sanity: all branches reported grandparent (map op id) as their parent. + assert all(parent_id == "map-op-id" for _prefix, parent_id in run_a) - executor_context._create_step_id_for_logical_step = create_step_id # noqa SLF001 - def create_child_context(operation_id, *, is_virtual=False): - child_ctx = Mock() - child_ctx.state = execution_state - return child_ctx +# region Branch state transitions - executor_context.create_child_context = create_child_context - with patch( - "aws_durable_execution_sdk_python.concurrency.executor.child_handler", - patched_child_handler, - ): - executor.execute(execution_state, executor_context) +def test_branch_initial_state(): + branch = Branch(Executable(3, lambda: "test")) + assert branch.status is BranchStatus.PENDING + assert branch.index == 3 + assert branch.result is None + assert branch.error is None + assert branch.resume_at is None - associations_per_run.append(captured_associations.copy()) - # first we will verify the validity of the test by ensuring that there exist at least 2 runs with different ordering - assert any( - assoc1 != assoc2 for assoc1, assoc2 in combinations(associations_per_run, 2) - ) - # then we will verify the invariant of association between step_id and result - associations_per_run = [dict(assoc) for assoc in associations_per_run] - assert all( - assoc1 == assoc2 for assoc1, assoc2 in combinations(associations_per_run, 2) - ) +def test_branch_start_clears_resume_at(): + branch = Branch(Executable(0, lambda: "test")) + branch.suspend_until(123.0) + assert branch.status is BranchStatus.SUSPENDED_WITH_TIMEOUT + assert branch.resume_at == 123.0 + branch.start() + assert branch.status is BranchStatus.RUNNING + assert branch.resume_at is None -def test_concurrent_executor_replay_with_succeeded_operations(): - """Test ConcurrentExecutor replay method with succeeded operations.""" +def test_branch_complete_stores_result(): + branch = Branch(Executable(0, lambda: "test")) + branch.start() + branch.complete("result") + assert branch.status is BranchStatus.COMPLETED + assert branch.result == "result" - def func1(ctx, item, idx, items): - return f"result_{item}" - items = ["a", "b"] - config = MapConfig() - - executor = MapExecutor.from_items( - items=items, - func=func1, - config=config, - ) - - # Mock execution state with succeeded operations - mock_execution_state = Mock() - mock_execution_state.durable_execution_arn = ( - "arn:aws:durable:us-east-1:123456789012:execution/test" - ) - - def mock_get_checkpoint_result(operation_id): - mock_result = Mock() - mock_result.is_succeeded.return_value = True - mock_result.is_failed.return_value = False - mock_result.is_replay_children.return_value = False - mock_result.is_existent.return_value = True - # Provide properly serialized JSON data - mock_result.result = f'"cached_result_{operation_id}"' # JSON string - return mock_result - - mock_execution_state.get_checkpoint_result = mock_get_checkpoint_result - - def mock_create_step_id_for_logical_step(step): - return f"op_{step}" - - # Mock executor context - mock_executor_context = Mock() - mock_executor_context._create_step_id_for_logical_step = ( # noqa - mock_create_step_id_for_logical_step - ) - - # Mock child context that has the same execution state - mock_child_context = Mock() - mock_child_context.state = mock_execution_state - mock_executor_context.create_child_context = Mock(return_value=mock_child_context) - mock_executor_context._parent_id = "parent_id" # noqa - - result = executor.replay(mock_execution_state, mock_executor_context) - - assert isinstance(result, BatchResult) - assert len(result.all) == 2 - assert result.all[0].status == BatchItemStatus.SUCCEEDED - assert result.all[0].result == "cached_result_op_0" - assert result.all[1].status == BatchItemStatus.SUCCEEDED - assert result.all[1].result == "cached_result_op_1" - - -def test_concurrent_executor_replay_with_failed_operations(): - """Test ConcurrentExecutor replay method with failed operations.""" - - def func1(ctx, item, idx, items): - return f"result_{item}" - - items = ["a"] - config = MapConfig() - - executor = MapExecutor.from_items( - items=items, - func=func1, - config=config, - ) - - # Mock execution state with failed operation - mock_execution_state = Mock() - - def mock_get_checkpoint_result(operation_id): - mock_result = Mock() - mock_result.is_succeeded.return_value = False - mock_result.is_failed.return_value = True - mock_result.error = Exception("Test error") - return mock_result - - mock_execution_state.get_checkpoint_result = mock_get_checkpoint_result - - # Mock executor context - mock_executor_context = Mock() - mock_executor_context._create_step_id_for_logical_step = Mock(return_value="op_1") # noqa: SLF001 - - result = executor.replay(mock_execution_state, mock_executor_context) - - assert isinstance(result, BatchResult) - assert len(result.all) == 1 - assert result.all[0].status == BatchItemStatus.FAILED - assert result.all[0].error is not None - - -def test_concurrent_executor_replay_with_replay_children(): - """Test ConcurrentExecutor replay method when children need re-execution.""" - - def func1(ctx, item, idx, items): - return f"result_{item}" - - items = ["a"] - config = MapConfig() - - executor = MapExecutor.from_items( - items=items, - func=func1, - config=config, - ) - - # Mock execution state with succeeded operation that needs replay - mock_execution_state = Mock() - - def mock_get_checkpoint_result(operation_id): - mock_result = Mock() - mock_result.is_succeeded.return_value = True - mock_result.is_failed.return_value = False - mock_result.is_replay_children.return_value = True - return mock_result - - mock_execution_state.get_checkpoint_result = mock_get_checkpoint_result - - # Mock executor context - mock_executor_context = Mock() - mock_executor_context._create_step_id_for_logical_step = Mock(return_value="op_1") # noqa: SLF001 - - # Mock _execute_item_in_child_context to return a result - with patch.object( - executor, "_execute_item_in_child_context", return_value="re_executed_result" - ): - result = executor.replay(mock_execution_state, mock_executor_context) - - assert isinstance(result, BatchResult) - assert len(result.all) == 1 - assert result.all[0].status == BatchItemStatus.SUCCEEDED - assert result.all[0].result == "re_executed_result" - - -def test_batch_item_from_dict_with_error(): - """Test BatchItem.from_dict() with error.""" - data = { - "index": 3, - "status": "FAILED", - "result": None, - "error": { - "ErrorType": "ValueError", - "ErrorMessage": "bad value", - "StackTrace": [], - }, - } - - item = BatchItem.from_dict(data) - - assert item.index == 3 - assert item.status == BatchItemStatus.FAILED - assert item.error.type == "ValueError" - assert item.error.message == "bad value" - - -def test_batch_result_with_mixed_statuses(): - """Test BatchResult serialization with mixed item statuses.""" - result = BatchResult( - all=[ - BatchItem(0, BatchItemStatus.SUCCEEDED, result="success"), - BatchItem( - 1, - BatchItemStatus.FAILED, - error=ErrorObject(message="msg", type="E", data=None, stack_trace=[]), - ), - BatchItem(2, BatchItemStatus.STARTED), - ], - completion_reason=CompletionReason.FAILURE_TOLERANCE_EXCEEDED, - ) - - serialized = json.dumps(result.to_dict()) - deserialized = BatchResult.from_dict(json.loads(serialized)) - - assert len(deserialized.all) == 3 - assert deserialized.all[0].status == BatchItemStatus.SUCCEEDED - assert deserialized.all[1].status == BatchItemStatus.FAILED - assert deserialized.all[2].status == BatchItemStatus.STARTED - assert deserialized.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED - - -def test_batch_result_empty_list(): - """Test BatchResult serialization with empty items list.""" - result = BatchResult(all=[], completion_reason=CompletionReason.ALL_COMPLETED) - - serialized = json.dumps(result.to_dict()) - deserialized = BatchResult.from_dict(json.loads(serialized)) - - assert len(deserialized.all) == 0 - assert deserialized.completion_reason == CompletionReason.ALL_COMPLETED - - -def test_batch_result_complex_nested_data(): - """Test BatchResult with complex nested data structures.""" - complex_result = { - "users": [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}], - "metadata": {"count": 2, "timestamp": "2025-10-31"}, - } - - result = BatchResult( - all=[BatchItem(0, BatchItemStatus.SUCCEEDED, result=complex_result)], - completion_reason=CompletionReason.ALL_COMPLETED, - ) - - serialized = json.dumps(result.to_dict()) - deserialized = BatchResult.from_dict(json.loads(serialized)) - - assert deserialized.all[0].result == complex_result - assert deserialized.all[0].result["users"][0]["name"] == "Alice" - - -def test_executor_does_not_deadlock_when_all_tasks_terminal_but_completion_config_allows_failures(): - """Ensure executor returns when all tasks are terminal even if completion rules are confusing.""" - - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - if executable.index == 0: - # fail one task - raise Exception("boom") # noqa EM101 TRY002 - return f"ok_{executable.index}" - - # Two tasks, min_successful=2 but tolerated failure_count set to 1. - # After one fail + one success, counters.is_complete() should return true, - # should_continue() should return false. counters.is_complete was failing to - # stop early, which caused map to hang. - executables = [Executable(0, lambda: "a"), Executable(1, lambda: "b")] - completion_config = CompletionConfig( - min_successful=2, - tolerated_failure_count=1, - tolerated_failure_percentage=None, - ) - - executor = TestExecutor( - executables=executables, - max_concurrency=2, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) - - execution_state = Mock() - execution_state.create_checkpoint = Mock() - executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 - child_context = Mock() - child_context.state.wrap_user_function = lambda func, *args, **kwargs: func - executor_context.create_child_context = lambda *args, **kwargs: child_context - - # Should return (not hang) and batch should reflect one FAILED and one SUCCEEDED - result = executor.execute(execution_state, executor_context) - statuses = {item.index: item.status for item in result.all} - assert statuses[0] == BatchItemStatus.FAILED - assert statuses[1] == BatchItemStatus.SUCCEEDED - - -def test_executor_terminates_quickly_when_impossible_to_succeed(): - """Test that executor terminates when min_successful becomes impossible.""" - executed_count = {"value": 0} - - def task_func(ctx, item, idx, items): - executed_count["value"] += 1 - if idx < 2: - raise Exception(f"fail_{idx}") # noqa EM102 TRY002 - time.sleep(0.05) - return f"ok_{idx}" - - items = list(range(100)) - config = MapConfig( - max_concurrency=10, - completion_config=CompletionConfig( - min_successful=99, tolerated_failure_count=1 - ), - ) - - executor = MapExecutor.from_items( - items=items, - func=task_func, - config=config, - ) - - execution_state = Mock() - execution_state.create_checkpoint = Mock() - executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 - child_context = Mock() - child_context.state.wrap_user_function = lambda func, *args, **kwargs: func - executor_context.create_child_context = lambda *args, **kwargs: child_context - - result = executor.execute(execution_state, executor_context) - - # With tolerated_failure_count=1, executor stops when failure_count > 1 (at 2 failures) - # Executor terminates early rather than executing all 100 tasks - assert executed_count["value"] < 100 - assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED, ( - executed_count - ) - assert sum(1 for item in result.all if item.status == BatchItemStatus.FAILED) == 2 - assert ( - sum(1 for item in result.all if item.status == BatchItemStatus.SUCCEEDED) < 98 - ) - - -def test_executor_exits_early_with_min_successful(): - """Test that parallel exits immediately when min_successful is reached without waiting for other branches.""" - - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return executable.func() - - execution_times = [] - - def fast_branch(): - execution_times.append(("fast", time.time())) - return "fast_result" - - def slow_branch(): - execution_times.append(("slow_start", time.time())) - time.sleep(2) # Long sleep - execution_times.append(("slow_end", time.time())) - return "slow_result" - - executables = [ - Executable(0, fast_branch), - Executable(1, slow_branch), - ] - - completion_config = CompletionConfig(min_successful=1) - - executor = TestExecutor( - executables=executables, - max_concurrency=2, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) - - execution_state = Mock() - execution_state.create_checkpoint = Mock() - executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda idx: f"step_{idx}" # noqa: SLF001 - executor_context._parent_id = "parent" # noqa: SLF001 - - def create_child_context(op_id, *, is_virtual=False): - child = Mock() - child.state = execution_state - child.state.wrap_user_function = lambda func, *args, **kwargs: func - return child +def test_branch_fail_stores_error(): + branch = Branch(Executable(0, lambda: "test")) + branch.start() + error = ValueError("boom") + branch.fail(error) + assert branch.status is BranchStatus.FAILED + assert branch.error is error - executor_context.create_child_context = create_child_context - start_time = time.time() - result = executor.execute(execution_state, executor_context) - elapsed_time = time.time() - start_time +def test_branch_suspend_not_terminal(): + branch = Branch(Executable(0, lambda: "test")) + branch.start() + branch.suspend() + assert branch.status is BranchStatus.SUSPENDED + assert branch.resume_at is None - # Should complete in less than 1.5 second (not wait for 2-second sleep) - assert elapsed_time < 1.5, f"Took {elapsed_time}s, expected < 1.5s" - # Result should show MIN_SUCCESSFUL_REACHED - assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED +def test_branch_suspend_until_not_terminal(): + branch = Branch(Executable(0, lambda: "test")) + branch.start() + branch.suspend_until(456.0) + assert branch.status is BranchStatus.SUSPENDED_WITH_TIMEOUT + assert branch.resume_at == 456.0 - # Fast branch should succeed - assert result.all[0].status == BatchItemStatus.SUCCEEDED - assert result.all[0].result == "fast_result" - # Slow branch should be marked as STARTED (incomplete) - assert result.all[1].status == BatchItemStatus.STARTED +def test_branch_start_rejects_invalid_source_states(): + branch = Branch(Executable(0, lambda: "test")) + branch.start() + branch.complete("done") + with pytest.raises(InvalidStateError, match="Cannot start branch 0"): + branch.start() - # Verify counts - assert result.success_count == 1 - assert result.failure_count == 0 - assert result.started_count == 1 - assert result.total_count == 2 + suspended = Branch(Executable(1, lambda: "test")) + suspended.start() + suspended.suspend() + with pytest.raises(InvalidStateError, match="Cannot start branch 1"): + suspended.start() -def test_executor_returns_with_incomplete_branches(): - """Test that executor returns when min_successful is reached, leaving other branches incomplete.""" +def test_branch_start_allows_timed_resume(): + branch = Branch(Executable(0, lambda: "test")) + branch.start() + branch.suspend_until(123.0) + branch.start() + assert branch.status is BranchStatus.RUNNING - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return executable.func() - operation_tracker = Mock() +def test_create_result_raises_on_failed_branch_without_error(): + executables = [Executable(0, lambda: "test")] + executor = _make_executor(executables, max_concurrency=1) + branch = Branch(executables[0]) + branch.status = BranchStatus.FAILED + executor.branches = [branch] + with pytest.raises(InvalidStateError, match="unexpected state"): + executor._create_result() # noqa: SLF001 - def fast_branch(): - operation_tracker.fast_executed() - return "fast_result" - def slow_branch(): - operation_tracker.slow_started() - time.sleep(2) # Long sleep - operation_tracker.slow_completed() - return "slow_result" +# endregion Branch state transitions - executables = [ - Executable(0, fast_branch), - Executable(1, slow_branch), - ] - completion_config = CompletionConfig(min_successful=1) +# region CompletionPolicy - executor = TestExecutor( - executables=executables, - max_concurrency=2, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, + +def test_completion_policy_from_config_none(): + policy = CompletionPolicy.from_config(5, None) + assert policy.total == 5 + assert policy.min_successful is None + assert policy.tolerated_failure_count is None + assert policy.tolerated_failure_percentage is None + assert not policy.has_criteria + + +def test_completion_policy_from_config_copies_fields(): + config = CompletionConfig( + min_successful=2, + tolerated_failure_count=3, + tolerated_failure_percentage=50.0, ) + policy = CompletionPolicy.from_config(10, config) + assert policy.total == 10 + assert policy.min_successful == 2 + assert policy.tolerated_failure_count == 3 + assert policy.tolerated_failure_percentage == 50.0 + assert policy.has_criteria - execution_state = Mock() - execution_state.create_checkpoint = Mock() - execution_state.wrap_user_function = lambda func, *args, **kwargs: func - executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda idx: f"step_{idx}" # noqa: SLF001 - executor_context._parent_id = "parent" # noqa: SLF001 - executor_context.create_child_context = lambda op_id, *, is_virtual=False: Mock( - state=execution_state + +def test_completion_policy_fail_fast_without_criteria(): + policy = CompletionPolicy.from_config(5, CompletionConfig()) + assert policy.should_continue(failed=0) + assert not policy.should_continue(failed=1) + assert policy.is_tolerance_exceeded(failed=1) + + +def test_completion_policy_min_successful_only_does_not_fail_fast(): + policy = CompletionPolicy.from_config(5, CompletionConfig(min_successful=2)) + assert policy.should_continue(failed=1) + assert policy.should_continue(failed=4) + + +def test_completion_policy_tolerated_failure_count_boundary(): + policy = CompletionPolicy.from_config( + 10, CompletionConfig(tolerated_failure_count=2) ) + assert policy.should_continue(failed=2) + assert not policy.should_continue(failed=3) - result = executor.execute(execution_state, executor_context) - # Verify fast branch executed - assert operation_tracker.fast_executed.call_count == 1 +def test_completion_policy_tolerated_failure_percentage_boundary(): + policy = CompletionPolicy.from_config( + 10, CompletionConfig(tolerated_failure_percentage=30) + ) + assert policy.should_continue(failed=3) + assert not policy.should_continue(failed=4) - # Slow branch may or may not have started (depends on thread scheduling) - # but it definitely should not have completed - assert operation_tracker.slow_completed.call_count == 0, ( - "Executor should return before slow branch completes" + +def test_completion_policy_percentage_with_zero_total(): + policy = CompletionPolicy.from_config( + 0, CompletionConfig(tolerated_failure_percentage=30) ) + assert not policy.is_tolerance_exceeded(failed=0) - # Result should show MIN_SUCCESSFUL_REACHED - assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - # Verify counts - one succeeded, one incomplete - assert result.success_count == 1 - assert result.failure_count == 0 - assert result.started_count == 1 - assert result.total_count == 2 +def test_completion_policy_is_complete_all_done(): + policy = CompletionPolicy.from_config(3, CompletionConfig()) + assert not policy.is_complete(succeeded=1, failed=1) + assert policy.is_complete(succeeded=2, failed=1) -def test_executor_returns_before_slow_branch_completes(): - """Test that executor returns immediately when min_successful is reached, not waiting for slow branches.""" +def test_completion_policy_is_complete_min_successful(): + policy = CompletionPolicy.from_config(5, CompletionConfig(min_successful=2)) + assert not policy.is_complete(succeeded=1, failed=0) + assert policy.is_complete(succeeded=2, failed=0) - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return executable.func() - slow_branch_mock = Mock() +def test_completion_policy_reason_tolerance_checked_first(): + policy = CompletionPolicy.from_config( + 2, CompletionConfig(tolerated_failure_count=0) + ) + assert ( + policy.reason(succeeded=1, failed=1) + is CompletionReason.FAILURE_TOLERANCE_EXCEEDED + ) - def fast_func(): - return "fast" - def slow_func(): - time.sleep(3) # Sleep - slow_branch_mock.completed() # Should not be called before executor returns - return "slow" +def test_completion_policy_reason_all_completed(): + policy = CompletionPolicy.from_config( + 2, CompletionConfig(tolerated_failure_count=1) + ) + assert policy.reason(succeeded=1, failed=1) is CompletionReason.ALL_COMPLETED - executables = [Executable(0, fast_func), Executable(1, slow_func)] - completion_config = CompletionConfig(min_successful=1) - executor = TestExecutor( - executables=executables, - max_concurrency=2, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, +def test_completion_policy_reason_min_successful(): + policy = CompletionPolicy.from_config(5, CompletionConfig(min_successful=2)) + assert ( + policy.reason(succeeded=2, failed=0) is CompletionReason.MIN_SUCCESSFUL_REACHED ) + +def test_completion_policy_reason_defaults_to_all_completed(): + policy = CompletionPolicy.from_config(5, CompletionConfig(min_successful=3)) + assert policy.reason(succeeded=1, failed=0) is CompletionReason.ALL_COMPLETED + + +# endregion CompletionPolicy + + +# region In-flight window semantics + + +def _make_executor_mocks(): + """Build the execution_state / executor_context mock pair used by executor tests.""" execution_state = Mock() execution_state.create_checkpoint = Mock() execution_state.wrap_user_function = lambda func, *args, **kwargs: func @@ -3298,491 +2643,614 @@ def slow_func(): executor_context.create_child_context = lambda op_id, *, is_virtual=False: Mock( state=execution_state ) + return execution_state, executor_context - result = executor.execute(execution_state, executor_context) - # Executor should have returned before slow branch completed - assert not slow_branch_mock.completed.called, ( - "Executor should return before slow branch completes" +class _RecordingExecutor(ConcurrentExecutor): + """Executor whose items delegate to the Executable's own callable.""" + + def execute_item(self, child_context, executable): + return executable.func() + + +def _make_executor(executables, max_concurrency, completion_config=None): + return _RecordingExecutor( + executables=executables, + max_concurrency=max_concurrency, + completion_config=completion_config or CompletionConfig(), + sub_type_top="TOP", + sub_type_iteration="ITER", + name_prefix="test_", + serdes=None, ) - # Result should show MIN_SUCCESSFUL_REACHED - assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - # Verify counts - assert result.success_count == 1 - assert result.failure_count == 0 - assert result.started_count == 1 - assert result.total_count == 2 +def test_max_concurrency_bounds_simultaneous_execution(): + """At most max_concurrency items execute simultaneously.""" + lock = threading.Lock() + active = 0 + peak = 0 + def tracked(): + nonlocal active, peak + with lock: + active += 1 + peak = max(peak, active) + time.sleep(0.05) + with lock: + active -= 1 + return "ok" -# region TimerScheduler edge cases with exact same reschedule time + executables = [Executable(i, tracked) for i in range(6)] + executor = _make_executor(executables, max_concurrency=2) + execution_state, executor_context = _make_executor_mocks() + result = executor.execute(execution_state, executor_context) -def test_timer_scheduler_same_timestamp_with_counter_tiebreaker(): - """ - Test that scheduling two tasks with the exact same resume_time works. + assert peak <= 2 + assert result.success_count == 6 + assert result.completion_reason is CompletionReason.ALL_COMPLETED - This verifies the fix where a counter is used as a tie-breaker to prevent - TypeError when heapq tries to compare ExecutableWithState objects. - """ - resubmit_callback = Mock() - with TimerScheduler(resubmit_callback) as scheduler: - # Create two different ExecutableWithState objects - exe_state1 = ExecutableWithState(Executable(index=0, func=lambda: "test1")) - exe_state2 = ExecutableWithState(Executable(index=1, func=lambda: "test2")) +def test_suspended_branch_holds_concurrency_slot(): + """A suspended branch keeps its slot: pending items must not start (issue #279).""" + started: list[int] = [] - # Use the exact same timestamp for both - same_timestamp = time.time() + 10.0 + def suspending(index: int): + started.append(index) + msg = "awaiting callback" + raise SuspendExecution(msg) - # Both schedules should work fine now - scheduler.schedule_resume(exe_state1, same_timestamp) - scheduler.schedule_resume(exe_state2, same_timestamp) + executables = [ + Executable(0, partial(suspending, 0)), + Executable(1, partial(suspending, 1)), + Executable(2, partial(suspending, 2)), + ] + executor = _make_executor(executables, max_concurrency=2) + execution_state, executor_context = _make_executor_mocks() - # Verify both are in the heap - assert len(scheduler._pending_resumes) == 2 # noqa: SLF001 + with pytest.raises(SuspendExecution): + executor.execute(execution_state, executor_context) - # Verify FIFO ordering (first scheduled should be first in heap) - first_item = scheduler._pending_resumes[0] # noqa: SLF001 - assert first_item[0] == same_timestamp # timestamp - assert first_item[1] == 0 # counter (first scheduled) - assert first_item[2] == exe_state1 # first exe_state + assert sorted(started) == [0, 1] -def test_timer_scheduler_multiple_same_timestamps(): - """ - Test that scheduling many tasks with the same timestamp works correctly. +def test_slot_freed_on_completion_admits_next_item(): + """A terminal completion frees a slot; a suspension does not.""" + started: list[int] = [] - Verifies FIFO ordering is maintained when multiple tasks have identical timestamps. - """ - resubmit_callback = Mock() + def completing(index: int): + started.append(index) + return f"result_{index}" - with TimerScheduler(resubmit_callback) as scheduler: - same_timestamp = time.time() + 10.0 + def suspending(index: int): + started.append(index) + msg = "awaiting callback" + raise SuspendExecution(msg) - # Create and schedule 10 tasks with the same timestamp - exe_states = [ - ExecutableWithState(Executable(index=i, func=lambda i=i: f"test{i}")) - for i in range(10) - ] + executables = [ + Executable(0, partial(completing, 0)), + Executable(1, partial(suspending, 1)), + Executable(2, partial(completing, 2)), + ] + executor = _make_executor(executables, max_concurrency=2) + execution_state, executor_context = _make_executor_mocks() - for exe_state in exe_states: - scheduler.schedule_resume(exe_state, same_timestamp) + with pytest.raises(SuspendExecution): + executor.execute(execution_state, executor_context) - # All should be scheduled successfully - assert len(scheduler._pending_resumes) == 10 # noqa: SLF001 + assert sorted(started) == [0, 1, 2] - # Verify the heap maintains proper ordering - # The first item should have counter 0 - assert scheduler._pending_resumes[0][1] == 0 # noqa: SLF001 +def test_branches_start_in_index_order(): + started: list[int] = [] -def test_timer_scheduler_counter_increments(): - """Test that the schedule counter increments correctly.""" - resubmit_callback = Mock() + def record(index: int): + started.append(index) + return index - with TimerScheduler(resubmit_callback) as scheduler: - exe_state1 = ExecutableWithState(Executable(0, lambda: "test1")) - exe_state2 = ExecutableWithState(Executable(1, lambda: "test2")) - exe_state3 = ExecutableWithState(Executable(2, lambda: "test3")) + executables = [Executable(i, partial(record, i)) for i in range(4)] + executor = _make_executor(executables, max_concurrency=1) + execution_state, executor_context = _make_executor_mocks() - # Schedule with different times - scheduler.schedule_resume(exe_state1, time.time() + 1.0) - scheduler.schedule_resume(exe_state2, time.time() + 2.0) - scheduler.schedule_resume(exe_state3, time.time() + 3.0) + result = executor.execute(execution_state, executor_context) - # Counter should have incremented to 3 - assert scheduler._schedule_counter == 3 # noqa: SLF001 + assert started == [0, 1, 2, 3] + assert result.success_count == 4 -def test_timer_scheduler_fifo_ordering_with_same_timestamp(): - """ - Test that FIFO ordering is maintained when timestamps are equal. +def test_timed_suspend_resumes_in_process_while_sibling_runs(): + """A due timed suspend is resumed in-process while another branch is running.""" + resumed = threading.Event() + call_counts = {0: 0} - When multiple tasks have the same timestamp, they should be processed - in the order they were scheduled (FIFO). The timer thread processes - items synchronously, so callback order is deterministic. - """ - results = [] - resubmit_callback = Mock( - side_effect=lambda batch: results.extend(exe.index for exe in batch) - ) + def timed_then_ok(): + call_counts[0] += 1 + if call_counts[0] == 1: + msg = "retry shortly" + raise TimedSuspendExecution(msg, time.time() + 0.3) + resumed.set() + return "ok_after_resume" - with TimerScheduler(resubmit_callback) as scheduler: - # Use a past timestamp so they trigger immediately - past_time = time.time() - 0.1 + def slow_sibling(): + resumed.wait(timeout=5.0) + return "sibling" - exe_state1 = ExecutableWithState(Executable(0, lambda: "first")) - exe_state2 = ExecutableWithState(Executable(1, lambda: "second")) - exe_state3 = ExecutableWithState(Executable(2, lambda: "third")) + executables = [Executable(0, timed_then_ok), Executable(1, slow_sibling)] + executor = _make_executor(executables, max_concurrency=2) + execution_state, executor_context = _make_executor_mocks() - # Make them all resumable - exe_state1.suspend() - exe_state2.suspend() - exe_state3.suspend() + result = executor.execute(execution_state, executor_context) - # Schedule all with same timestamp - scheduler.schedule_resume(exe_state1, past_time) - scheduler.schedule_resume(exe_state2, past_time) - scheduler.schedule_resume(exe_state3, past_time) + assert call_counts[0] == 2 + assert result.success_count == 2 + assert result.completion_reason is CompletionReason.ALL_COMPLETED + # A resume wave refreshes state once before resubmitting. + execution_state.create_checkpoint.assert_called() - # Wait for timer thread to process them - time.sleep(0.3) - # Verify FIFO order - they should be resubmitted in order 0, 1, 2 - assert results == [0, 1, 2] +def test_all_timed_suspended_parent_suspends_with_earliest_timestamp(): + earliest = time.time() + 100 + latest = time.time() + 200 + def suspend_at(timestamp: float): + msg = "retry later" + raise TimedSuspendExecution(msg, timestamp) -# endregion TimerScheduler edge cases with exact same reschedule time + executables = [ + Executable(0, partial(suspend_at, latest)), + Executable(1, partial(suspend_at, earliest)), + ] + executor = _make_executor(executables, max_concurrency=2) + execution_state, executor_context = _make_executor_mocks() + with pytest.raises(TimedSuspendExecution) as exc_info: + executor.execute(execution_state, executor_context) -# region Completion Reason Inference Tests (from_items) + assert exc_info.value.scheduled_timestamp == earliest -def test_from_items_no_config_with_failures(): - """Validates: Requirements 2.4 - Fail-fast with no config.""" - items = [ - BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), - BatchItem( - 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) - ), - ] - result = BatchResult.from_items(items, completion_config=None) - assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED +def test_never_started_branches_omitted_from_result(): + """Branches that never started are omitted from the batch result (JS parity).""" + release = threading.Event() + def fast(): + return "fast" -def test_from_items_empty_config_with_failures(): - """Validates: Requirements 2.5 - Fail-fast with empty config.""" - items = [ - BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), - BatchItem( - 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) - ), + def slow(): + release.wait(timeout=5.0) + return "slow" + + executables = [ + Executable(0, fast), + Executable(1, slow), + Executable(2, fast), + Executable(3, fast), ] - config = CompletionConfig() # All fields None - result = BatchResult.from_items(items, completion_config=config) - assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED + executor = _make_executor( + executables, + max_concurrency=2, + completion_config=CompletionConfig(min_successful=1), + ) + execution_state, executor_context = _make_executor_mocks() + try: + result = executor.execute(execution_state, executor_context) + finally: + release.set() -def test_from_items_tolerance_checked_before_all_completed(): - """Validates: Requirements 2.1, 2.2 - Tolerance priority.""" - items = [ - BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), - BatchItem( - 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) - ), - BatchItem( - 2, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) - ), - ] - config = CompletionConfig(tolerated_failure_count=1) - result = BatchResult.from_items(items, completion_config=config) - # All completed but tolerance exceeded - should return TOLERANCE_EXCEEDED - assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED + assert result.completion_reason is CompletionReason.MIN_SUCCESSFUL_REACHED + assert result.success_count == 1 + assert result.started_count == 1 + assert result.total_count == 2 + indexes = {item.index for item in result.all} + assert indexes == {0, 1} -def test_from_items_all_completed_within_tolerance(): - """Validates: Requirements 1.1 - All completed.""" - items = [ - BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), - BatchItem( - 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) - ), - ] - config = CompletionConfig(tolerated_failure_count=1) - result = BatchResult.from_items(items, completion_config=config) - assert result.completion_reason == CompletionReason.ALL_COMPLETED +def test_tolerance_breach_stops_submission_of_pending_items(): + """Fail-fast: a failure with no completion criteria stops further submissions.""" + started: list[int] = [] + def failing(index: int): + started.append(index) + msg = "boom" + raise ValueError(msg) -def test_from_items_min_successful_reached(): - """Validates: Requirements 1.3 - Min successful.""" - items = [ - BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), - BatchItem(1, BatchItemStatus.SUCCEEDED, result="ok"), - BatchItem(2, BatchItemStatus.STARTED), - ] - config = CompletionConfig(min_successful=2) - result = BatchResult.from_items(items, completion_config=config) - assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED + executables = [Executable(i, partial(failing, i)) for i in range(4)] + executor = _make_executor(executables, max_concurrency=1) + execution_state, executor_context = _make_executor_mocks() + result = executor.execute(execution_state, executor_context) -def test_from_items_tolerance_count_exceeded(): - """Validates: Requirements 1.2 - Tolerance count.""" - items = [ - BatchItem( - 0, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) - ), - BatchItem( - 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) - ), - BatchItem(2, BatchItemStatus.STARTED), - ] - config = CompletionConfig(tolerated_failure_count=1) - result = BatchResult.from_items(items, completion_config=config) - assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED + assert started == [0] + assert result.completion_reason is CompletionReason.FAILURE_TOLERANCE_EXCEEDED + assert result.failure_count == 1 + assert result.total_count == 1 -def test_from_items_tolerance_percentage_exceeded(): - """Validates: Requirements 1.2 - Tolerance percentage.""" - items = [ - BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), - BatchItem( - 1, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) - ), - BatchItem( - 2, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) - ), - BatchItem( - 3, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) +def test_orphaned_branch_stops_scheduling(): + """An orphaned branch means an ancestor completed: stop starting new work.""" + started: list[int] = [] + + def orphaned(index: int): + started.append(index) + msg = "parent done" + raise OrphanedChildException(msg, operation_id=f"step_{index}") + + executables = [Executable(i, partial(orphaned, i)) for i in range(4)] + executor = _make_executor(executables, max_concurrency=1) + execution_state, executor_context = _make_executor_mocks() + + result = executor.execute(execution_state, executor_context) + + assert started == [0] + assert result.started_count == 1 + assert result.total_count == 1 + + +def test_branch_worker_maps_outcomes_to_events(): + """_branch_worker converts each branch outcome into the matching event.""" + outcomes: dict[int, Callable] = { + 0: lambda: "done", + 1: lambda: (_ for _ in ()).throw(ValueError("boom")), + 2: lambda: (_ for _ in ()).throw(SuspendExecution("cb")), + 3: lambda: (_ for _ in ()).throw(TimedSuspendExecution("retry", 1234.5)), + 4: lambda: (_ for _ in ()).throw( + OrphanedChildException("parent done", operation_id="step_4") ), - ] - config = CompletionConfig(tolerated_failure_percentage=50.0) - # 3 failures out of 4 = 75% > 50% - result = BatchResult.from_items(items, completion_config=config) - assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED + } + class OutcomeExecutor(ConcurrentExecutor): + def execute_item(self, child_context, executable): + return outcomes[executable.index]() -def test_from_items_tolerance_priority_over_min_successful(): - """Validates: Requirements 2.3 - Tolerance takes precedence.""" - items = [ - BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok"), - BatchItem(1, BatchItemStatus.SUCCEEDED, result="ok"), - BatchItem( - 2, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + executables = [Executable(i, outcomes[i]) for i in range(5)] + execution_state, executor_context = _make_executor_mocks() + executor = OutcomeExecutor( + executables=executables, + max_concurrency=5, + completion_config=CompletionConfig(), + sub_type_top="TOP", + sub_type_iteration="ITER", + name_prefix="test_", + serdes=None, + ) + + events: queue.Queue = queue.Queue() + for executable in executables: + executor._branch_worker(executor_context, events, executable) # noqa: SLF001 + + collected = {} + while not events.empty(): + event = events.get_nowait() + collected[event.index] = event + + assert collected[0].kind is BranchEventKind.COMPLETED + assert collected[0].result == "done" + assert collected[1].kind is BranchEventKind.FAILED + # child_handler wraps user exceptions in CallableRuntimeError before + # they reach the worker. + assert isinstance(collected[1].error, CallableRuntimeError) + assert collected[2].kind is BranchEventKind.SUSPENDED + assert collected[3].kind is BranchEventKind.SUSPENDED_UNTIL + assert collected[3].resume_at == 1234.5 + assert collected[4].kind is BranchEventKind.ORPHANED + + +# region Completion record and replay reconstruction + + +def test_completion_record_from_summary_payload_started_indexes(): + payload = json.dumps( + { + "totalCount": 3, + "completionReason": "MIN_SUCCESSFUL_REACHED", + "startedIndexes": [1], + } + ) + record = CompletionRecord.from_summary_payload(payload) + assert record is not None + assert record.completion_reason is CompletionReason.MIN_SUCCESSFUL_REACHED + assert record.started_total == 3 + assert record.started_indexes == frozenset({1}) + + +def test_completion_record_from_summary_payload_completed_indexes(): + payload = json.dumps( + { + "totalCount": 4, + "completionReason": "ALL_COMPLETED", + "completedIndexes": [0, 2], + } + ) + record = CompletionRecord.from_summary_payload(payload) + assert record is not None + assert record.started_indexes == frozenset({1, 3}) + + +@pytest.mark.parametrize( + "payload", + [ + None, + "", + "not json", + '"a json string"', + json.dumps({"totalCount": 3, "completionReason": "MIN_SUCCESSFUL_REACHED"}), + json.dumps({"totalCount": 3, "startedIndexes": [1]}), + json.dumps( + { + "totalCount": 3, + "completionReason": "SOME_FUTURE_REASON", + "startedIndexes": [1], + } ), - BatchItem( - 3, BatchItemStatus.FAILED, error=ErrorObject("msg", "Error", None, None) + json.dumps( + { + "totalCount": "3", + "completionReason": "ALL_COMPLETED", + "startedIndexes": [1], + } ), - ] - config = CompletionConfig(min_successful=2, tolerated_failure_count=1) - # Min successful reached (2) but tolerance exceeded (2 > 1) - result = BatchResult.from_items(items, completion_config=config) - assert result.completion_reason == CompletionReason.FAILURE_TOLERANCE_EXCEEDED - + ], +) +def test_completion_record_rejects_payloads_without_record(payload): + assert CompletionRecord.from_summary_payload(payload) is None -def test_from_items_empty_array(): - """Validates: Edge case - empty items.""" - items = [] - result = BatchResult.from_items(items, completion_config=None) - assert result.completion_reason == CompletionReason.ALL_COMPLETED - assert result.total_count == 0 +def test_summary_index_fields_records_smaller_side(): + few_started = [ + BatchItem(0, BatchItemStatus.SUCCEEDED, result="a"), + BatchItem(1, BatchItemStatus.STARTED), + BatchItem(2, BatchItemStatus.SUCCEEDED, result="b"), + ] + assert CompletionRecord.summary_index_fields(few_started) == {"startedIndexes": [1]} -def test_from_items_all_succeeded(): - """Validates: All items succeeded.""" - items = [ - BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok1"), - BatchItem(1, BatchItemStatus.SUCCEEDED, result="ok2"), + few_completed = [ + BatchItem(0, BatchItemStatus.SUCCEEDED, result="a"), + BatchItem(1, BatchItemStatus.STARTED), + BatchItem(2, BatchItemStatus.STARTED), ] - result = BatchResult.from_items(items, completion_config=None) - assert result.completion_reason == CompletionReason.ALL_COMPLETED - assert result.success_count == 2 + assert CompletionRecord.summary_index_fields(few_completed) == { + "completedIndexes": [0] + } -# endregion Completion Reason Inference Tests +def _make_replay_mocks(branch_checkpoints: dict): + """Mocks for replay: op id 'op_{index}', checkpoints per index.""" + execution_state = Mock() + execution_state.durable_execution_arn = ( + "arn:aws:durable:us-east-1:123456789012:execution/test" + ) -# region Virtual-context wire-format tests + def get_checkpoint_result(operation_id: str): + checkpoint = branch_checkpoints.get(operation_id) + if checkpoint is not None: + return checkpoint + absent = Mock() + absent.is_succeeded.return_value = False + absent.is_failed.return_value = False + absent.is_existent.return_value = False + absent.is_replay_children.return_value = False + return absent + execution_state.get_checkpoint_result = get_checkpoint_result + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda idx: f"op_{idx}" # noqa: SLF001 + executor_context._parent_id = "parent" # noqa: SLF001 + child_context = Mock() + child_context.state = execution_state + executor_context.create_child_context = Mock(return_value=child_context) + return execution_state, executor_context -def test_flat_mode_stamps_grandparent_as_inner_op_parent_id(): - """In FLAT mode, inner operations in a branch stamp the map/parallel op id as parent_id. - This is the core FLAT-mode invariant. Inner operations must not - stamp the branch's own operation id (that would reproduce the - NESTED hierarchy) — they must stamp the enclosing map/parallel op - id, so the branch is collapsed out of the observable hierarchy - even though it still exists as a logical scope for concurrency - and step-id prefixing. +def _succeeded_checkpoint(serialized_result: str): + checkpoint = Mock() + checkpoint.is_succeeded.return_value = True + checkpoint.is_failed.return_value = False + checkpoint.is_existent.return_value = True + checkpoint.is_replay_children.return_value = False + checkpoint.result = serialized_result + return checkpoint - The test drives `_execute_item_in_child_context` with a real - non-virtual executor context, captures the child context the - executor builds for the branch, and asserts the branch's - `_parent_id` equals the executor_context's own `_parent_id`. - """ - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - # Record the child context we receive so the assertions below can - # inspect its identity fields. - self.last_child_context = child_context - return executable.func(child_context) +def test_replay_round_trip_matches_live_result(): + """execute() -> summary -> replay() reconstructs the exact live result.""" - execution_state = Mock() - execution_state.create_checkpoint = Mock() - execution_state.wrap_user_function = lambda func, *args, **kwargs: func + def succeed(index: int): + return f"live_{index}" - # Mock out the checkpoint so the real child_handler reports "not - # existent" (non-existent checkpoint -> normal execution path). - mock_checkpoint = Mock() - mock_checkpoint.is_succeeded.return_value = False - mock_checkpoint.is_failed.return_value = False - mock_checkpoint.is_existent.return_value = False - mock_checkpoint.is_replay_children.return_value = False - execution_state.get_checkpoint_result.return_value = mock_checkpoint + def suspend(): + msg = "awaiting callback" + raise SuspendExecution(msg) - # Build a real DurableContext that represents the map/parallel op. - map_op_id = "map-op-id" - execution_context = ExecutionContext( - durable_execution_arn="arn:aws:durable:us-east-1:0:execution/test" - ) - executor_context = DurableContext( - state=execution_state, - execution_context=execution_context, - parent_id=map_op_id, # This context *is* the map/parallel op. + executables = [ + Executable(0, partial(succeed, 0)), + Executable(1, suspend), + Executable(2, partial(succeed, 2)), + Executable(3, partial(succeed, 3)), + ] + executor = _make_executor( + executables, + max_concurrency=2, + completion_config=CompletionConfig(min_successful=2), ) + execution_state, executor_context = _make_executor_mocks() + live: BatchResult = executor.execute(execution_state, executor_context) - executables = [Executable(index=0, func=lambda ctx: "ok")] - executor = TestExecutor( - executables=executables, - max_concurrency=1, - completion_config=CompletionConfig(min_successful=1), - sub_type_top="MAP", - sub_type_iteration="MAP_ITER", - name_prefix="branch-", - serdes=None, - nesting_type=NestingType.FLAT, + assert [(i.index, i.status) for i in live.all] == [ + (0, BatchItemStatus.SUCCEEDED), + (1, BatchItemStatus.STARTED), + (2, BatchItemStatus.SUCCEEDED), + ] + assert live.completion_reason is CompletionReason.MIN_SUCCESSFUL_REACHED + + summary: str = MapSummaryGenerator()(live) + top_checkpoint = Mock() + top_checkpoint.result = summary + + # Branch 1's checkpoint claims SUCCEEDED: the recorded STARTED must win + # (a terminal checkpoint can land after the completion decision). + replay_state, replay_context = _make_replay_mocks( + { + "op_0": _succeeded_checkpoint('"replayed_0"'), + "op_1": _succeeded_checkpoint('"raced_1"'), + "op_2": _succeeded_checkpoint('"replayed_2"'), + } + ) + replay_executor = _make_executor( + executables, + max_concurrency=2, + completion_config=CompletionConfig(min_successful=2), + ) + replayed: BatchResult = replay_executor.replay( + replay_state, replay_context, top_checkpoint ) - executor._execute_item_in_child_context(executor_context, executables[0]) # noqa: SLF001 + assert [(i.index, i.status) for i in replayed.all] == [ + (i.index, i.status) for i in live.all + ] + assert replayed.completion_reason is live.completion_reason + assert replayed.all[0].result == "replayed_0" + assert replayed.all[1].result is None + assert replayed.all[2].result == "replayed_2" - # The branch's child context must be virtual AND propagate the - # map/parallel op id as its _parent_id. Inner operations stamping - # self._parent_id will therefore report to the map/parallel op. - branch_ctx = executor.last_child_context - assert branch_ctx.is_virtual is True - assert branch_ctx._parent_id == map_op_id # noqa: SLF001 - # The step-id prefix is the branch's own operation id (stable replay id). - assert branch_ctx._step_id_prefix != map_op_id # noqa: SLF001 +def test_replay_without_record_falls_back_to_checkpoint_derivation(): + """No record: every executable is represented, old-style.""" + executables = [Executable(i, lambda: "x") for i in range(3)] + executor = _make_executor(executables, max_concurrency=2) + replay_state, replay_context = _make_replay_mocks( + {"op_0": _succeeded_checkpoint('"done_0"')} + ) -def test_nested_mode_stamps_branch_op_as_inner_op_parent_id(): - """In NESTED mode, inner operations in a branch stamp the branch's own operation id as parent_id.""" + replayed: BatchResult = executor.replay(replay_state, replay_context) - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - self.last_child_context = child_context - return executable.func(child_context) + assert [(i.index, i.status) for i in replayed.all] == [ + (0, BatchItemStatus.SUCCEEDED), + (1, BatchItemStatus.STARTED), + (2, BatchItemStatus.STARTED), + ] - execution_state = Mock() - execution_state.create_checkpoint = Mock() - execution_state.wrap_user_function = lambda func, *args, **kwargs: func - mock_checkpoint = Mock() - mock_checkpoint.is_succeeded.return_value = False - mock_checkpoint.is_failed.return_value = False - mock_checkpoint.is_existent.return_value = False - mock_checkpoint.is_replay_children.return_value = False - execution_state.get_checkpoint_result.return_value = mock_checkpoint +def test_replay_summary_without_index_keys_falls_back(): + """Summaries written before the record existed derive from checkpoints.""" + executables = [Executable(i, lambda: "x") for i in range(2)] + executor = _make_executor(executables, max_concurrency=2) + top_checkpoint = Mock() + top_checkpoint.result = json.dumps( + { + "totalCount": 2, + "successCount": 2, + "failureCount": 0, + "completionReason": "ALL_COMPLETED", + "status": "SUCCEEDED", + "type": "MapResult", + } + ) + replay_state, replay_context = _make_replay_mocks( + { + "op_0": _succeeded_checkpoint('"done_0"'), + "op_1": _succeeded_checkpoint('"done_1"'), + } + ) - map_op_id = "map-op-id" - execution_context = ExecutionContext( - durable_execution_arn="arn:aws:durable:us-east-1:0:execution/test" + replayed: BatchResult = executor.replay( + replay_state, replay_context, top_checkpoint ) - executor_context = DurableContext( - state=execution_state, - execution_context=execution_context, - parent_id=map_op_id, + + assert [(i.index, i.status) for i in replayed.all] == [ + (0, BatchItemStatus.SUCCEEDED), + (1, BatchItemStatus.SUCCEEDED), + ] + + +def test_replay_recorded_terminal_with_missing_checkpoint_is_started(): + """Nested branch recorded terminal but checkpoint absent: conservative STARTED.""" + executables = [Executable(0, lambda: "x")] + executor = _make_executor(executables, max_concurrency=1) + top_checkpoint = Mock() + top_checkpoint.result = json.dumps( + { + "totalCount": 1, + "completionReason": "ALL_COMPLETED", + "startedIndexes": [], + } ) + replay_state, replay_context = _make_replay_mocks({}) - executables = [Executable(index=0, func=lambda ctx: "ok")] - executor = TestExecutor( - executables=executables, - max_concurrency=1, - completion_config=CompletionConfig(min_successful=1), - sub_type_top="MAP", - sub_type_iteration="MAP_ITER", - name_prefix="branch-", - serdes=None, - nesting_type=NestingType.NESTED, + replayed: BatchResult = executor.replay( + replay_state, replay_context, top_checkpoint ) - executor._execute_item_in_child_context(executor_context, executables[0]) # noqa: SLF001 + assert [(i.index, i.status) for i in replayed.all] == [(0, BatchItemStatus.STARTED)] + assert replayed.completion_reason is CompletionReason.ALL_COMPLETED - # In NESTED mode, the branch is a regular child — its _parent_id is - # its own operation id, not the grandparent. - branch_ctx = executor.last_child_context - assert branch_ctx.is_virtual is False - assert branch_ctx._parent_id == branch_ctx._step_id_prefix # noqa: SLF001 - assert branch_ctx._parent_id != map_op_id # noqa: SLF001 +def test_replay_flat_reconstructs_terminal_statuses_by_reexecution(): + """FLAT branches have no checkpoints: re-execution discriminates outcomes.""" -# endregion Virtual-context wire-format tests + def succeed(): + return "flat_ok" + def fail(): + msg = "flat boom" + raise ValueError(msg) -def test_flat_mode_produces_deterministic_step_ids_across_runs(): - """Step ids and inner parent_ids must be deterministic under FLAT mode. + executables = [Executable(0, succeed), Executable(1, fail)] + executor = _RecordingExecutor( + executables=executables, + max_concurrency=2, + completion_config=CompletionConfig( + tolerated_failure_count=1, + ), + sub_type_top="TOP", + sub_type_iteration="ITER", + name_prefix="test_", + serdes=None, + nesting_type=NestingType.FLAT, + ) + top_checkpoint = Mock() + top_checkpoint.result = json.dumps( + { + "totalCount": 2, + "completionReason": "ALL_COMPLETED", + "startedIndexes": [], + } + ) + replay_state, replay_context = _make_replay_mocks({}) + replay_state.wrap_user_function = lambda func, *args, **kwargs: func - Replay depends on regenerating the same operation ids for the same - logical inputs. This test runs the same executor twice against - fresh executor contexts and asserts that the resulting set of - (step_id_prefix, parent_id) pairs is identical. Any source of - non-determinism (e.g. step prefixes that depend on thread - completion order, or parent-id propagation that's different on - the second run) would show up here as a mismatch and would cause - `NonDeterministicExecutionException` at replay time in production. - """ + replayed: BatchResult = executor.replay( + replay_state, replay_context, top_checkpoint + ) - class TestExecutor(ConcurrentExecutor): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.captured = [] + assert [(i.index, i.status) for i in replayed.all] == [ + (0, BatchItemStatus.SUCCEEDED), + (1, BatchItemStatus.FAILED), + ] + assert replayed.all[0].result == "flat_ok" + assert replayed.all[1].error is not None + assert replayed.completion_reason is CompletionReason.ALL_COMPLETED - def execute_item(self, child_context, executable): - self.captured.append( - ( - child_context._step_id_prefix, # noqa: SLF001 - child_context._parent_id, # noqa: SLF001 - ) - ) - return executable.func(child_context) - def make_run(): - execution_state = Mock() - execution_state.create_checkpoint = Mock() +# endregion Completion record and replay reconstruction - mock_checkpoint = Mock() - mock_checkpoint.is_succeeded.return_value = False - mock_checkpoint.is_failed.return_value = False - mock_checkpoint.is_existent.return_value = False - mock_checkpoint.is_replay_children.return_value = False - execution_state.get_checkpoint_result.return_value = mock_checkpoint - execution_context = ExecutionContext( - durable_execution_arn="arn:aws:durable:us-east-1:0:execution/test" - ) - executor_context = DurableContext( - state=execution_state, - execution_context=execution_context, - parent_id="map-op-id", - ) +def test_unlimited_concurrency_starts_all_items(): + """max_concurrency=None keeps the previous start-everything behavior.""" + barrier = threading.Barrier(3, timeout=5.0) - executables = [ - Executable(index=i, func=lambda ctx, i=i: f"r{i}") for i in range(3) - ] - executor = TestExecutor( - executables=executables, - max_concurrency=3, - completion_config=CompletionConfig(min_successful=3), - sub_type_top="MAP", - sub_type_iteration="MAP_ITER", - name_prefix="branch-", - serdes=None, - nesting_type=NestingType.FLAT, - ) - executor.execute(execution_state, executor_context) - return executor.captured + def synchronized(): + barrier.wait() + return "ok" - run_a = make_run() - run_b = make_run() + executables = [Executable(i, synchronized) for i in range(3)] + executor = _make_executor(executables, max_concurrency=None) + execution_state, executor_context = _make_executor_mocks() - # Ordering of captured items is non-deterministic because branches run - # on a ThreadPoolExecutor. What matters is that the SET of (prefix, - # parent_id) pairs is identical across runs — i.e. replay reconstructs - # the same branch identity regardless of completion order. - assert sorted(run_a) == sorted(run_b), ( - "FLAT-mode branch step-id prefixes and parent_ids must be identical " - f"across runs. Run A: {run_a!r}; Run B: {run_b!r}" - ) - # Sanity: all branches reported grandparent (map op id) as their parent. - assert all(parent_id == "map-op-id" for _prefix, parent_id in run_a) + result = executor.execute(execution_state, executor_context) + + assert result.success_count == 3 + + +# endregion In-flight window semantics diff --git a/packages/aws-durable-execution-sdk-python/tests/config_test.py b/packages/aws-durable-execution-sdk-python/tests/config_test.py index 17069355..2275fcb9 100644 --- a/packages/aws-durable-execution-sdk-python/tests/config_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/config_test.py @@ -3,6 +3,8 @@ from concurrent.futures import ThreadPoolExecutor from unittest.mock import Mock +import pytest + from aws_durable_execution_sdk_python.config import ( BatchedInput, CallbackConfig, @@ -19,6 +21,7 @@ StepSemantics, TerminationMode, ) +from aws_durable_execution_sdk_python.exceptions import ValidationError from aws_durable_execution_sdk_python.waits import ( WaitForConditionConfig, WaitForConditionDecision, @@ -277,3 +280,62 @@ def test_invoke_config_with_tenant_id(): """Test InvokeConfig with explicit tenant_id.""" config = InvokeConfig(tenant_id="test-tenant") assert config.tenant_id == "test-tenant" + + +# region Config validation + + +def test_completion_config_rejects_min_successful_below_one(): + with pytest.raises(ValidationError, match="min_successful must be at least 1"): + CompletionConfig(min_successful=0) + + +def test_completion_config_rejects_negative_tolerated_failure_count(): + with pytest.raises( + ValidationError, match="tolerated_failure_count must be non-negative" + ): + CompletionConfig(tolerated_failure_count=-1) + + +def test_completion_config_rejects_out_of_range_failure_percentage(): + with pytest.raises( + ValidationError, match="tolerated_failure_percentage must be between 0 and 100" + ): + CompletionConfig(tolerated_failure_percentage=101) + with pytest.raises( + ValidationError, match="tolerated_failure_percentage must be between 0 and 100" + ): + CompletionConfig(tolerated_failure_percentage=-0.1) + + +def test_completion_config_accepts_boundary_values(): + assert CompletionConfig(min_successful=1).min_successful == 1 + assert CompletionConfig(tolerated_failure_count=0).tolerated_failure_count == 0 + assert ( + CompletionConfig(tolerated_failure_percentage=0).tolerated_failure_percentage + == 0 + ) + assert ( + CompletionConfig(tolerated_failure_percentage=100).tolerated_failure_percentage + == 100 + ) + + +def test_map_config_rejects_max_concurrency_below_one(): + with pytest.raises(ValidationError, match="max_concurrency must be at least 1"): + MapConfig(max_concurrency=0) + with pytest.raises(ValidationError, match="max_concurrency must be at least 1"): + MapConfig(max_concurrency=-1) + + +def test_parallel_config_rejects_max_concurrency_below_one(): + with pytest.raises(ValidationError, match="max_concurrency must be at least 1"): + ParallelConfig(max_concurrency=0) + + +def test_configs_accept_none_max_concurrency_as_unlimited(): + assert MapConfig().max_concurrency is None + assert ParallelConfig(max_concurrency=1).max_concurrency == 1 + + +# endregion Config validation diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py index 5ebae34f..843529b3 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py @@ -671,7 +671,8 @@ def get_checkpoint_result(self, operation_id): ) # Verify replay was called instead of execute - mock_replay.assert_called_once_with(execution_state, map_context) + mock_replay.assert_called_once() + assert mock_replay.call_args.args[:2] == (execution_state, map_context) assert result == expected_batch_result @@ -732,7 +733,8 @@ def get_checkpoint_result(self, operation_id): operation_identifier, ) - mock_replay.assert_called_once_with(execution_state, map_context) + mock_replay.assert_called_once() + assert mock_replay.call_args.args[:2] == (execution_state, map_context) assert result == expected_batch_result diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py index 0a7da40a..dfb02c20 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py @@ -638,7 +638,8 @@ def get_checkpoint_result(self, operation_id): ) # Verify replay was called instead of execute - mock_replay.assert_called_once_with(execution_state, parallel_context) + mock_replay.assert_called_once() + assert mock_replay.call_args.args[:2] == (execution_state, parallel_context) assert result == expected_batch_result @@ -695,7 +696,8 @@ def get_checkpoint_result(self, operation_id): callables, config, execution_state, parallel_context, operation_identifier ) - mock_replay.assert_called_once_with(execution_state, parallel_context) + mock_replay.assert_called_once() + assert mock_replay.call_args.args[:2] == (execution_state, parallel_context) assert result == expected_batch_result