From 9f04c8de0104db5e8291ccd6135c7cf036ec0777 Mon Sep 17 00:00:00 2001 From: yaythomas Date: Mon, 13 Jul 2026 18:43:53 +0000 Subject: [PATCH] feat: limit in-flight operations in map/parallel max_concurrency previously sized the thread pool only, so branches that suspended (invoke, wait, callbacks) released their thread and the next item started immediately. All items were initiated at once regardless of the configured limit. Rework the concurrency executor around a single coordinator loop: - The calling thread owns all branch state and schedules work; worker threads only run branches and report outcomes on a queue. No locks remain in the module. - max_concurrency now bounds in-flight branches. A suspended branch keeps its slot until it reaches a terminal state, matching the JS and Java SDKs. When every in-flight branch is suspended, the parent suspends with the earliest resume timestamp, or indefinitely. - Branches start in index order and operation ids derive from the branch index, so scheduling stays deterministic across invocations. - Timed suspends resume in-process via the event queue timeout, replacing the TimerScheduler background thread. - Consolidate completion logic (thresholds and reason inference) into CompletionPolicy, replacing ExecutionCounters and the duplicated logic in BatchResult. - Omit never-started branches from BatchResult, matching the JS SDK. - Stop scheduling when a branch is orphaned: an ancestor completed, so all further checkpoints under it would be rejected. - Record the completion decision in map/parallel summaries and obey it on replay: summaries carry the started-branch index set (or its complement, whichever is smaller) and replay reconstructs the exact live result instead of re-deriving statuses from child checkpoints. Summaries without the record fall back to checkpoint derivation. This also reconstructs FLAT results on replay, which checkpoint derivation could not discriminate. - Validate config bounds at construction, matching the Java SDK: max_concurrency and min_successful must be at least 1, tolerated_failure_count non-negative, tolerated_failure_percentage within 0-100. Previously max_concurrency=0 silently meant unlimited and min_successful=0 silently meant all items. Closes #279 --- .../test/map/test_map_completion.py | 9 +- .../concurrency/executor.py | 621 +-- .../concurrency/models.py | 551 ++- .../config.py | 32 + .../operation/map.py | 4 +- .../operation/parallel.py | 8 +- .../tests/concurrency_test.py | 3630 +++++++---------- .../tests/config_test.py | 62 + .../tests/operation/map_test.py | 6 +- .../tests/operation/parallel_test.py | 6 +- 10 files changed, 2214 insertions(+), 2715 deletions(-) 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