From e762f7efd4938bd201007d7f9f68fa1e24cc1312 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 now bounds in-flight branches via a coordinator loop, the completion decision is recorded in an SDK-owned summary envelope and obeyed on replay, and completion config bounds are validated before the operation starts. Branch threads abandoned by early completion are joined at invocation end, so no SDK-created thread outlives its invocation. BREAKING CHANGE: max_concurrency=0 and min_successful=0 now raise ValidationError (previously unlimited / all-items). min_successful greater than the item count now raises ValidationError. BatchResult omits never-started branches. Large-result summary payloads use the SDK envelope format; custom summary_generator output is stored under the summary key instead of replacing the payload. The internal MapSummaryGenerator and ParallelSummaryGenerator classes are removed; the envelope supersedes them. CompletionConfig.all_completed() now tolerates all failures instead of failing fast. --- .../test/map/test_map_completion.py | 9 +- .../concurrency/executor.py | 735 +-- .../concurrency/models.py | 631 +-- .../config.py | 106 +- .../context.py | 48 +- .../identifier.py | 18 + .../operation/map.py | 73 +- .../operation/parallel.py | 64 +- .../aws_durable_execution_sdk_python/state.py | 43 + .../tests/concurrency_test.py | 3979 ++++++++--------- .../tests/config_test.py | 64 +- .../tests/context_test.py | 70 +- .../tests/operation/map_test.py | 296 +- .../tests/operation/parallel_test.py | 306 +- .../tests/state_test.py | 101 + 15 files changed, 3552 insertions(+), 2991 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 256589ae..c0e98060 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,21 +4,24 @@ 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, cast from aws_durable_execution_sdk_python.concurrency.models import ( BatchItem, BatchItemStatus, BatchResult, + Branch, + BranchEvent, + BranchEventKind, BranchStatus, + CompletionPolicy, + CompletionReason, + CompletionRecord, Executable, - ExecutableWithState, - ExecutionCounters, - SuspendResult, ) from aws_durable_execution_sdk_python.config import ( ChildConfig, @@ -26,24 +29,29 @@ ) from aws_durable_execution_sdk_python.exceptions import ( DurableOperationError, + InvalidStateError, OrphanedChildException, SuspendExecution, TimedSuspendExecution, ) -from aws_durable_execution_sdk_python.identifier import OperationIdentifier +from aws_durable_execution_sdk_python.identifier import ( + OperationIdentifier, + OperationIdNamespace, +) from aws_durable_execution_sdk_python.lambda_service import ErrorObject from aws_durable_execution_sdk_python.operation.child import child_handler 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.types import SummaryGenerator + from aws_durable_execution_sdk_python.state import ( + CheckpointedResult, + ExecutionState, + ) logger = logging.getLogger(__name__) @@ -55,99 +63,44 @@ ResultType = TypeVar("ResultType") -# region concurrency logic -class TimerScheduler: - """Manage timed suspend tasks with a background timer thread.""" +def _branch_error_object(err: Exception) -> ErrorObject: + """Convert a failed branch's error for the batch result. + + Records the raw escaping type so live results, branch FAIL checkpoints, + and replay reconstruction all carry the same discriminator; + ``from_exception`` on the ChildContextError wrapper would instead + record the wrapper class name. + """ + if isinstance(err, DurableOperationError): + return ErrorObject( + message=err.message, + type=err.error_type, + data=err.data, + stack_trace=err.stack_trace, + ) + return ErrorObject.from_exception(err) - 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 +class ConcurrentExecutor(Generic[CallableType, ResultType]): + """Execute durable operations concurrently. This contains the execution logic for Map and Parallel. - 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) + 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. -class ConcurrentExecutor(ABC, Generic[CallableType, ResultType]): - """Execute durable operations concurrently. This contains the execution logic for Map and Parallel.""" + 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, @@ -158,300 +111,309 @@ def __init__( sub_type_iteration: OperationSubType, name_prefix: str, serdes: SerDes | None, + operation_id_namespace: OperationIdNamespace, item_serdes: SerDes | None = None, - summary_generator: SummaryGenerator | None = None, nesting_type: NestingType = NestingType.NESTED, ): - """Initialize ConcurrentExecutor. - - Args: - summary_generator: Optional function to generate compact summaries for large results. - When the serialized result exceeds 256KB, this generator creates a JSON summary - instead of checkpointing the full result. Used by map/parallel operations to - handle large BatchResult payloads efficiently. Matches TypeScript behavior in - run-in-child-context-handler.ts. - """ self.executables = executables + self.operation_id_namespace = operation_id_namespace self.max_concurrency = max_concurrency self.completion_config = completion_config self.sub_type_top = sub_type_top self.sub_type_iteration = sub_type_iteration 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 - @abstractmethod - def execute_item( + self.policy: CompletionPolicy = CompletionPolicy.from_config( + len(executables), completion_config + ) + self.branches: list[Branch[CallableType, ResultType]] = [] + + def execute_item( # noqa: PLR6301 self, child_context: DurableContext, executable: Executable[CallableType] ) -> ResultType: """Execute a single executable in a child context and return the result.""" - raise NotImplementedError + logger.debug("▶️ Processing branch: %s", executable.index) + func = cast("Callable[[DurableContext], ResultType]", executable.func) + result: ResultType = func(child_context) + logger.debug("✅ Processed branch: %s", executable.index) + return result def get_iteration_name(self, index: int) -> str: """Get the display name for an iteration/branch at the given index. - Subclasses can override this to provide custom naming (e.g., from item_namer - or branch names). The default returns "{name_prefix}{index}". + Returns the Executable's bound name when present, else + "{name_prefix}{index}". An explicitly provided empty name is + preserved. """ - return f"{self.name_prefix}{index}" + name: str | None = self.executables[index].name + return name if name is not None else f"{self.name_prefix}{index}" 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) - - def on_done(future: Future) -> None: - self._on_task_complete(executable_with_state, future, scheduler) - - future.add_done_callback(on_done) - return future - - # Submit initial tasks - futures = [ - submit_task(exe_state) for exe_state in self.executables_with_state - ] - - # 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 - - # Suspend execution if everything done and at least one of the tasks raised a suspend exception. - if self._suspend_exception: - raise self._suspend_exception - - 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) - - # 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, - ) + max_in_flight: int = self.max_concurrency or len(self.executables) + self.branches = [Branch(executable=exe) for exe in self.executables] + + 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 + } + + in_flight: int = 0 + running: int = 0 + succeeded: int = 0 + failed: int = 0 + + pool: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=max_in_flight) + # Registered so ExecutionState.close() joins any branches still + # running after early completion before the invocation returns. + execution_state.register_branch_pool(pool) + + def submit(branch: Branch[CallableType, ResultType]) -> None: + branch.start() + pool.submit( + self._branch_worker, executor_context, events, branch.executable ) - 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() + 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." + ) - def _create_result(self) -> BatchResult[ResultType]: - """ - Build the final BatchResult. + 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 BranchEventKind.FATAL if event.fatal_error is not None: + # System-level failure: propagate immediately without + # counting a branch failure or checkpointing further. + raise event.fatal_error + 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 continue in the background of this invocation + # and raise OrphanedChildException on their next attempt to + # checkpoint. ExecutionState.close() joins them before the + # invocation returns, so no branch thread outlives the + # invocation. + pool.shutdown(wait=False, cancel_futures=True) + + # The decision that ended the loop determines the reason. Captured + # before the drain so raced terminal events update item statuses + # without flipping the reason (and the recorded summary) to a + # decision that never fired. + completion_reason: CompletionReason = self.policy.reason(succeeded, failed) + + # 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) + elif raced.kind is BranchEventKind.FATAL and raced.fatal_error is not None: + # A straggler hit a system-level failure after the completion + # decision. The same failure would reject the parent's own + # checkpoint, so propagate instead of returning a result. + raise raced.fatal_error + + return self._create_result(completion_reason) + + def _create_result( + self, completion_reason: CompletionReason | None = None + ) -> BatchResult[ResultType]: + """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 the one captured when the completion + decision fired. When absent it is computed from the branch states + against the true batch total. """ + 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: - err = executable.error - # Record the raw escaping type so first-run matches the - # branch's FAIL checkpoint that replay reads back; - # from_exception on the ChildContextError wrapper would - # instead record the wrapper class name. - error_object = ( - ErrorObject( - message=err.message, - type=err.error_type, - data=err.data, - stack_trace=err.stack_trace, - ) - if isinstance(err, DurableOperationError) - else ErrorObject.from_exception(err) - ) + case BranchStatus.FAILED if branch.error is not None: + failed += 1 batch_items.append( BatchItem( - executable.index, + branch.index, BatchItemStatus.FAILED, - error=error_object, + error=_branch_error_object(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) + + if completion_reason is None: + completion_reason = self.policy.reason(succeeded, failed) + return BatchResult(batch_items, completion_reason) + + 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)) + except BaseException as e: + # System-level failure (background checkpoint failure, SystemExit). + # Post a fatal event so the coordinator re-raises it on the + # calling thread instead of blocking forever on the queue, then + # let the exception propagate to the worker thread. + events.put(BranchEvent.fatal(executable.index, e)) + raise + else: + events.put(BranchEvent.completed(executable.index, result)) def _execute_item_in_child_context( self, @@ -472,7 +434,7 @@ def _execute_item_in_child_context( and execution-order invariant. """ - operation_id: str = executor_context._create_step_id_for_logical_step( # noqa: SLF001 + operation_id: str = self.operation_id_namespace.create_id_for_step( executable.index ) name: str = self.get_iteration_name(executable.index) @@ -522,23 +484,105 @@ def run_in_child_handler() -> ResultType: config=ChildConfig( serdes=self.item_serdes or self.serdes, sub_type=self.sub_type_iteration, - summary_generator=self.summary_generator, is_virtual=is_virtual, ), ) 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 (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 = self.operation_id_namespace.create_id_for_step( + 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=_branch_error_object(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: - operation_id = executor_context._create_step_id_for_logical_step( # noqa: SLF001 + operation_id = self.operation_id_namespace.create_id_for_step( executable.index ) checkpoint = execution_state.get_checkpoint_result(operation_id) @@ -561,6 +605,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 4a842b69..948af732 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,10 +1,9 @@ -"""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 @@ -13,20 +12,17 @@ from aws_durable_execution_sdk_python.exceptions import ( ChildContextError, InvalidStateError, - SuspendExecution, ) 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 + from aws_durable_execution_sdk_python.types import SummaryGenerator logger = logging.getLogger(__name__) -T = TypeVar("T") R = TypeVar("R") CallableType = TypeVar("CallableType") @@ -47,17 +43,232 @@ 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. + + 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 valid decision + record: absent or empty payloads, payloads written before the + record existed, unknown completion reasons, and malformed or + out-of-range index fields. Never raises on malformed input. + """ + 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 + if isinstance(started_total, bool) or started_total < 0: + 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): + indexes = cls._validated_indexes(started_raw, started_total) + if indexes is None: + return None + started_indexes = indexes + elif isinstance(completed_raw, list): + indexes = cls._validated_indexes(completed_raw, started_total) + if indexes is None: + return None + started_indexes = frozenset(range(started_total)) - indexes + else: + return None + + return cls( + completion_reason=completion_reason, + started_total=started_total, + started_indexes=started_indexes, + ) @staticmethod - def do_not_suspend() -> SuspendResult: - return SuspendResult(should_suspend=False) + def _validated_indexes( + raw: list[object], started_total: int + ) -> frozenset[int] | None: + """Validate a raw index list: ints (not bools) within [0, started_total).""" + validated: set[int] = set() + for element in raw: + if isinstance(element, bool) or not isinstance(element, int): + return None + if element < 0 or element >= started_total: + return None + validated.add(element) + return frozenset(validated) @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} + + @staticmethod + def summary_envelope( + result: BatchResult, + result_type: str, + summary: str | None, + ) -> str: + """Serialize a batch operation's large-result summary payload. + + The payload is an SDK-owned JSON envelope: the completion decision + record (``totalCount``, ``completionReason``, and the started or + completed index set) that replay requires, informational view + fields, and the optional customer summary verbatim under + ``summary``. Written and parsed with plain JSON, independent of + any configured serdes. A payload exceeding the checkpoint size + limit fails the operation when checkpointed. + """ + fields: dict[str, object] = { + "type": result_type, + "totalCount": result.total_count, + "completionReason": result.completion_reason.value, + **CompletionRecord.summary_index_fields(result.all), + "startedCount": result.started_count, + "successCount": result.success_count, + "failureCount": result.failure_count, + "status": result.status.value, + } + if summary is not None: + fields["summary"] = summary if isinstance(summary, str) else str(summary) + return json.dumps(fields) + + +def envelope_summary_generator( + result_type: str, + custom_generator: SummaryGenerator | None, +) -> SummaryGenerator: + """Build the summary generator for a batch operation's parent context. + + The returned generator wraps the optional customer generator: the SDK + always writes the envelope with the completion decision record, and + the customer generator only contributes the ``summary`` field. An + exception raised by the customer generator propagates and fails the + operation. + """ + + def generate(result: BatchResult) -> str: + summary: str | None = None + if custom_generator is not None: + summary = custom_generator(result) + return CompletionRecord.summary_envelope(result, result_type, summary) + + return generate @dataclass(frozen=True) @@ -115,115 +326,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 { @@ -307,6 +430,7 @@ def total_count(self) -> int: class Executable(Generic[CallableType]): index: int func: CallableType + name: str | None = None class BranchStatus(Enum): @@ -318,228 +442,111 @@ 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 - - @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 - - @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 +class Branch(Generic[CallableType, ResultType]): + """Execution state of a single branch. - @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 + Owned exclusively by the coordinator thread: every transition happens + there, so no synchronization is needed. + """ - @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 complete(self, result: ResultType | None) -> None: + self.status = BranchStatus.COMPLETED + self.result = result - 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 fail(self, error: Exception) -> None: + self.status = BranchStatus.FAILED + self.error = error - def complete(self, result: ResultType) -> None: - """Transition to COMPLETED state.""" - self._status = BranchStatus.COMPLETED - self._result = result - self._is_result_set = True + def suspend(self) -> None: + """Suspend indefinitely, pending an external callback.""" + self.status = BranchStatus.SUSPENDED + self.resume_at = None - def fail(self, error: Exception) -> None: - """Transition to FAILED state.""" - self._status = BranchStatus.FAILED - self._error = error + 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 - def reset_to_pending(self) -> None: - """Reset to PENDING state for resubmission.""" - self._status = BranchStatus.PENDING - self._future = None - self._suspend_until = None - # endregion State transitions +class BranchEventKind(Enum): + COMPLETED = "completed" + FAILED = "failed" + SUSPENDED = "suspended" + SUSPENDED_UNTIL = "suspended_until" + ORPHANED = "orphaned" + FATAL = "fatal" -class ExecutionCounters: - """Thread-safe counters for tracking execution state.""" +@dataclass(frozen=True) +class BranchEvent(Generic[ResultType]): + """Outcome of one branch execution attempt. - 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 + The only message worker threads send to the coordinator. Workers never + touch branch state directly. + """ - return True + index: int + kind: BranchEventKind + result: ResultType | None = None + error: Exception | None = None + resume_at: float | None = None + fatal_error: BaseException | None = None - 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 + @classmethod + def completed( + cls, index: int, result: ResultType | None + ) -> BranchEvent[ResultType]: + return cls(index=index, kind=BranchEventKind.COMPLETED, result=result) - # All tasks completed - if completed_count == self.total_tasks: - return True + @classmethod + def failed(cls, index: int, error: Exception) -> BranchEvent[ResultType]: + return cls(index=index, kind=BranchEventKind.FAILED, error=error) - # when we breach min successful, we've completed - return self.success_count >= self.min_successful + @classmethod + def suspended(cls, index: int) -> BranchEvent[ResultType]: + return cls(index=index, kind=BranchEventKind.SUSPENDED) - 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 suspended_until(cls, index: int, resume_at: float) -> BranchEvent[ResultType]: + return cls( + index=index, kind=BranchEventKind.SUSPENDED_UNTIL, resume_at=resume_at + ) - 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 orphaned(cls, index: int) -> BranchEvent[ResultType]: + return cls(index=index, kind=BranchEventKind.ORPHANED) - # 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 fatal(cls, index: int, error: BaseException) -> BranchEvent[ResultType]: + """A system-level error that must propagate to the calling thread. - return False + Carries the original BaseException (for example a background + checkpoint failure) so the coordinator re-raises it instead of + recording a branch failure. + """ + return cls(index=index, kind=BranchEventKind.FATAL, fatal_error=error) -# 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 f4e062b5..e25d1b98 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 @@ -33,7 +33,7 @@ class Duration: seconds: int = 0 - def __post_init__(self): + def __post_init__(self) -> None: if self.seconds < 0: msg = "Duration seconds must be positive" raise ValidationError(msg) @@ -137,6 +137,43 @@ class CompletionConfig: tolerated_failure_count: int | None = None tolerated_failure_percentage: int | float | None = None + def __post_init__(self) -> None: + 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) + + def _validate_for_total(self, total: int) -> None: + """Validate this config against the number of items it will govern. + + SDK-internal: called by DurableContext.map and DurableContext.parallel + before the operation's child context starts, so the error surfaces as + a bare ValidationError (matching wait and wait_for_condition + validation) instead of a checkpointed operation failure. + """ + if self.min_successful is not None and self.min_successful > total: + msg = ( + f"min_successful cannot be greater than total items: " + f"{self.min_successful} > {total}" + ) + raise ValidationError(msg) + # TODO: reevaluate this # @staticmethod # def first_completed(): @@ -154,10 +191,12 @@ def first_successful(): @staticmethod def all_completed(): + # 100% tolerated failures: every item runs regardless of failures. + # All-None fields would select the fail-fast default instead. return CompletionConfig( min_successful=None, tolerated_failure_count=None, - tolerated_failure_percentage=None, + tolerated_failure_percentage=100, ) @staticmethod @@ -202,13 +241,15 @@ class ParallelConfig: - item_serdes: Used for individual function results in child contexts - serdes: Used for the entire BatchResult at handler level - summary_generator: Function to generate compact summaries for large results (>256KB). - When the serialized result exceeds CHECKPOINT_SIZE_LIMIT, this generator - creates a JSON summary instead of checkpointing the full result. The operation - is marked with ReplayChildren=true to reconstruct the full result during replay. - - Used internally by map/parallel operations to handle large BatchResult payloads. - Signature: (result: T) -> str + summary_generator: Function contributing a customer-facing summary for large + results (>256KB). When the serialized result exceeds CHECKPOINT_SIZE_LIMIT, + the SDK checkpoints a compact JSON payload instead of the full result and + marks the operation ReplayChildren=true so the full result is reconstructed + during replay. The SDK always writes the fields replay requires; the + generator's return value is stored verbatim under the payload's "summary" + key for observability and is never read by the SDK. The summary is + checkpointed as provided. An exception raised by the generator fails + the operation. Signature: (result: T) -> str nesting_type: How child operations should inherit context from their parent. - NESTED: Each branch runs in its own isolated context (default) @@ -231,6 +272,11 @@ class ParallelConfig: summary_generator: SummaryGenerator | None = None nesting_type: NestingType = NestingType.NESTED + def __post_init__(self) -> None: + 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]): @@ -297,12 +343,14 @@ class ChildConfig(Generic[T]): Examples: OperationSubType.MAP_ITERATION, OperationSubType.PARALLEL_BRANCH. Used internally by the execution engine for operation classification. - summary_generator: Function to generate compact summaries for large results (>256KB). - When the serialized result exceeds CHECKPOINT_SIZE_LIMIT, this generator - creates a JSON summary instead of checkpointing the full result. The operation - is marked with ReplayChildren=true to reconstruct the full result during replay. - - Used internally by map/parallel operations to handle large BatchResult payloads. + summary_generator: Function generating the checkpoint payload for large + results (>256KB). When the serialized result exceeds CHECKPOINT_SIZE_LIMIT, + the SDK checkpoints the generator's output instead of the full result and + marks the operation ReplayChildren=true so the full result is reconstructed + during replay. The output is checkpointed as provided. For map and + parallel operations the SDK supplies a generator that writes the + completion-record envelope; see MapConfig and ParallelConfig. An + exception raised by the generator fails the operation. Signature: (result: T) -> str is_virtual: When True, skip all checkpoints (START, SUCCEED, @@ -358,13 +406,15 @@ class MapConfig(Generic[T]): - item_serdes: Used for individual item results in child contexts - serdes: Used for the entire BatchResult at handler level - summary_generator: Function to generate compact summaries for large results (>256KB). - When the serialized result exceeds CHECKPOINT_SIZE_LIMIT, this generator - creates a JSON summary instead of checkpointing the full result. The operation - is marked with ReplayChildren=true to reconstruct the full result during replay. - - Used internally by map/parallel operations to handle large BatchResult payloads. - Signature: (result: T) -> str + summary_generator: Function contributing a customer-facing summary for large + results (>256KB). When the serialized result exceeds CHECKPOINT_SIZE_LIMIT, + the SDK checkpoints a compact JSON payload instead of the full result and + marks the operation ReplayChildren=true so the full result is reconstructed + during replay. The SDK always writes the fields replay requires; the + generator's return value is stored verbatim under the payload's "summary" + key for observability and is never read by the SDK. The summary is + checkpointed as provided. An exception raised by the generator fails + the operation. Signature: (result: T) -> str nesting_type: How child operations should inherit context from their parent. - NESTED: Each item runs in its own isolated context (default) @@ -372,8 +422,11 @@ class MapConfig(Generic[T]): item_namer: Optional callable to generate custom names for each map iteration. When provided, replaces the default "map-item-{index}" naming scheme. - Receives the item and its index, and returns a string name for that iteration. - This affects observability (execution history names) but not replay determinism. + Receives the item and its index, and returns a string name for that + iteration. Called eagerly for every input when the map starts, + including items that never run due to early completion, and again + on every replay, so it must be deterministic and side-effect-free. + An exception raised by the callable fails the map operation. If None, uses the default naming: "map-item-{index}". Example: @@ -398,6 +451,11 @@ class MapConfig(Generic[T]): nesting_type: NestingType = NestingType.NESTED item_namer: Callable[[T, int], str] | None = None + def __post_init__(self) -> None: + 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/context.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py index de1dc616..ff1ab6bc 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py @@ -1,6 +1,5 @@ from __future__ import annotations -import hashlib import logging from contextlib import contextmanager from dataclasses import dataclass @@ -26,6 +25,9 @@ StepConfig, WaitForCallbackConfig, ) +from aws_durable_execution_sdk_python.concurrency.models import ( + envelope_summary_generator, +) from aws_durable_execution_sdk_python.exceptions import ( CallbackError, CallbackExternalError, @@ -36,7 +38,10 @@ SuspendExecution, ValidationError, ) -from aws_durable_execution_sdk_python.identifier import OperationIdentifier +from aws_durable_execution_sdk_python.identifier import ( + OperationIdentifier, + OperationIdNamespace, +) from aws_durable_execution_sdk_python.lambda_service import ( OperationSubType, OperationType, @@ -339,6 +344,9 @@ def __init__( self._step_id_prefix: str | None = ( step_id_prefix if step_id_prefix is not None else parent_id ) + self._operation_id_namespace: OperationIdNamespace = OperationIdNamespace( + self._step_id_prefix + ) # cached at construction to make invariant even if parent/prefix mutates. self._is_virtual: bool = self._parent_id != self._step_id_prefix self._step_counter: OrderedCounter = OrderedCounter() @@ -469,9 +477,7 @@ def _create_step_id_for_logical_step(self, step: int) -> str: This allows us to recover operation ids or even look forward without changing the internal state of this context. """ - prefix: str | None = self._step_id_prefix - step_id: str = f"{prefix}-{step}" if prefix else str(step) - return hashlib.blake2b(step_id.encode()).hexdigest()[:64] + return self._operation_id_namespace.create_id_for_step(step) def _create_step_id(self) -> str: """Generate a thread-safe step id, incrementing in order of invocation. @@ -662,6 +668,12 @@ def map( """Execute a callable for each item in parallel.""" map_name: str | None = self._resolve_step_name(name, func) + # Validate before the child context starts, so the error surfaces as + # a bare ValidationError (matching wait and wait_for_condition) with + # no STARTED/FAILED map operation in history. + if config is not None: + config.completion_config._validate_for_total(len(inputs)) + with self._replay_aware(): operation_id = self._create_step_id() operation_identifier = OperationIdentifier( @@ -684,6 +696,7 @@ def map_in_child_context() -> BatchResult[R]: execution_state=self.state, map_context=map_context, operation_identifier=operation_identifier, + operation_id_namespace=OperationIdNamespace(operation_id), ) return child_handler( @@ -692,7 +705,14 @@ def map_in_child_context() -> BatchResult[R]: operation_identifier=operation_identifier, config=ChildConfig( sub_type=OperationSubType.MAP, - serdes=getattr(config, "serdes", None), + serdes=config.serdes if config else None, + # The SDK-owned envelope records the completion decision for + # replay. A configured summary_generator only contributes the + # customer-facing summary field. + summary_generator=envelope_summary_generator( + "MapResult", + config.summary_generator if config else None, + ), ), ) @@ -703,6 +723,12 @@ def parallel( config: ParallelConfig | None = None, ) -> BatchResult[T]: """Execute multiple callables in parallel.""" + # Validate before the child context starts, so the error surfaces as + # a bare ValidationError (matching wait and wait_for_condition) with + # no STARTED/FAILED parallel operation in history. + if config is not None: + config.completion_config._validate_for_total(len(functions)) + with self._replay_aware(): # _create_step_id() is thread-safe. rest of method is safe, since using local copy of parent id operation_id = self._create_step_id() @@ -725,6 +751,7 @@ def parallel_in_child_context() -> BatchResult[T]: execution_state=self.state, parallel_context=parallel_context, operation_identifier=operation_identifier, + operation_id_namespace=OperationIdNamespace(operation_id), ) return child_handler( @@ -733,7 +760,14 @@ def parallel_in_child_context() -> BatchResult[T]: operation_identifier=operation_identifier, config=ChildConfig( sub_type=OperationSubType.PARALLEL, - serdes=getattr(config, "serdes", None), + serdes=config.serdes if config else None, + # The SDK-owned envelope records the completion decision for + # replay. A configured summary_generator only contributes the + # customer-facing summary field. + summary_generator=envelope_summary_generator( + "ParallelResult", + config.summary_generator if config else None, + ), ), ) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/identifier.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/identifier.py index 89d07727..9a46bae1 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/identifier.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/identifier.py @@ -2,6 +2,7 @@ from __future__ import annotations +import hashlib from dataclasses import dataclass from aws_durable_execution_sdk_python.lambda_service import ( @@ -10,6 +11,23 @@ ) +@dataclass(frozen=True) +class OperationIdNamespace: + """The operation-id namespace of one context. + + Maps a logical step position to its deterministic operation id. + Pure: the same position always yields the same id, so ids can be + derived concurrently and ahead of execution without mutating + context state. + """ + + prefix: str | None = None + + def create_id_for_step(self, step: int) -> str: + step_id: str = f"{self.prefix}-{step}" if self.prefix else str(step) + return hashlib.blake2b(step_id.encode()).hexdigest()[:64] + + @dataclass(frozen=True) class OperationIdentifier: """Container for operation id, parent id, and name.""" 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..c6b6a102 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 @@ -2,8 +2,6 @@ from __future__ import annotations -import json -import logging from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, Generic, TypeVar @@ -18,15 +16,15 @@ if TYPE_CHECKING: from aws_durable_execution_sdk_python.context import DurableContext - from aws_durable_execution_sdk_python.identifier import OperationIdentifier + from aws_durable_execution_sdk_python.identifier import ( + OperationIdentifier, + OperationIdNamespace, + ) from aws_durable_execution_sdk_python.serdes import SerDes from aws_durable_execution_sdk_python.state import ( CheckpointedResult, ExecutionState, ) - from aws_durable_execution_sdk_python.types import SummaryGenerator - -logger = logging.getLogger(__name__) # Input item type T = TypeVar("T") @@ -45,10 +43,9 @@ def __init__( iteration_sub_type: OperationSubType, name_prefix: str, serdes: SerDes | None, - summary_generator: SummaryGenerator | None = None, + operation_id_namespace: OperationIdNamespace, item_serdes: SerDes | None = None, nesting_type: NestingType = NestingType.NESTED, - item_namer: Callable[[T, int], str] | None = None, ): super().__init__( executables=executables, @@ -58,12 +55,11 @@ def __init__( sub_type_iteration=iteration_sub_type, name_prefix=name_prefix, serdes=serdes, - summary_generator=summary_generator, + operation_id_namespace=operation_id_namespace, item_serdes=item_serdes, nesting_type=nesting_type, ) self.items = items - self._item_namer = item_namer @classmethod def from_items( @@ -71,10 +67,24 @@ def from_items( items: Sequence[T], func: Callable, config: MapConfig[T], + operation_id_namespace: OperationIdNamespace, ) -> MapExecutor[T, R]: """Create MapExecutor from items and a callable.""" + + def bind(i: int, item: T) -> Callable: + def run(child_context: DurableContext) -> R: + result: R = func(child_context, item, i, items) + return result + + return run + executables: list[Executable[Callable]] = [ - Executable(index=i, func=func) for i in range(len(items)) + Executable( + index=i, + func=bind(i, item), + name=config.item_namer(item, i) if config.item_namer else None, + ) + for i, item in enumerate(items) ] return cls( @@ -86,25 +96,11 @@ def from_items( iteration_sub_type=OperationSubType.MAP_ITERATION, name_prefix="map-item-", serdes=config.serdes, - summary_generator=config.summary_generator, + operation_id_namespace=operation_id_namespace, item_serdes=config.item_serdes, nesting_type=config.nesting_type, - item_namer=config.item_namer, ) - def get_iteration_name(self, index: int) -> str: - """Return custom item name if item_namer is provided, otherwise default.""" - if self._item_namer is not None: - return self._item_namer(self.items[index], index) - return super().get_iteration_name(index) - - def execute_item(self, child_context, executable: Executable[Callable]) -> R: - logger.debug("🗺️ Processing map item: %s", executable.index) - item = self.items[executable.index] - result: R = executable.func(child_context, item, executable.index, self.items) - logger.debug("✅ Processed map item: %s", executable.index) - return result - def map_handler( items: Sequence[T], @@ -113,18 +109,16 @@ def map_handler( execution_state: ExecutionState, map_context: DurableContext, operation_identifier: OperationIdentifier, + operation_id_namespace: OperationIdNamespace, ) -> BatchResult[R]: """Execute a callable for each item in parallel.""" - # Summary Generator Construction (matches TypeScript implementation): - # Construct the summary generator at the handler level, just like TypeScript does in map-handler.ts. - # This matches the pattern where handlers are responsible for configuring operation-specific behavior. - # - # See TypeScript reference: aws-durable-execution-sdk-js/src/handlers/map-handler/map-handler.ts (~line 79) + map_config: MapConfig = config or MapConfig() executor: MapExecutor[T, R] = MapExecutor.from_items( items=items, func=func, - config=config or MapConfig(summary_generator=MapSummaryGenerator()), + config=map_config, + operation_id_namespace=operation_id_namespace, ) checkpoint: CheckpointedResult = execution_state.get_checkpoint_result( @@ -132,19 +126,6 @@ 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) - - -class MapSummaryGenerator: - def __call__(self, result: BatchResult) -> str: - fields = { - "totalCount": result.total_count, - "successCount": result.success_count, - "failureCount": result.failure_count, - "completionReason": result.completion_reason.value, - "status": result.status.value, - "type": "MapResult", - } - 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..d4cb2703 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 @@ -2,7 +2,6 @@ from __future__ import annotations -import json import logging from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, TypeVar @@ -20,10 +19,12 @@ if TYPE_CHECKING: from aws_durable_execution_sdk_python.concurrency.models import BatchResult from aws_durable_execution_sdk_python.context import DurableContext - from aws_durable_execution_sdk_python.identifier import OperationIdentifier + from aws_durable_execution_sdk_python.identifier import ( + OperationIdentifier, + OperationIdNamespace, + ) from aws_durable_execution_sdk_python.serdes import SerDes from aws_durable_execution_sdk_python.state import ExecutionState - from aws_durable_execution_sdk_python.types import SummaryGenerator logger = logging.getLogger(__name__) @@ -41,7 +42,7 @@ def __init__( iteration_sub_type: OperationSubType, name_prefix: str, serdes: SerDes | None, - summary_generator: SummaryGenerator | None = None, + operation_id_namespace: OperationIdNamespace, item_serdes: SerDes | None = None, nesting_type: NestingType = NestingType.NESTED, ): @@ -53,7 +54,7 @@ def __init__( sub_type_iteration=iteration_sub_type, name_prefix=name_prefix, serdes=serdes, - summary_generator=summary_generator, + operation_id_namespace=operation_id_namespace, item_serdes=item_serdes, nesting_type=nesting_type, ) @@ -63,15 +64,20 @@ def from_callables( cls, callables: Sequence[Callable | ParallelBranch], config: ParallelConfig, + operation_id_namespace: OperationIdNamespace, ) -> ParallelExecutor: """Create ParallelExecutor from a sequence of callables or ParallelBranch instances. Since ParallelBranch is callable, it is stored directly as the func in - each Executable. The get_iteration_name method inspects the func to - extract the branch name when available. + each Executable, with its name bound when provided. """ executables: list[Executable[Callable]] = [ - Executable(index=i, func=func) for i, func in enumerate(callables) + Executable( + index=i, + func=func, + name=func.name if isinstance(func, ParallelBranch) else None, + ) + for i, func in enumerate(callables) ] return cls( @@ -82,24 +88,11 @@ def from_callables( iteration_sub_type=OperationSubType.PARALLEL_BRANCH, name_prefix="parallel-branch-", serdes=config.serdes, - summary_generator=config.summary_generator, + operation_id_namespace=operation_id_namespace, item_serdes=config.item_serdes, nesting_type=config.nesting_type, ) - def get_iteration_name(self, index: int) -> str: - """Return custom branch name if the callable is a ParallelBranch with a name.""" - func = self.executables[index].func - if isinstance(func, ParallelBranch) and func.name is not None: - return func.name - return super().get_iteration_name(index) - - def execute_item(self, child_context, executable: Executable[Callable]) -> R: # noqa: PLR6301 - logger.debug("🔀 Processing parallel branch: %s", executable.index) - result: R = executable.func(child_context) - logger.debug("✅ Processed parallel branch: %s", executable.index) - return result - def parallel_handler( callables: Sequence[Callable | ParallelBranch], @@ -107,37 +100,20 @@ def parallel_handler( execution_state: ExecutionState, parallel_context: DurableContext, operation_identifier: OperationIdentifier, + operation_id_namespace: OperationIdNamespace, ) -> BatchResult[R]: """Execute multiple operations in parallel.""" - # Summary Generator Construction (matches TypeScript implementation): - # Construct the summary generator at the handler level, just like TypeScript does in parallel-handler.ts. - # This matches the pattern where handlers are responsible for configuring operation-specific behavior. - # - # See TypeScript reference: aws-durable-execution-sdk-js/src/handlers/parallel-handler/parallel-handler.ts (~line 112) + parallel_config: ParallelConfig = config or ParallelConfig() executor = ParallelExecutor.from_callables( callables, - config or ParallelConfig(summary_generator=ParallelSummaryGenerator()), + parallel_config, + operation_id_namespace=operation_id_namespace, ) checkpoint = execution_state.get_checkpoint_result( 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) - - -class ParallelSummaryGenerator: - def __call__(self, result: BatchResult) -> str: - fields = { - "totalCount": result.total_count, - "successCount": result.success_count, - "failureCount": result.failure_count, - "completionReason": result.completion_reason.value, - "status": result.status.value, - "startedCount": result.started_count, - "type": "ParallelResult", - } - - return json.dumps(fields) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py index baec099b..cb174441 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py @@ -8,6 +8,7 @@ import queue import threading import time +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from enum import Enum from threading import Lock @@ -304,6 +305,15 @@ def __init__( # Protects parent_to_children and parent_done self._parent_done_lock: Lock = Lock() + # Branch thread pools created by concurrency coordinators. A pool + # abandoned on early completion can still have branches running user + # code; close() joins every registered pool so no SDK-created thread + # outlives the invocation (a thread still running at handler return + # is frozen with the execution environment and resumes mid-flight + # during the next warm invocation). + self._branch_pools: list[ThreadPoolExecutor] = [] + self._branch_pools_lock: Lock = Lock() + # Dedup set so each operation's replay plugin hook fires at most once. # Replay status itself is tracked per-context on DurableContext; the # context decides WHEN to emit (only while replaying) and calls @@ -828,6 +838,18 @@ def stop_checkpointing(self) -> None: logger.debug("Signaling background thread to stop checkpointing") self._checkpointing_stopped.set() + def register_branch_pool(self, pool: ThreadPoolExecutor) -> None: + """Register a branch thread pool for joining at invocation end. + + Concurrency coordinators register their pool on creation. close() + joins every registered pool before stopping the checkpoint batcher, + so branches abandoned by early completion finish (or unwind on + OrphanedChildException at their next checkpoint attempt) before the + invocation returns. + """ + with self._branch_pools_lock: + self._branch_pools.append(pool) + def _collect_checkpoint_batch(self) -> list[QueuedOperation]: """Collect multiple checkpoint operations into a batch for API efficiency. @@ -980,6 +1002,27 @@ def _calculate_operation_size(queued_op: QueuedOperation) -> int: return len(serialized) def close(self): + """Release invocation-scoped resources. + + Joins still-running branch threads BEFORE stopping the checkpoint + batcher: a branch blocked on an in-flight synchronous checkpoint + needs the batcher alive to receive its response (a rejection for an + orphaned branch) and unwind. Stopping the batcher first would + deadlock that branch until the Lambda timeout. + + Drains registrations until no pools remain: a branch joined in one + batch can start a nested map or parallel operation, whose + coordinator registers a new pool mid-join. The lock is never held + across shutdown(wait=True), so those registrations do not block. + """ + while True: + with self._branch_pools_lock: + pools: list[ThreadPoolExecutor] = list(self._branch_pools) + self._branch_pools.clear() + if not pools: + break + for pool in pools: + pool.shutdown(wait=True) self.stop_checkpointing() def wrap_user_function( 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 d532d98d..0cf8f0ab 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,12 @@ """Tests for the concurrency module.""" +import hashlib 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,22 +15,26 @@ 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, + envelope_summary_generator, ) from aws_durable_execution_sdk_python.config import ( ChildConfig, CompletionConfig, MapConfig, + ParallelBranch, + ParallelConfig, NestingType, ) from aws_durable_execution_sdk_python.context import ( @@ -36,8 +42,11 @@ ExecutionContext, ) from aws_durable_execution_sdk_python.exceptions import ( + BackgroundThreadError, ChildContextError, + ValidationError, InvalidStateError, + OrphanedChildException, SuspendExecution, TimedSuspendExecution, ) @@ -48,10 +57,23 @@ OperationSubType, OperationType, ) +from aws_durable_execution_sdk_python.identifier import OperationIdNamespace + + from aws_durable_execution_sdk_python.operation.map import MapExecutor +from aws_durable_execution_sdk_python.operation.parallel import ( + ParallelExecutor, +) from aws_durable_execution_sdk_python.state import CheckpointedResult +class _StubNamespace(OperationIdNamespace): + """Test namespace producing readable ids matching checkpoint fixtures.""" + + def create_id_for_step(self, step: int) -> str: + return f"op_{step}" + + def test_batch_item_status_enum(): """Test BatchItemStatus enum values.""" assert BatchItemStatus.SUCCEEDED.value == "SUCCEEDED" @@ -453,14 +475,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 +569,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 = [ @@ -887,6 +604,7 @@ def execute_item(self, child_context, executable): name_prefix="test_", serdes=None, nesting_type=NestingType.NESTED, + operation_id_namespace=_StubNamespace(), ) assert executor_nested.nesting_type is NestingType.NESTED @@ -900,6 +618,7 @@ def execute_item(self, child_context, executable): name_prefix="test_", serdes=None, nesting_type=NestingType.FLAT, + operation_id_namespace=_StubNamespace(), ) assert executor_flat.nesting_type is NestingType.FLAT @@ -922,6 +641,7 @@ def execute_item(self, child_context, executable): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) assert executor.nesting_type is NestingType.NESTED @@ -947,6 +667,7 @@ def execute_item(self, child_context, executable): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) execution_state = Mock() @@ -963,137 +684,9 @@ def mock_run_in_child_context(func, name, config): result = executor.execute(execution_state, mock_run_in_child_context) 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) + # The pool is registered so ExecutionState.close() joins branches + # still running after early completion at invocation end. + execution_state.register_branch_pool.assert_called_once() def test_concurrent_executor_create_result_with_early_exit(): @@ -1130,13 +723,14 @@ def failure_callable(): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) execution_state = Mock() execution_state.create_checkpoint = Mock() executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context @@ -1173,10 +767,11 @@ def execute_item(self, child_context, executable): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context @@ -1214,12 +809,13 @@ def execute_item(self, child_context, executable): name_prefix="test_", serdes=None, nesting_type=NestingType.NESTED, + operation_id_namespace=_StubNamespace(), ) # Parent (executor) context is replaying. executor_context = Mock() executor_context.is_replaying = lambda: True - executor_context._create_step_id_for_logical_step = lambda *args: "branch-1" # noqa: SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "branch-1" executor_context._parent_id = "parent" # noqa: SLF001 # Brand-new branch: no checkpoint for the branch op. @@ -1256,11 +852,12 @@ def execute_item(self, child_context, executable): name_prefix="test_", serdes=None, nesting_type=NestingType.NESTED, + operation_id_namespace=_StubNamespace(), ) executor_context = Mock() executor_context.is_replaying = lambda: True - executor_context._create_step_id_for_logical_step = lambda *args: "branch-1" # noqa: SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "branch-1" executor_context._parent_id = "parent" # noqa: SLF001 # Replayed branch: a terminal checkpoint exists for the branch op. @@ -1303,11 +900,12 @@ def execute_item(self, child_context, executable): name_prefix="test_", serdes=None, nesting_type=NestingType.FLAT, + operation_id_namespace=_StubNamespace(), ) executor_context = Mock() executor_context.is_replaying = lambda: True - executor_context._create_step_id_for_logical_step = lambda *args: "branch-1" # noqa: SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "branch-1" executor_context._parent_id = "parent" # noqa: SLF001 child_context = Mock() @@ -1320,18 +918,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.""" @@ -1358,6 +944,7 @@ def failure_callable(): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) execution_state = Mock() @@ -1395,13 +982,14 @@ def execute_item(self, child_context, executable): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) execution_state = Mock() execution_state.create_checkpoint = Mock() executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context @@ -1443,13 +1031,14 @@ def execute_item(self, child_context, executable): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) execution_state = Mock() execution_state.create_checkpoint = Mock() executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context @@ -1490,13 +1079,14 @@ def execute_item(self, child_context, executable): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) execution_state = Mock() execution_state.create_checkpoint = Mock() executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context @@ -1549,6 +1139,7 @@ def execute_item(self, child_context, executable): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) execution_state = Mock() @@ -1563,7 +1154,7 @@ def checkpoint(*args, **kwargs): execution_state.create_checkpoint = Mock(side_effect=checkpoint) executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa: SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context @@ -1633,13 +1224,14 @@ def execute_item(self, child_context, executable): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) execution_state = Mock() execution_state.create_checkpoint = Mock() executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context @@ -1658,36 +1250,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, ) @@ -1699,120 +1276,42 @@ def execute_item(self, child_context, executable): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) - # 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_context = Mock() + executor_context._create_step_id_for_logical_step = lambda *args: "1" + 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.executables_with_state = [exe_state] + result = executor.execute(execution_state, executor_context) - result = executor.should_execution_suspend() + 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" - assert result.should_suspend - assert isinstance(result.exception, TimedSuspendExecution) - assert result.exception.scheduled_timestamp == future_time +def test_failed_branch_error_type_matches_across_execute_and_replay(): + """A failed branch's error_type must be identical on first run and replay. -def test_concurrent_executor_should_execution_suspend_indefinite(): - """Test should_execution_suspend with indefinite SUSPENDED state.""" + First run builds the item error in _create_result; replay reads the branch's + FAIL checkpoint (which records the raw escaping type). Both must surface the + raw type ("ValueError"), not the ChildContextError wrapper class name. + """ 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")] 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, - ) - - # Create executable with state in SUSPENDED (indefinite) - exe_state = ExecutableWithState(executables[0]) - exe_state.suspend() - - executor.executables_with_state = [exe_state] - - result = executor.should_execution_suspend() - - assert result.should_suspend - assert isinstance(result.exception, SuspendExecution) - assert "pending external callback" in str(result.exception) - - -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): - msg = "Test error" - raise ValueError(msg) - - def failure_callable(): - return "test" - - executables = [Executable(0, failure_callable)] - completion_config = CompletionConfig( - min_successful=1, - tolerated_failure_count=0, - 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, - ) - - 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) - - 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_failed_branch_error_type_matches_across_execute_and_replay(): - """A failed branch's error_type must be identical on first run and replay. - - First run builds the item error in _create_result; replay reads the branch's - FAIL checkpoint (which records the raw escaping type). Both must surface the - raw type ("ValueError"), not the ChildContextError wrapper class name. - """ - - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - msg = "Test error" - raise ValueError(msg) - - completion_config = CompletionConfig( - min_successful=1, - tolerated_failure_count=0, + tolerated_failure_count=0, tolerated_failure_percentage=None, ) @@ -1825,13 +1324,14 @@ def make_executor(): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) # First run: execute() runs the branch and builds the result live. execution_state = Mock() execution_state.create_checkpoint = Mock() executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa: SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context @@ -1849,7 +1349,7 @@ def make_executor(): replay_state = Mock() replay_state.get_checkpoint_result = Mock(return_value=failed_checkpoint) replay_context = Mock() - replay_context._create_step_id_for_logical_step = lambda *args: "1" # noqa: SLF001 + replay_context._create_step_id_for_logical_step = lambda *args: "1" replay_error = make_executor().replay(replay_state, replay_context).all[0].error assert replay_error is not None @@ -1859,190 +1359,6 @@ def make_executor(): assert live_error.type == replay_error.type -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() - - -def test_concurrent_executor_mixed_suspend_states(): - """Test should_execution_suspend with mixed suspend states.""" - - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" - - 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, - ) - - # 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 - - executor.executables_with_state = [exe_state1, exe_state2] - - result = executor.should_execution_suspend() - - # Should return timed suspend (earliest timestamp takes precedence) - assert result.should_suspend - assert isinstance(result.exception, TimedSuspendExecution) - - -def test_concurrent_executor_multiple_timed_suspends(): - """Test should_execution_suspend with multiple timed suspends to find earliest.""" - - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" - - 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, - ) - - # Create two with different timed suspends - exe_state1 = ExecutableWithState(executables[0]) - exe_state2 = ExecutableWithState(executables[1]) - - later_time = time.time() + 10 - earlier_time = time.time() + 5 - - exe_state1.suspend_with_timeout(later_time) - exe_state2.suspend_with_timeout(earlier_time) - - executor.executables_with_state = [exe_state1, exe_state2] - - result = executor.should_execution_suspend() - - # Should return the earlier timestamp - assert result.should_suspend - assert isinstance(result.exception, TimedSuspendExecution) - assert result.exception.scheduled_timestamp == earlier_time - - -def test_timer_scheduler_double_check_condition_race(): - """Test TimerScheduler double-check condition when heap changes between checks.""" - callback = Mock() - - with TimerScheduler(callback) as scheduler: - exe_state1 = ExecutableWithState(Executable(0, lambda: "test")) - exe_state2 = ExecutableWithState(Executable(1, lambda: "test")) - - exe_state1.suspend() - exe_state2.suspend() - - # Schedule first task with past time - past_time = time.time() - 1 - scheduler.schedule_resume(exe_state1, past_time) - - # Brief delay to let timer thread see the first task - time.sleep(0.05) - - # 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) - - # Wait for processing - time.sleep(0.2) - - assert callback.call_count >= 1 - - -def test_should_execution_suspend_earliest_timestamp_comparison(): - """Test should_execution_suspend timestamp comparison logic (line 554).""" - - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" - - 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, - ) - - 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 three executables with different suspend times - exe_state1 = ExecutableWithState(executables[0]) - exe_state2 = ExecutableWithState(executables[1]) - exe_state3 = ExecutableWithState(executables[2]) - - time1 = time.time() + 10 - time2 = time.time() + 5 # Earliest - time3 = time.time() + 15 - - exe_state1.suspend_with_timeout(time1) - exe_state2.suspend_with_timeout(time2) - exe_state3.suspend_with_timeout(time3) - - executor.executables_with_state = [exe_state1, exe_state2, exe_state3] - - result = executor.should_execution_suspend() - - assert result.should_suspend - assert isinstance(result.exception, TimedSuspendExecution) - assert result.exception.scheduled_timestamp == time2 - - def test_concurrent_executor_execute_with_failing_task(): """Test execute() with a task that fails using public execute method.""" @@ -2067,13 +1383,14 @@ def failure_callable(): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) execution_state = Mock() execution_state.create_checkpoint = Mock() executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context @@ -2085,27 +1402,6 @@ def failure_callable(): assert result.all[0].error.message == "Task failed" -def test_timer_scheduler_cannot_resume_branch(): - """Test TimerScheduler when exe_state cannot resume (434->433 branch).""" - callback = Mock() - - with TimerScheduler(callback) as scheduler: - exe_state = ExecutableWithState(Executable(0, lambda: "test")) - - # Set to completed state so can_resume returns False - exe_state.complete("done") - - # Schedule with past time - past_time = time.time() - 1 - scheduler.schedule_resume(exe_state, past_time) - - # Wait for processing - time.sleep(0.2) - - # Callback should not be called since can_resume is False - callback.assert_not_called() - - def test_create_result_no_failed_executables(): """Test when no executables are failed using public execute method.""" @@ -2131,13 +1427,14 @@ def success_callable(): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) execution_state = Mock() execution_state.create_checkpoint = Mock() executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" executor_context.create_child_context = lambda *args, **kwargs: Mock() result = executor.execute(execution_state, executor_context) @@ -2173,13 +1470,14 @@ def suspend_callable(): sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) execution_state = Mock() execution_state.create_checkpoint = Mock() executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context @@ -2190,1661 +1488,2246 @@ def suspend_callable(): # Tests for _create_result method match statement branches -def test_create_result_completed_branch(): - """Test _create_result with COMPLETED status branch.""" - - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" +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 + } - executables = [Executable(0, lambda: "test")] + # With started items, should infer MIN_SUCCESSFUL_REACHED 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, - ) + 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() - # Create executable with COMPLETED status - exe_state = ExecutableWithState(executables[0]) - exe_state.complete("test_result") - executor.executables_with_state = [exe_state] - result = executor._create_result() # noqa: SLF001 +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 + } - 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 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_create_result_failed_branch(): - """Test _create_result with FAILED status branch.""" +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", + } - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + # Should work without completion_config parameter + result = BatchResult.from_dict(data) + assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - executables = [Executable(0, lambda: "test")] - completion_config = CompletionConfig(min_successful=1) + # Should also work with None completion_config + result2 = BatchResult.from_dict(data, None) + assert result2.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - 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] +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 + + # 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 - result = executor._create_result() # noqa: SLF001 + # Test empty items - should be ALL_COMPLETED + batch = BatchResult.from_dict({"all": []}, CompletionConfig(1)) + assert batch.completion_reason == CompletionReason.ALL_COMPLETED - 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 +def test_operation_id_determinism_across_shuffles(): + """Test that operation_id depends on Executable.index, not execution order.""" -def test_create_result_pending_branch(): - """Test _create_result with PENDING status branch.""" + def index_based_function(index, ctx): + """Function that returns a result based on the executable index.""" + return f"result_for_index_{index}" class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + """Custom executor for testing operation_id determinism.""" - executables = [Executable(0, lambda: "test")] - completion_config = CompletionConfig(min_successful=1) + def execute_item(self, child_context, executable): + return executable.func(child_context) - 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 executables with specific indices using partial + num_executables = 50 + funcs = [partial(index_based_function, i) for i in range(num_executables)] - # 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] + # Track operation_id -> result associations + captured_associations = [] - result = executor._create_result() # noqa: SLF001 + 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 - 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() + completion_config = CompletionConfig(min_successful=num_executables) -def test_create_result_running_branch(): - """Test _create_result with RUNNING status branch.""" + # Run multiple times with different shuffle orders + associations_per_run = [] - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + for run in range(10): # Test 10 different shuffle orders + captured_associations.clear() - executables = [Executable(0, lambda: "test")] - completion_config = CompletionConfig(min_successful=1) + # 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) - 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 = 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, + operation_id_namespace=_StubNamespace(), + ) - # Create executable with RUNNING status - exe_state = ExecutableWithState(executables[0]) - future = Future() - exe_state.run(future) - executor.executables_with_state = [exe_state] + # Create executor context mock + executor_context = Mock() + executor_context._parent_id = "parent_123" # noqa SLF001 - result = executor._create_result() # noqa: SLF001 + def create_step_id(index): + return f"step_{index}" - 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 + executor_context._create_step_id_for_logical_step = create_step_id + def create_child_context(operation_id, *, is_virtual=False): + child_ctx = Mock() + child_ctx.state = execution_state + return child_ctx -def test_create_result_suspended_branch(): - """Test _create_result with SUSPENDED 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}" + with patch( + "aws_durable_execution_sdk_python.concurrency.executor.child_handler", + patched_child_handler, + ): + executor.execute(execution_state, executor_context) - executables = [Executable(0, lambda: "test")] - completion_config = CompletionConfig(min_successful=1) + associations_per_run.append(captured_associations.copy()) - executor = TestExecutor( - executables=executables, - max_concurrency=1, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=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) + ) + # 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) ) - # 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 +def test_concurrent_executor_replay_with_succeeded_operations(): + """Test ConcurrentExecutor replay method with succeeded operations.""" - 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 + 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, + operation_id_namespace=_StubNamespace(), + ) -def test_create_result_suspended_with_timeout_branch(): - """Test _create_result with SUSPENDED_WITH_TIMEOUT status branch.""" + # 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" + ) - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + 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 - executables = [Executable(0, lambda: "test")] - completion_config = CompletionConfig(min_successful=1) + mock_execution_state.get_checkpoint_result = mock_get_checkpoint_result - 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 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 ) - # 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] + # 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._create_result() # noqa: SLF001 + result = executor.replay(mock_execution_state, mock_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 + 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_create_result_mixed_statuses(): - """Test _create_result with mixed executable statuses covering all branches.""" +def test_concurrent_executor_replay_with_failed_operations(): + """Test ConcurrentExecutor replay method with failed operations.""" - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + def func1(ctx, item, idx, items): + return f"result_{item}" - 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 - ] - completion_config = CompletionConfig(min_successful=1) + items = ["a"] + config = MapConfig() - executor = TestExecutor( - executables=executables, - max_concurrency=6, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, + executor = MapExecutor.from_items( + items=items, + func=func1, + config=config, + operation_id_namespace=_StubNamespace(), ) - # Create executables with different statuses - exe_states = [ExecutableWithState(exe) for exe in executables] + # Mock execution state with failed operation + mock_execution_state = Mock() - # COMPLETED - exe_states[0].complete("completed_result") + 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 - # FAILED - exe_states[1].fail(RuntimeError("Test failure")) + mock_execution_state.get_checkpoint_result = mock_get_checkpoint_result - # PENDING (default state, no change needed) + # Mock executor context + mock_executor_context = Mock() + mock_executor_context._create_step_id_for_logical_step = Mock(return_value="op_1") - # RUNNING - future = Future() - exe_states[3].run(future) + result = executor.replay(mock_execution_state, mock_executor_context) - # SUSPENDED - exe_states[4].suspend() + assert isinstance(result, BatchResult) + assert len(result.all) == 1 + assert result.all[0].status == BatchItemStatus.FAILED + assert result.all[0].error is not None - # SUSPENDED_WITH_TIMEOUT - exe_states[5].suspend_with_timeout(time.time() + 10) - executor.executables_with_state = exe_states +def test_concurrent_executor_replay_with_replay_children(): + """Test ConcurrentExecutor replay method when children need re-execution.""" - result = executor._create_result() # noqa: SLF001 + def func1(ctx, item, idx, items): + return f"result_{item}" - assert len(result.all) == 6 + items = ["a"] + config = MapConfig() - # Check COMPLETED -> SUCCEEDED - assert result.all[0].status == BatchItemStatus.SUCCEEDED - assert result.all[0].result == "completed_result" - assert result.all[0].error is None + executor = MapExecutor.from_items( + items=items, + func=func1, + config=config, + operation_id_namespace=_StubNamespace(), + ) - # 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. - assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED + # 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 -def test_create_result_multiple_completed(): - """Test _create_result with multiple COMPLETED executables.""" + mock_execution_state.get_checkpoint_result = mock_get_checkpoint_result - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return f"result_{executable.index}" + # Mock executor context + mock_executor_context = Mock() + mock_executor_context._create_step_id_for_logical_step = Mock(return_value="op_1") - executables = [ - Executable(0, lambda: "test0"), - Executable(1, lambda: "test1"), - Executable(2, lambda: "test2"), - ] - completion_config = CompletionConfig(min_successful=3) + # 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) - executor = TestExecutor( - executables=executables, - max_concurrency=3, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) + 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": [], + }, + } - # 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") + item = BatchItem.from_dict(data) - executor.executables_with_state = exe_states + assert item.index == 3 + assert item.status == BatchItemStatus.FAILED + assert item.error.type == "ValueError" + assert item.error.message == "bad value" - 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_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)) -def test_create_result_multiple_failed(): - """Test _create_result with multiple FAILED executables.""" + 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 - 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"), - ] - completion_config = CompletionConfig(min_successful=1) +def test_batch_result_empty_list(): + """Test BatchResult serialization with empty items list.""" + result = BatchResult(all=[], completion_reason=CompletionReason.ALL_COMPLETED) - executor = TestExecutor( - executables=executables, - max_concurrency=3, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) + serialized = json.dumps(result.to_dict()) + deserialized = BatchResult.from_dict(json.loads(serialized)) - # 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")) + 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"}, + } - executor.executables_with_state = exe_states + result = BatchResult( + all=[BatchItem(0, BatchItemStatus.SUCCEEDED, result=complex_result)], + completion_reason=CompletionReason.ALL_COMPLETED, + ) - result = executor._create_result() # noqa: SLF001 + serialized = json.dumps(result.to_dict()) + deserialized = BatchResult.from_dict(json.loads(serialized)) - 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 + assert deserialized.all[0].result == complex_result + assert deserialized.all[0].result["users"][0]["name"] == "Alice" -def test_create_result_multiple_started_states(): - """Test _create_result with multiple executables in STARTED states.""" +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): - return f"result_{executable.index}" + if executable.index == 0: + # fail one task + raise Exception("boom") # noqa EM101 TRY002 + return f"ok_{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 - ] - completion_config = CompletionConfig(min_successful=1) + # 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=4, + max_concurrency=2, completion_config=completion_config, sub_type_top="TOP", sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) - # Create executables with different STARTED states - exe_states = [ExecutableWithState(exe) for exe in executables] + execution_state = Mock() + execution_state.create_checkpoint = Mock() + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda *args: "1" + 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 - # PENDING (default state) - # RUNNING - future = Future() - exe_states[1].run(future) +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}" - # SUSPENDED - exe_states[2].suspend() + items = list(range(100)) + config = MapConfig( + max_concurrency=10, + completion_config=CompletionConfig( + min_successful=99, tolerated_failure_count=1 + ), + ) - # SUSPENDED_WITH_TIMEOUT - exe_states[3].suspend_with_timeout(time.time() + 5) + executor = MapExecutor.from_items( + items=items, + func=task_func, + config=config, + operation_id_namespace=_StubNamespace(), + ) - executor.executables_with_state = exe_states + execution_state = Mock() + execution_state.create_checkpoint = Mock() + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda *args: "1" + 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) == 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 + # 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_empty_executables(): - """Test _create_result with no executables.""" +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}" + 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), + ] - executables = [] - completion_config = CompletionConfig(min_successful=0) + 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", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) - executor.executables_with_state = [] + execution_state = Mock() + execution_state.create_checkpoint = Mock() + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda idx: f"step_{idx}" + executor_context._parent_id = "parent" # noqa: SLF001 - result = executor._create_result() # 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 - assert len(result.all) == 0 - assert result.completion_reason == CompletionReason.ALL_COMPLETED + 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_timer_scheduler_future_time_condition_false(): - """Test TimerScheduler when scheduled time is in future (434->433 branch).""" - callback = Mock() + # 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" - with TimerScheduler(callback) as scheduler: - exe_state = ExecutableWithState(Executable(0, lambda: "test")) - exe_state.suspend() + # Result should show MIN_SUCCESSFUL_REACHED + assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - # Schedule with future time so condition will be False - future_time = time.time() + 10 - scheduler.schedule_resume(exe_state, future_time) + # Fast branch should succeed + assert result.all[0].status == BatchItemStatus.SUCCEEDED + assert result.all[0].result == "fast_result" - # Wait briefly for timer thread to check and find condition False - time.sleep(0.1) + # Slow branch should be marked as STARTED (incomplete) + assert result.all[1].status == BatchItemStatus.STARTED - # Callback should not be called since time is in future - callback.assert_not_called() + # Verify counts + assert result.success_count == 1 + assert result.failure_count == 0 + assert result.started_count == 1 + assert result.total_count == 2 -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_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 executable.func() + + operation_tracker = Mock() + + 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" + + executables = [ + Executable(0, fast_branch), + Executable(1, slow_branch), + ] - # 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() + executor = TestExecutor( + executables=executables, + max_concurrency=2, + completion_config=completion_config, + sub_type_top="TOP", + sub_type_iteration="ITER", + name_prefix="test_", + serdes=None, + operation_id_namespace=_StubNamespace(), + ) + 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}" + executor_context._parent_id = "parent" # noqa: SLF001 + executor_context.create_child_context = lambda op_id, *, is_virtual=False: Mock( + state=execution_state + ) -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 - } + result = executor.execute(execution_state, executor_context) - # 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() + # 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" + ) -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", - } + # Result should show MIN_SUCCESSFUL_REACHED + assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - # Should work without completion_config parameter - result = BatchResult.from_dict(data) + # 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_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 executable.func() + + 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=2, + completion_config=completion_config, + sub_type_top="TOP", + sub_type_iteration="ITER", + name_prefix="test_", + serdes=None, + operation_id_namespace=_StubNamespace(), + ) + + 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}" + executor_context._parent_id = "parent" # noqa: SLF001 + executor_context.create_child_context = lambda op_id, *, is_virtual=False: Mock( + state=execution_state + ) + + 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" + ) + + # Result should show MIN_SUCCESSFUL_REACHED assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - # Should also work with None completion_config - result2 = BatchResult.from_dict(data, None) - assert result2.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_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_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_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_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 + + +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_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 + + +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_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 + + +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) + ), + ] + 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_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_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 + + +# endregion Completion Reason Inference Tests + +# region Virtual-context wire-format tests + + +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. + + 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) + + 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 + + # 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. + ) - # 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 + 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, + operation_id_namespace=_StubNamespace(), + ) - # Test empty items - should be ALL_COMPLETED - batch = BatchResult.from_dict({"all": []}, CompletionConfig(1)) - assert batch.completion_reason == CompletionReason.ALL_COMPLETED + executor._execute_item_in_child_context(executor_context, executables[0]) # noqa: SLF001 + # 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_operation_id_determinism_across_shuffles(): - """Test that operation_id depends on Executable.index, not execution order.""" - def index_based_function(index, ctx): - """Function that returns a result based on the executable index.""" - return f"result_for_index_{index}" +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.""" class TestExecutor(ConcurrentExecutor): - """Custom executor for testing operation_id determinism.""" - def execute_item(self, child_context, executable): + self.last_child_context = child_context 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() + execution_state.wrap_user_function = lambda func, *args, **kwargs: func - completion_config = CompletionConfig(min_successful=num_executables) - - # Run multiple times with different shuffle orders - associations_per_run = [] + 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 - for run in range(10): # Test 10 different shuffle orders - captured_associations.clear() + 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, + ) - # 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) + 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, + operation_id_namespace=_StubNamespace(), + ) - 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._execute_item_in_child_context(executor_context, executables[0]) # noqa: SLF001 - # Create executor context mock - executor_context = Mock() - executor_context._parent_id = "parent_123" # noqa SLF001 + # 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 create_step_id(index): - return f"step_{index}" - executor_context._create_step_id_for_logical_step = create_step_id # noqa SLF001 +# endregion Virtual-context wire-format tests - def create_child_context(operation_id, *, is_virtual=False): - child_ctx = Mock() - child_ctx.state = execution_state - return child_ctx - executor_context.create_child_context = create_child_context +def test_flat_mode_produces_deterministic_step_ids_across_runs(): + """Step ids and inner parent_ids must be deterministic under FLAT mode. - with patch( - "aws_durable_execution_sdk_python.concurrency.executor.child_handler", - patched_child_handler, - ): - executor.execute(execution_state, executor_context) + 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. + """ - associations_per_run.append(captured_associations.copy()) + class TestExecutor(ConcurrentExecutor): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.captured = [] - # 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 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() -def test_concurrent_executor_replay_with_succeeded_operations(): - """Test ConcurrentExecutor replay method with succeeded operations.""" + 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 func1(ctx, item, idx, items): - return f"result_{item}" + 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", + ) - items = ["a", "b"] - config = MapConfig() + 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, + operation_id_namespace=_StubNamespace(), + ) + executor.execute(execution_state, executor_context) + return executor.captured - executor = MapExecutor.from_items( - items=items, - func=func1, - config=config, - ) + run_a = make_run() + run_b = make_run() - # 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" + # 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) - 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 +# region Branch state transitions - 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 - ) +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 - # 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) +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 - 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_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 test_concurrent_executor_replay_with_failed_operations(): - """Test ConcurrentExecutor replay method with failed operations.""" - def func1(ctx, item, idx, items): - return f"result_{item}" +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 - items = ["a"] - config = MapConfig() - executor = MapExecutor.from_items( - items=items, - func=func1, - config=config, - ) +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 - # 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 +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 - 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 +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() - result = executor.replay(mock_execution_state, mock_executor_context) + suspended = Branch(Executable(1, lambda: "test")) + suspended.start() + suspended.suspend() + with pytest.raises(InvalidStateError, match="Cannot start branch 1"): + suspended.start() - 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_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 -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}" +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 - items = ["a"] - config = MapConfig() - executor = MapExecutor.from_items( - items=items, - func=func1, - config=config, - ) +# endregion Branch state transitions - # 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 +# region CompletionPolicy - 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 +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 - # 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_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 -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": [], - }, - } +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) - 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_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_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, +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) - 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_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) -def test_batch_result_empty_list(): - """Test BatchResult serialization with empty items list.""" - result = BatchResult(all=[], completion_reason=CompletionReason.ALL_COMPLETED) +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) - 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_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_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"}, - } +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) - result = BatchResult( - all=[BatchItem(0, BatchItemStatus.SUCCEEDED, result=complex_result)], - completion_reason=CompletionReason.ALL_COMPLETED, + +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 ) - 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_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 -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.""" +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 + ) - 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, +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 + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda idx: f"step_{idx}" + executor_context._parent_id = "parent" # noqa: SLF001 + executor_context.create_child_context = lambda op_id, *, is_virtual=False: Mock( + state=execution_state ) + return execution_state, executor_context - executor = TestExecutor( + +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=2, - completion_config=completion_config, + max_concurrency=max_concurrency, + completion_config=completion_config or CompletionConfig(), sub_type_top="TOP", sub_type_iteration="ITER", name_prefix="test_", serdes=None, + operation_id_namespace=_StubNamespace(), ) - 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_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" -def test_executor_terminates_quickly_when_impossible_to_succeed(): - """Test that executor terminates when min_successful becomes impossible.""" - executed_count = {"value": 0} + executables = [Executable(i, tracked) for i in range(6)] + executor = _make_executor(executables, max_concurrency=2) + execution_state, executor_context = _make_executor_mocks() - 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}" + result = executor.execute(execution_state, executor_context) - items = list(range(100)) - config = MapConfig( - max_concurrency=10, - completion_config=CompletionConfig( - min_successful=99, tolerated_failure_count=1 - ), - ) + assert peak <= 2 + assert result.success_count == 6 + assert result.completion_reason is CompletionReason.ALL_COMPLETED - 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 +def test_suspended_branch_holds_concurrency_slot(): + """A suspended branch keeps its slot: pending items must not start (issue #279).""" + started: list[int] = [] - result = executor.execute(execution_state, executor_context) + def suspending(index: int): + started.append(index) + msg = "awaiting callback" + raise SuspendExecution(msg) - # 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 - ) + 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() + with pytest.raises(SuspendExecution): + executor.execute(execution_state, executor_context) -def test_executor_exits_early_with_min_successful(): - """Test that parallel exits immediately when min_successful is reached without waiting for other branches.""" + assert sorted(started) == [0, 1] - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return executable.func() - execution_times = [] +def test_slot_freed_on_completion_admits_next_item(): + """A terminal completion frees a slot; a suspension does not.""" + started: list[int] = [] - def fast_branch(): - execution_times.append(("fast", time.time())) - return "fast_result" + def completing(index: int): + started.append(index) + return f"result_{index}" - 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" + def suspending(index: int): + started.append(index) + msg = "awaiting callback" + raise SuspendExecution(msg) executables = [ - Executable(0, fast_branch), - Executable(1, slow_branch), + 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() - completion_config = CompletionConfig(min_successful=1) + with pytest.raises(SuspendExecution): + executor.execute(execution_state, executor_context) - executor = TestExecutor( - executables=executables, - max_concurrency=2, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) + assert sorted(started) == [0, 1, 2] - 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_branches_start_in_index_order(): + started: list[int] = [] + + def record(index: int): + started.append(index) + return index - executor_context.create_child_context = create_child_context + 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() - start_time = time.time() result = executor.execute(execution_state, executor_context) - elapsed_time = time.time() - start_time - # 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" + assert started == [0, 1, 2, 3] + assert result.success_count == 4 - # Result should show MIN_SUCCESSFUL_REACHED - assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED - # Fast branch should succeed - assert result.all[0].status == BatchItemStatus.SUCCEEDED - assert result.all[0].result == "fast_result" +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} - # Slow branch should be marked as STARTED (incomplete) - assert result.all[1].status == BatchItemStatus.STARTED + 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" - # Verify counts - assert result.success_count == 1 - assert result.failure_count == 0 - assert result.started_count == 1 - assert result.total_count == 2 + def slow_sibling(): + resumed.wait(timeout=5.0) + return "sibling" + executables = [Executable(0, timed_then_ok), Executable(1, slow_sibling)] + executor = _make_executor(executables, max_concurrency=2) + execution_state, executor_context = _make_executor_mocks() -def test_executor_returns_with_incomplete_branches(): - """Test that executor returns when min_successful is reached, leaving other branches incomplete.""" + result = executor.execute(execution_state, executor_context) - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return executable.func() + 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() - operation_tracker = Mock() - def fast_branch(): - operation_tracker.fast_executed() - return "fast_result" +def test_all_timed_suspended_parent_suspends_with_earliest_timestamp(): + earliest = time.time() + 100 + latest = time.time() + 200 - def slow_branch(): - operation_tracker.slow_started() - time.sleep(2) # Long sleep - operation_tracker.slow_completed() - return "slow_result" + def suspend_at(timestamp: float): + msg = "retry later" + raise TimedSuspendExecution(msg, timestamp) executables = [ - Executable(0, fast_branch), - Executable(1, slow_branch), + 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() - completion_config = CompletionConfig(min_successful=1) + with pytest.raises(TimedSuspendExecution) as exc_info: + executor.execute(execution_state, executor_context) - executor = TestExecutor( - executables=executables, - max_concurrency=2, - completion_config=completion_config, - sub_type_top="TOP", - sub_type_iteration="ITER", - name_prefix="test_", - serdes=None, - ) + assert exc_info.value.scheduled_timestamp == earliest - 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.execute(execution_state, executor_context) +def test_never_started_branches_omitted_from_result(): + """Branches that never started are omitted from the batch result (JS parity).""" + release = threading.Event() - # Verify fast branch executed - assert operation_tracker.fast_executed.call_count == 1 + def fast(): + return "fast" - # 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 slow(): + release.wait(timeout=5.0) + return "slow" + + executables = [ + Executable(0, fast), + Executable(1, slow), + Executable(2, fast), + Executable(3, fast), + ] + executor = _make_executor( + executables, + max_concurrency=2, + completion_config=CompletionConfig(min_successful=1), ) + execution_state, executor_context = _make_executor_mocks() - # Result should show MIN_SUCCESSFUL_REACHED - assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED + try: + result = executor.execute(execution_state, executor_context) + finally: + release.set() - # Verify counts - one succeeded, one incomplete + assert result.completion_reason is CompletionReason.MIN_SUCCESSFUL_REACHED assert result.success_count == 1 - assert result.failure_count == 0 assert result.started_count == 1 assert result.total_count == 2 + indexes = {item.index for item in result.all} + assert indexes == {0, 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_tolerance_breach_stops_submission_of_pending_items(): + """Fail-fast: a failure with no completion criteria stops further submissions.""" + started: list[int] = [] - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - return executable.func() + def failing(index: int): + started.append(index) + msg = "boom" + raise ValueError(msg) - slow_branch_mock = Mock() + 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() - def fast_func(): - return "fast" + result = executor.execute(execution_state, executor_context) - def slow_func(): - time.sleep(3) # Sleep - slow_branch_mock.completed() # Should not be called before executor returns - return "slow" + assert started == [0] + assert result.completion_reason is CompletionReason.FAILURE_TOLERANCE_EXCEEDED + assert result.failure_count == 1 + assert result.total_count == 1 - executables = [Executable(0, fast_func), Executable(1, slow_func)] - completion_config = CompletionConfig(min_successful=1) - executor = TestExecutor( +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") + ), + } + + class OutcomeExecutor(ConcurrentExecutor): + def execute_item(self, child_context, executable): + return outcomes[executable.index]() + + executables = [Executable(i, outcomes[i]) for i in range(5)] + execution_state, executor_context = _make_executor_mocks() + executor = OutcomeExecutor( executables=executables, - max_concurrency=2, - completion_config=completion_config, + max_concurrency=5, + completion_config=CompletionConfig(), sub_type_top="TOP", sub_type_iteration="ITER", name_prefix="test_", serdes=None, - ) + operation_id_namespace=_StubNamespace(), + ) + + 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 ChildContextError before + # they reach the worker. + assert isinstance(collected[1].error, ChildContextError) + 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], + } + ), + json.dumps( + { + "totalCount": "3", + "completionReason": "ALL_COMPLETED", + "startedIndexes": [1], + } + ), + ], +) +def test_completion_record_rejects_payloads_without_record(payload): + assert CompletionRecord.from_summary_payload(payload) is None + + +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]} + few_completed = [ + BatchItem(0, BatchItemStatus.SUCCEEDED, result="a"), + BatchItem(1, BatchItemStatus.STARTED), + BatchItem(2, BatchItemStatus.STARTED), + ] + assert CompletionRecord.summary_index_fields(few_completed) == { + "completedIndexes": [0] + } + + +def _make_replay_mocks(branch_checkpoints: dict): + """Mocks for replay: op id 'op_{index}', checkpoints per index.""" execution_state = Mock() - execution_state.create_checkpoint = Mock() - execution_state.wrap_user_function = lambda func, *args, **kwargs: func + execution_state.durable_execution_arn = ( + "arn:aws:durable:us-east-1:123456789012:execution/test" + ) + + 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"step_{idx}" # noqa: SLF001 + executor_context._create_step_id_for_logical_step = lambda idx: f"op_{idx}" executor_context._parent_id = "parent" # noqa: SLF001 - executor_context.create_child_context = lambda op_id, *, is_virtual=False: Mock( - state=execution_state + child_context = Mock() + child_context.state = execution_state + executor_context.create_child_context = Mock(return_value=child_context) + return execution_state, executor_context + + +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 + + +def test_replay_round_trip_matches_live_result(): + """execute() -> summary -> replay() reconstructs the exact live result.""" + + def succeed(index: int): + return f"live_{index}" + + def suspend(): + msg = "awaiting callback" + raise SuspendExecution(msg) + + 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) - result = executor.execute(execution_state, executor_context) + 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 = envelope_summary_generator("MapResult", None)(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 should have returned before slow branch completed - assert not slow_branch_mock.completed.called, ( - "Executor should return before slow branch completes" + 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" + + +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"')} ) - # Result should show MIN_SUCCESSFUL_REACHED - assert result.completion_reason == CompletionReason.MIN_SUCCESSFUL_REACHED + replayed: BatchResult = executor.replay(replay_state, replay_context) - # Verify counts - assert result.success_count == 1 - assert result.failure_count == 0 - assert result.started_count == 1 - assert result.total_count == 2 + assert [(i.index, i.status) for i in replayed.all] == [ + (0, BatchItemStatus.SUCCEEDED), + (1, BatchItemStatus.STARTED), + (2, BatchItemStatus.STARTED), + ] -# region TimerScheduler edge cases with exact same reschedule time +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"'), + } + ) + replayed: BatchResult = executor.replay( + replay_state, replay_context, top_checkpoint + ) -def test_timer_scheduler_same_timestamp_with_counter_tiebreaker(): - """ - Test that scheduling two tasks with the exact same resume_time works. + assert [(i.index, i.status) for i in replayed.all] == [ + (0, BatchItemStatus.SUCCEEDED), + (1, BatchItemStatus.SUCCEEDED), + ] - 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_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({}) - # Use the exact same timestamp for both - same_timestamp = time.time() + 10.0 + replayed: BatchResult = executor.replay( + replay_state, replay_context, top_checkpoint + ) - # Both schedules should work fine now - scheduler.schedule_resume(exe_state1, same_timestamp) - scheduler.schedule_resume(exe_state2, same_timestamp) + assert [(i.index, i.status) for i in replayed.all] == [(0, BatchItemStatus.STARTED)] + assert replayed.completion_reason is CompletionReason.ALL_COMPLETED - # Verify both are in the heap - assert len(scheduler._pending_resumes) == 2 # noqa: SLF001 - # 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 +def test_replay_flat_reconstructs_terminal_statuses_by_reexecution(): + """FLAT branches have no checkpoints: re-execution discriminates outcomes.""" + def succeed(): + return "flat_ok" -def test_timer_scheduler_multiple_same_timestamps(): - """ - Test that scheduling many tasks with the same timestamp works correctly. + def fail(): + msg = "flat boom" + raise ValueError(msg) - Verifies FIFO ordering is maintained when multiple tasks have identical timestamps. - """ - resubmit_callback = Mock() + 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, + operation_id_namespace=_StubNamespace(), + ) + 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 - with TimerScheduler(resubmit_callback) as scheduler: - same_timestamp = time.time() + 10.0 + replayed: BatchResult = executor.replay( + replay_state, replay_context, top_checkpoint + ) - # 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) - ] + 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.all[1].error.type == "ValueError" + assert replayed.completion_reason is CompletionReason.ALL_COMPLETED - for exe_state in exe_states: - scheduler.schedule_resume(exe_state, same_timestamp) - # All should be scheduled successfully - assert len(scheduler._pending_resumes) == 10 # noqa: SLF001 +# endregion Completion record and replay reconstruction - # Verify the heap maintains proper ordering - # The first item should have counter 0 - assert scheduler._pending_resumes[0][1] == 0 # noqa: SLF001 +def test_create_result_uses_captured_completion_reason(): + """The reason captured at the completion decision wins over recount.""" + executables = [Executable(0, lambda: "x"), Executable(1, lambda: "y")] + executor = _make_executor(executables, max_concurrency=2) + completed = Branch(executables[0]) + completed.start() + completed.complete("done") + failed = Branch(executables[1]) + failed.start() + failed.fail(ValueError("boom")) + executor.branches = [completed, failed] -def test_timer_scheduler_counter_increments(): - """Test that the schedule counter increments correctly.""" - resubmit_callback = Mock() + result: BatchResult = executor._create_result( # noqa: SLF001 + CompletionReason.MIN_SUCCESSFUL_REACHED + ) - 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")) + assert result.completion_reason is CompletionReason.MIN_SUCCESSFUL_REACHED - # 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) - # Counter should have incremented to 3 - assert scheduler._schedule_counter == 3 # noqa: SLF001 +def test_branch_base_exception_propagates_to_caller(): + """A BaseException in a branch propagates from execute(): no hang, no conversion.""" + def exits(): + raise SystemExit(3) -def test_timer_scheduler_fifo_ordering_with_same_timestamp(): - """ - Test that FIFO ordering is maintained when timestamps are equal. + executables = [Executable(0, exits)] + executor = _make_executor(executables, max_concurrency=1) + execution_state, executor_context = _make_executor_mocks() + + with pytest.raises(SystemExit): + executor.execute(execution_state, executor_context) - 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. + +def test_background_thread_error_propagates_to_caller(): + """A fatal checkpoint-subsystem failure in a branch reaches the calling thread. + + BackgroundThreadError must not be recorded as an ordinary branch failure: + tolerance logic must not absorb it and the parent must not attempt + further checkpoints. """ - results = [] - resubmit_callback = Mock( - side_effect=lambda batch: results.extend(exe.index for exe in batch) - ) - with TimerScheduler(resubmit_callback) as scheduler: - # Use a past timestamp so they trigger immediately - past_time = time.time() - 0.1 + def checkpoint_dead(): + msg = "background checkpoint thread failed" + raise BackgroundThreadError(msg, RuntimeError("worker died")) - exe_state1 = ExecutableWithState(Executable(0, lambda: "first")) - exe_state2 = ExecutableWithState(Executable(1, lambda: "second")) - exe_state3 = ExecutableWithState(Executable(2, lambda: "third")) + executables = [Executable(0, checkpoint_dead), Executable(1, lambda: "ok")] + executor = _make_executor( + executables, + max_concurrency=2, + completion_config=CompletionConfig(tolerated_failure_count=5), + ) + execution_state, executor_context = _make_executor_mocks() - # Make them all resumable - exe_state1.suspend() - exe_state2.suspend() - exe_state3.suspend() + with pytest.raises(BackgroundThreadError): + 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) - # Wait for timer thread to process them - time.sleep(0.3) +def test_unlimited_concurrency_starts_all_items(): + """max_concurrency=None keeps the previous start-everything behavior.""" + barrier = threading.Barrier(3, timeout=5.0) - # Verify FIFO order - they should be resubmitted in order 0, 1, 2 - assert results == [0, 1, 2] + def synchronized(): + barrier.wait() + return "ok" + executables = [Executable(i, synchronized) for i in range(3)] + executor = _make_executor(executables, max_concurrency=None) + execution_state, executor_context = _make_executor_mocks() -# endregion TimerScheduler edge cases with exact same reschedule time + result = executor.execute(execution_state, executor_context) + assert result.success_count == 3 -# region Completion Reason Inference Tests (from_items) +# endregion In-flight window semantics -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 +# region summary envelope -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 _live_result_for_envelope() -> BatchResult: + return BatchResult( + [ + BatchItem(0, BatchItemStatus.SUCCEEDED, "r0"), + BatchItem(1, BatchItemStatus.STARTED), + ], + CompletionReason.MIN_SUCCESSFUL_REACHED, + ) -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 +def test_envelope_default_fields_and_order(): + """Without a custom generator the envelope has record + view fields, no summary.""" + payload = envelope_summary_generator("MapResult", None)(_live_result_for_envelope()) + data = json.loads(payload) + assert data == { + "type": "MapResult", + "totalCount": 2, + "completionReason": "MIN_SUCCESSFUL_REACHED", + "startedIndexes": [1], + "startedCount": 1, + "successCount": 1, + "failureCount": 0, + "status": "SUCCEEDED", + } + assert list(data.keys())[:3] == ["type", "totalCount", "completionReason"] -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_envelope_preserves_custom_summary_verbatim(): + """The customer generator output is stored verbatim under 'summary'.""" + custom = "plain text, not JSON: 100% " + payload = envelope_summary_generator("ParallelResult", lambda _r: custom)( + _live_result_for_envelope() + ) + data = json.loads(payload) + assert data["summary"] == custom + assert data["type"] == "ParallelResult" + record = CompletionRecord.from_summary_payload(payload) + assert record is not None + assert record.started_total == 2 + assert record.started_indexes == frozenset({1}) -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 + +def test_envelope_coerces_non_string_summary(): + """A generator returning a non-str is coerced so the payload stays a flat envelope.""" + payload = envelope_summary_generator("MapResult", lambda _r: {"not": "a str"})( + _live_result_for_envelope() + ) + data = json.loads(payload) + assert isinstance(data["summary"], str) -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_envelope_propagates_raising_generator(): + """An exception from the customer generator propagates to the caller.""" + def boom(_result): + msg = "customer bug" + raise ValueError(msg) -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 + with pytest.raises(ValueError, match="customer bug"): + envelope_summary_generator("MapResult", boom)(_live_result_for_envelope()) -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) - ), - ] - 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_envelope_never_truncates_summary(): + """The customer summary is checkpointed as provided regardless of size.""" + big = "x" * 500_000 + payload = envelope_summary_generator("MapResult", lambda _r: big)( + _live_result_for_envelope() + ) + data = json.loads(payload) + assert data["summary"] == big -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_replay_reads_record_from_envelope_with_custom_summary(): + """Replay obeys the record even when a custom summary is present.""" + executables = [Executable(0, lambda: "x"), Executable(1, lambda: "x")] + live = _live_result_for_envelope() + payload = envelope_summary_generator("MapResult", lambda _r: "digest")(live) + top_checkpoint = Mock() + top_checkpoint.result = payload + replay_state, replay_context = _make_replay_mocks( + {"op_0": _succeeded_checkpoint('"replayed_0"')} + ) + replayed = _make_executor( + executables, + max_concurrency=2, + completion_config=CompletionConfig(min_successful=1), + ).replay(replay_state, replay_context, top_checkpoint) -def test_from_items_all_succeeded(): - """Validates: All items succeeded.""" - items = [ - BatchItem(0, BatchItemStatus.SUCCEEDED, result="ok1"), - BatchItem(1, BatchItemStatus.SUCCEEDED, result="ok2"), + assert [(i.index, i.status) for i in replayed.all] == [ + (0, BatchItemStatus.SUCCEEDED), + (1, BatchItemStatus.STARTED), ] - result = BatchResult.from_items(items, completion_config=None) - assert result.completion_reason == CompletionReason.ALL_COMPLETED - assert result.success_count == 2 + assert replayed.completion_reason is CompletionReason.MIN_SUCCESSFUL_REACHED + + +def test_record_parses_pre_envelope_flat_format(): + """Payloads written before the envelope (record keys at top level) still parse.""" + flat = json.dumps( + { + "totalCount": 3, + "successCount": 2, + "failureCount": 0, + "completionReason": "MIN_SUCCESSFUL_REACHED", + "status": "SUCCEEDED", + "type": "MapResult", + "startedIndexes": [2], + } + ) + record = CompletionRecord.from_summary_payload(flat) + assert record is not None + assert record.started_total == 3 + assert record.started_indexes == frozenset({2}) + + +@pytest.mark.parametrize( + "payload", + [ + '{"totalCount": 3, "completionReason": "ALL_COMPLETED", "startedIndexes": [[]]}', + '{"totalCount": 3, "completionReason": "ALL_COMPLETED", "startedIndexes": ["1"]}', + '{"totalCount": 3, "completionReason": "ALL_COMPLETED", "startedIndexes": [true]}', + '{"totalCount": 3, "completionReason": "ALL_COMPLETED", "startedIndexes": [3]}', + '{"totalCount": 3, "completionReason": "ALL_COMPLETED", "startedIndexes": [-1]}', + '{"totalCount": 3, "completionReason": "ALL_COMPLETED", "completedIndexes": [null]}', + '{"totalCount": true, "completionReason": "ALL_COMPLETED", "startedIndexes": []}', + '{"totalCount": -1, "completionReason": "ALL_COMPLETED", "startedIndexes": []}', + ], +) +def test_record_rejects_malformed_index_fields(payload): + """Malformed or out-of-range index fields return None instead of raising.""" + assert CompletionRecord.from_summary_payload(payload) is None -# endregion Completion Reason Inference Tests +# endregion summary envelope -# region Virtual-context wire-format tests +def test_flat_replay_preserves_failed_item_error_type(): + """A failed FLAT branch carries the same error type live and on replay. -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. + Live construction records the raw escaping type (for example ValueError) + from the ChildContextError wrapper. FLAT replay re-executes the virtual + branch and must record the identical discriminator, not the wrapper + class name. + """ - 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 fails(): + msg = "boom" + raise ValueError(msg) - 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`. - """ + def succeed(): + return "big" * 1 - 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) + executables = [Executable(0, succeed), Executable(1, fails)] - execution_state = Mock() - execution_state.create_checkpoint = Mock() - execution_state.wrap_user_function = lambda func, *args, **kwargs: func + def make_flat_executor(): + return _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, + operation_id_namespace=_StubNamespace(), + ) - # 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 + execution_state, executor_context = _make_executor_mocks() + live: BatchResult = make_flat_executor().execute(execution_state, executor_context) + live_failed = next(i for i in live.all if i.status is BatchItemStatus.FAILED) + assert live_failed.error is not None - # 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" + payload = envelope_summary_generator("MapResult", None)(live) + top_checkpoint = Mock() + top_checkpoint.result = payload + + replay_state, replay_context = _make_replay_mocks({}) + replay_state.wrap_user_function = lambda func, *args, **kwargs: func + replayed: BatchResult = make_flat_executor().replay( + replay_state, replay_context, top_checkpoint ) - executor_context = DurableContext( - state=execution_state, - execution_context=execution_context, - parent_id=map_op_id, # This context *is* the map/parallel op. + replayed_failed = next( + i for i in replayed.all if i.status is BatchItemStatus.FAILED ) + assert replayed_failed.error is not None + assert replayed_failed.error.type == live_failed.error.type + assert replayed_failed.error.type == "ValueError" + assert replayed_failed.error.message == live_failed.error.message - executables = [Executable(index=0, func=lambda ctx: "ok")] - executor = TestExecutor( - executables=executables, + +def test_all_completed_runs_every_branch_through_failures(): + """all_completed() tolerates failures: every branch runs, reason ALL_COMPLETED.""" + + def fails(): + msg = "boom" + raise ValueError(msg) + + executables = [Executable(0, fails), Executable(1, lambda: "ok")] + executor = _make_executor( + 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, + completion_config=CompletionConfig.all_completed(), ) + execution_state, executor_context = _make_executor_mocks() - executor._execute_item_in_child_context(executor_context, executables[0]) # noqa: SLF001 + result: BatchResult = executor.execute(execution_state, executor_context) - # 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 + assert [(i.index, i.status) for i in result.all] == [ + (0, BatchItemStatus.FAILED), + (1, BatchItemStatus.SUCCEEDED), + ] + assert result.completion_reason is CompletionReason.ALL_COMPLETED -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.""" +def test_map_min_successful_greater_than_total_raises(): + """min_successful exceeding the total is rejected before the child context. - class TestExecutor(ConcurrentExecutor): - def execute_item(self, child_context, executable): - self.last_child_context = child_context - return executable.func(child_context) + The check lives on CompletionConfig and is called by context.map() / + context.parallel() before child_handler, so it surfaces as a bare + ValidationError with no STARTED/FAILED operation in history (matching + wait and wait_for_condition validation, and Java's constructor throw). + """ + with pytest.raises(ValidationError, match="min_successful cannot be greater"): + CompletionConfig(min_successful=3)._validate_for_total(2) - 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_parallel_min_successful_greater_than_total_raises(): + """min_successful exceeding the branch count is rejected pre-context.""" + with pytest.raises(ValidationError, match="min_successful cannot be greater"): + CompletionConfig(min_successful=2)._validate_for_total(1) - 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, - ) - 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, - ) +def test_validate_for_total_accepts_min_successful_equal_to_total(): + """min_successful equal to the total is valid.""" + CompletionConfig(min_successful=2)._validate_for_total(2) - executor._execute_item_in_child_context(executor_context, executables[0]) # noqa: SLF001 - # 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_validate_for_total_accepts_none_min_successful(): + """A config without min_successful passes any total.""" + CompletionConfig()._validate_for_total(0) -# endregion Virtual-context wire-format tests +def test_map_item_namer_empty_string_is_preserved(): + """An item_namer returning "" is honored, not replaced by the default.""" + executor: MapExecutor[str, str] = MapExecutor.from_items( + items=["a"], + func=lambda ctx, item, idx, items: item, + config=MapConfig(item_namer=lambda item, index: ""), + operation_id_namespace=OperationIdNamespace("test-prefix"), + ) + assert executor.get_iteration_name(0) == "" -def test_flat_mode_produces_deterministic_step_ids_across_runs(): - """Step ids and inner parent_ids must be deterministic under FLAT mode. +def test_parallel_branch_empty_name_is_preserved(): + """A ParallelBranch with name="" is honored, not replaced by the default.""" + executor: ParallelExecutor[int] = ParallelExecutor.from_callables( + [ParallelBranch(func=lambda ctx: 1, name="")], + ParallelConfig(), + operation_id_namespace=OperationIdNamespace("test-prefix"), + ) + assert executor.get_iteration_name(0) == "" - 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): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.captured = [] +def test_item_namer_called_eagerly_for_unscheduled_items(): + """item_namer runs for every input at map start, including items that + never get scheduled because of early completion.""" + named: list[int] = [] - 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 namer(item: int, index: int) -> str: + named.append(index) + return f"n-{index}" - def make_run(): - execution_state = Mock() - execution_state.create_checkpoint = Mock() + executor: MapExecutor[int, int] = MapExecutor.from_items( + items=[10, 20, 30], + func=lambda ctx, item, idx, items: item, + config=MapConfig( + max_concurrency=1, + completion_config=CompletionConfig(min_successful=1), + item_namer=namer, + ), + operation_id_namespace=OperationIdNamespace("test-prefix"), + ) - 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 + assert named == [0, 1, 2] + assert len(executor.executables) == 3 - 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=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 test_operation_id_namespace_derivation_is_stable(): + """The id scheme is pinned: prefix-step hashed with blake2b, 64 hex chars. - run_a = make_run() - run_b = make_run() + Checkpointed executions depend on this derivation staying byte-identical + across releases. + """ + expected: str = hashlib.blake2b(b"some-operation-id-7").hexdigest()[:64] + assert OperationIdNamespace("some-operation-id").create_id_for_step(7) == expected - # 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) + expected_no_prefix: str = hashlib.blake2b(b"7").hexdigest()[:64] + assert OperationIdNamespace(None).create_id_for_step(7) == expected_no_prefix 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 4983b0ea..156bae0a 100644 --- a/packages/aws-durable-execution-sdk-python/tests/config_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/config_test.py @@ -2,6 +2,8 @@ from unittest.mock import Mock +import pytest + from aws_durable_execution_sdk_python.config import ( CallbackConfig, ChildConfig, @@ -13,6 +15,7 @@ StepConfig, StepSemantics, ) +from aws_durable_execution_sdk_python.exceptions import ValidationError from aws_durable_execution_sdk_python.waits import ( WaitForConditionConfig, WaitForConditionDecision, @@ -45,7 +48,7 @@ def test_completion_config_all_completed(): config = CompletionConfig.all_completed() assert config.min_successful is None assert config.tolerated_failure_count is None - assert config.tolerated_failure_percentage is None + assert config.tolerated_failure_percentage == 100 def test_completion_config_all_successful(): @@ -201,3 +204,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/context_test.py b/packages/aws-durable-execution-sdk-python/tests/context_test.py index f4487748..6f7b2b13 100644 --- a/packages/aws-durable-execution-sdk-python/tests/context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/context_test.py @@ -11,6 +11,7 @@ from aws_durable_execution_sdk_python.config import ( CallbackConfig, ChildConfig, + CompletionConfig, Duration, InvokeConfig, MapConfig, @@ -1857,6 +1858,51 @@ def mock_run_side_effect(func, name=None, config=None): mock_map_handler.assert_called_once() +@patch("aws_durable_execution_sdk_python.context.child_handler") +def test_context_map_min_successful_greater_than_total_raises_bare( + mock_child_handler, +): + """ctx.map validates min_successful before the child context starts. + + The error is a bare ValidationError (not a checkpointed operation + failure wrapped in ChildContextError), matching wait and + wait_for_condition validation. + """ + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = ( + "arn:aws:durable:us-east-1:123456789012:execution/test" + ) + context = create_test_context(state=mock_state) + + with pytest.raises(ValidationError, match="min_successful cannot be greater"): + context.map( + [1, 2], + lambda ctx, item, idx, items: item, + config=MapConfig(completion_config=CompletionConfig(min_successful=3)), + ) + # No operation started: the child context was never entered. + mock_child_handler.assert_not_called() + + +@patch("aws_durable_execution_sdk_python.context.child_handler") +def test_context_parallel_min_successful_greater_than_total_raises_bare( + mock_child_handler, +): + """ctx.parallel validates min_successful before the child context starts.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = ( + "arn:aws:durable:us-east-1:123456789012:execution/test" + ) + context = create_test_context(state=mock_state) + + with pytest.raises(ValidationError, match="min_successful cannot be greater"): + context.parallel( + [lambda ctx: 1], + config=ParallelConfig(completion_config=CompletionConfig(min_successful=2)), + ) + mock_child_handler.assert_not_called() + + def test_context_parallel_handler_call(): """Test that parallel method calls through to parallel_handler (line 306).""" execution_calls = [] @@ -2201,7 +2247,7 @@ def test_should_use_step_id_prefix_when_generating_step_ids(): ) expected_prefixed = hashlib.blake2b(b"branch-op-1").hexdigest()[:64] - assert virtual._create_step_id_for_logical_step(1) == expected_prefixed # noqa: SLF001 + assert virtual._create_step_id_for_logical_step(1) == expected_prefixed def test_should_use_parent_id_as_step_prefix_when_non_virtual(): @@ -2228,7 +2274,7 @@ def test_should_use_parent_id_as_step_prefix_when_non_virtual(): ) expected = hashlib.blake2b(b"parent-op-1").hexdigest()[:64] - assert non_virtual._create_step_id_for_logical_step(1) == expected # noqa: SLF001 + assert non_virtual._create_step_id_for_logical_step(1) == expected assert non_virtual.is_virtual is False @@ -2282,7 +2328,7 @@ def test_should_create_virtual_child_with_none_parent_when_parent_is_root(): assert child.is_virtual is True expected = hashlib.blake2b(b"child-op-1").hexdigest()[:64] - assert child._create_step_id_for_logical_step(1) == expected # noqa: SLF001 + assert child._create_step_id_for_logical_step(1) == expected def test_should_propagate_outer_parent_id_when_virtual_is_nested_in_virtual(): @@ -2327,7 +2373,7 @@ def test_should_propagate_outer_parent_id_when_virtual_is_nested_in_virtual(): # own operation id; they must not leak the outer ancestor into the # step-id namespace. expected = hashlib.blake2b(b"inner-branch-op-1").hexdigest()[:64] - assert inner_branch._create_step_id_for_logical_step(1) == expected # noqa: SLF001 + assert inner_branch._create_step_id_for_logical_step(1) == expected # endregion Virtual-context identity tests @@ -2635,8 +2681,8 @@ def test_replay_aware_stays_replaying_between_two_completed_ops(): replay_status=ReplayStatus.REPLAY, ) # Both the wrapped op and the following op already completed. - first_id = ctx._create_step_id_for_logical_step(1) # noqa: SLF001 - second_id = ctx._create_step_id_for_logical_step(2) # noqa: SLF001 + first_id = ctx._create_step_id_for_logical_step(1) + second_id = ctx._create_step_id_for_logical_step(2) ctx.state._operations[first_id] = _step_op( # noqa: SLF001 first_id, OperationStatus.SUCCEEDED ) @@ -2677,8 +2723,8 @@ def test_replay_aware_terminal_non_success_op_stays_replaying(terminal_status): ) # op1: terminal-but-not-succeeded/failed (e.g. a handled invoke/callback timeout). # op2: a completed step that ran after it. - first_id = ctx._create_step_id_for_logical_step(1) # noqa: SLF001 - second_id = ctx._create_step_id_for_logical_step(2) # noqa: SLF001 + first_id = ctx._create_step_id_for_logical_step(1) + second_id = ctx._create_step_id_for_logical_step(2) ctx.state._operations[first_id] = Operation( # noqa: SLF001 operation_id=first_id, operation_type=OperationType.CHAINED_INVOKE, @@ -2734,8 +2780,8 @@ def test_replay_aware_step_stays_replaying_for_completed_op(): replay_status=ReplayStatus.REPLAY, ) # Wrapped op completed; a following op also completed so nothing flips. - first_id = ctx._create_step_id_for_logical_step(1) # noqa: SLF001 - second_id = ctx._create_step_id_for_logical_step(2) # noqa: SLF001 + first_id = ctx._create_step_id_for_logical_step(1) + second_id = ctx._create_step_id_for_logical_step(2) ctx.state._operations[first_id] = _step_op( # noqa: SLF001 first_id, OperationStatus.SUCCEEDED ) @@ -2885,8 +2931,8 @@ def on_operation_end(self, info): execution_context=ExecutionContext(durable_execution_arn="arn"), replay_status=ReplayStatus.REPLAY, ) - callback_id = ctx._create_step_id_for_logical_step(1) # noqa: SLF001 - following_id = ctx._create_step_id_for_logical_step(2) # noqa: SLF001 + callback_id = ctx._create_step_id_for_logical_step(1) + following_id = ctx._create_step_id_for_logical_step(2) state._operations[callback_id] = _callback_op( # noqa: SLF001 callback_id, OperationStatus.SUCCEEDED ) 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 a7a6024c..02f69a7f 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 @@ -14,6 +14,7 @@ CompletionReason, Executable, ) +from aws_durable_execution_sdk_python.identifier import OperationIdNamespace from aws_durable_execution_sdk_python.config import ( CompletionConfig, MapConfig, @@ -23,12 +24,22 @@ from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import OperationSubType from aws_durable_execution_sdk_python.operation import child # PLC0415 -from aws_durable_execution_sdk_python.operation.map import MapExecutor, map_handler +from aws_durable_execution_sdk_python.operation.map import ( + MapExecutor, + map_handler, +) from aws_durable_execution_sdk_python.serdes import serialize from aws_durable_execution_sdk_python.state import ExecutionState from tests.serdes_test import CustomStrSerDes +class _StubNamespace(OperationIdNamespace): + """Test namespace producing readable ids matching checkpoint fixtures.""" + + def create_id_for_step(self, step: int) -> str: + return f"op_{step}" + + def create_test_context( state: ExecutionState | None = None, parent_id: str | None = None ) -> DurableContext: @@ -62,6 +73,7 @@ def test_map_executor_init(): name_prefix="test-", serdes=None, nesting_type=NestingType.FLAT, + operation_id_namespace=_StubNamespace(), ) assert executor.items == items @@ -78,12 +90,18 @@ def callable_func(ctx, item, idx, items): config = MapConfig(max_concurrency=3, nesting_type=NestingType.FLAT) - executor = MapExecutor.from_items(items, callable_func, config) + executor = MapExecutor.from_items( + items, + callable_func, + config, + operation_id_namespace=_StubNamespace(), + ) assert len(executor.executables) == 3 assert executor.items == items - assert all(exe.func == callable_func for exe in executor.executables) assert [exe.index for exe in executor.executables] == [0, 1, 2] + # Each executable binds its item: calling it runs the user function. + assert executor.executables[1].func(Mock()) == "B" assert executor.nesting_type is NestingType.FLAT @@ -98,6 +116,7 @@ def callable_func(ctx, item, idx, items): items, callable_func, MapConfig(), + operation_id_namespace=_StubNamespace(), ) assert len(executor.executables) == 1 @@ -105,7 +124,7 @@ def callable_func(ctx, item, idx, items): assert executor.nesting_type is NestingType.NESTED -@patch("aws_durable_execution_sdk_python.operation.map.logger") +@patch("aws_durable_execution_sdk_python.concurrency.executor.logger") def test_map_executor_execute_item(mock_logger): """Test MapExecutor.execute_item method with logging.""" items = ["hello", "world"] @@ -117,6 +136,7 @@ def callable_func(ctx, item, idx, items): items, callable_func, MapConfig(), + operation_id_namespace=_StubNamespace(), ) executable = executor.executables[0] @@ -124,8 +144,8 @@ def callable_func(ctx, item, idx, items): assert result == "hello_0" assert mock_logger.debug.call_count == 2 - mock_logger.debug.assert_any_call("🗺️ Processing map item: %s", 0) - mock_logger.debug.assert_any_call("✅ Processed map item: %s", 0) + mock_logger.debug.assert_any_call("▶️ Processing branch: %s", 0) + mock_logger.debug.assert_any_call("✅ Processed branch: %s", 0) def test_map_executor_execute_item_with_context(): @@ -139,6 +159,7 @@ def callable_func(ctx, item, idx, items): items, callable_func, MapConfig(), + operation_id_namespace=_StubNamespace(), ) executable = executor.executables[1] @@ -159,6 +180,9 @@ def mock_run_in_child_context(func, name, config): # Create a minimal ExecutionState mock class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -177,6 +201,7 @@ def get_checkpoint_result(self, operation_id): execution_state, mock_run_in_child_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) assert isinstance(result, BatchResult) @@ -193,6 +218,9 @@ def mock_run_in_child_context(func, name, config): return func("mock_context") class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -213,6 +241,7 @@ def get_checkpoint_result(self, operation_id): execution_state, mock_run_in_child_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) assert isinstance(result, BatchResult) @@ -234,6 +263,7 @@ def callable_func(ctx, item, idx, items_list): items, callable_func, MapConfig(), + operation_id_namespace=_StubNamespace(), ) executable = executor.executables[2] @@ -253,6 +283,7 @@ def callable_func(ctx, item, idx, items): items, callable_func, MapConfig(), + operation_id_namespace=_StubNamespace(), ) assert len(executor.executables) == 0 @@ -270,6 +301,7 @@ def callable_func(ctx, item, idx, items): items, callable_func, MapConfig(), + operation_id_namespace=_StubNamespace(), ) assert len(executor.executables) == 1 @@ -288,6 +320,7 @@ def callable_func(ctx, item, idx, items): items, callable_func, MapConfig(), + operation_id_namespace=_StubNamespace(), ) # Verify it has inherited attributes from ConcurrentExecutor @@ -309,7 +342,7 @@ def callable_func(ctx, item, idx, items): ) executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" executor_context.create_child_context = lambda *args, **kwargs: Mock() with patch.object( @@ -317,6 +350,9 @@ def callable_func(ctx, item, idx, items): ) as mock_execute: class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -335,6 +371,7 @@ def get_checkpoint_result(self, operation_id): execution_state, executor_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify execute was called @@ -362,10 +399,13 @@ def callable_func(ctx, item, idx, items): mock_from_items.return_value = mock_executor executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" executor_context.create_child_context = lambda *args, **kwargs: Mock() class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -383,6 +423,7 @@ def get_checkpoint_result(self, operation_id): execution_state, executor_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify from_items was called with a MapConfig instance @@ -410,12 +451,15 @@ def callable_func(ctx, item, idx, items): return f"RESULT_{item.upper()}" executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -434,6 +478,7 @@ def get_checkpoint_result(self, operation_id): execution_state, executor_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify execute was called @@ -453,10 +498,13 @@ def mock_summary_generator(result): config = MapConfig(summary_generator=mock_summary_generator) executor_context = Mock() - executor_context._create_step_id_for_logical_step = Mock(side_effect=["1", "2"]) # noqa SLF001 + executor_context._create_step_id_for_logical_step = Mock(side_effect=["1", "2"]) executor_context.create_child_context = Mock(return_value=Mock()) class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -475,35 +523,12 @@ def get_checkpoint_result(self, operation_id): execution_state, executor_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify that create_child_context was called twice (N=2 items) assert executor_context.create_child_context.call_count == 2 - # Verify that _create_step_id_for_logical_step was called twice with unique values - assert executor_context._create_step_id_for_logical_step.call_count == 2 # noqa SLF001 - calls = executor_context._create_step_id_for_logical_step.call_args_list # noqa SLF001 - # Verify unique values were passed - assert calls[0] != calls[1] - - -def test_map_executor_from_items_with_summary_generator(): - """Test MapExecutor.from_items preserves summary_generator.""" - items = ["item1"] - - def callable_func(ctx, item, idx, items): - return f"result_{item}" - - def mock_summary_generator(result): - return f"Map summary: {result}" - - config = MapConfig(summary_generator=mock_summary_generator) - - executor = MapExecutor.from_items(items, callable_func, config) - - # Verify that the summary_generator is preserved in the executor - assert executor.summary_generator is mock_summary_generator - def test_map_handler_default_summary_generator(): """Test that map_handler calls executor_context methods correctly with default config.""" @@ -513,10 +538,13 @@ def callable_func(ctx, item, idx, items): return f"result_{item}" executor_context = Mock() - executor_context._create_step_id_for_logical_step = Mock(return_value="1") # noqa SLF001 + executor_context._create_step_id_for_logical_step = Mock(return_value="1") executor_context.create_child_context = Mock(return_value=Mock()) # SLF001 class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -535,39 +563,12 @@ def get_checkpoint_result(self, operation_id): execution_state, executor_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify that create_child_context was called once (N=1 item) assert executor_context.create_child_context.call_count == 1 - # Verify that _create_step_id_for_logical_step was called once - assert executor_context._create_step_id_for_logical_step.call_count == 1 # noqa SLF001 - - -def test_map_executor_init_with_summary_generator(): - """Test MapExecutor initialization with summary_generator.""" - items = ["item1"] - executables = [Executable(index=0, func=lambda: None)] - - def mock_summary_generator(result): - return f"Summary: {result}" - - executor = MapExecutor( - executables=executables, - items=items, - max_concurrency=2, - completion_config=CompletionConfig(), - top_level_sub_type=OperationSubType.MAP, - iteration_sub_type=OperationSubType.MAP_ITERATION, - name_prefix="test-", - serdes=None, - summary_generator=mock_summary_generator, - ) - - assert executor.summary_generator is mock_summary_generator - assert executor.items == items - assert executor.executables == executables - def test_map_handler_with_explicit_none_summary_generator(): """Test that map_handler calls executor_context methods correctly with explicit None summary_generator.""" @@ -580,6 +581,9 @@ def func(ctx, item, index, array): config = MapConfig(summary_generator=None) class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -591,7 +595,7 @@ def get_checkpoint_result(self, operation_id): ) executor_context = Mock() - executor_context._create_step_id_for_logical_step = Mock( # noqa: SLF001 + executor_context._create_step_id_for_logical_step = Mock( side_effect=["1", "2", "3"] ) executor_context.create_child_context = Mock(return_value=Mock()) @@ -604,6 +608,7 @@ def get_checkpoint_result(self, operation_id): execution_state=execution_state, map_context=executor_context, operation_identifier=operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify that create_child_context was called 3 times (N=3 items) @@ -619,6 +624,9 @@ def callable_func(ctx, item, idx, items): # Mock execution state that indicates operation already succeeded class MockExecutionState: + def register_branch_pool(self, pool): + pass + durable_execution_arn = "arn:aws:durable:us-east-1:123456789012:execution/test" def get_checkpoint_result(self, operation_id): @@ -637,7 +645,7 @@ def get_checkpoint_result(self, operation_id): # Mock map context map_context = Mock() - map_context._create_step_id_for_logical_step = Mock( # noqa: SLF001 + map_context._create_step_id_for_logical_step = Mock( side_effect=["child_1", "child_2"] ) @@ -667,10 +675,12 @@ def get_checkpoint_result(self, operation_id): execution_state, map_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) # 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 @@ -683,6 +693,9 @@ def callable_func(ctx, item, idx, items): # Mock execution state that indicates operation succeeded but children need replay class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() if operation_id == "test_op": @@ -700,7 +713,7 @@ def get_checkpoint_result(self, operation_id): # Mock map context map_context = Mock() - map_context._create_step_id_for_logical_step = Mock(return_value="child_1") # noqa: SLF001 + map_context._create_step_id_for_logical_step = Mock(return_value="child_1") # Mock the executor's replay method and _execute_item_in_child_context with ( @@ -729,9 +742,11 @@ def get_checkpoint_result(self, operation_id): execution_state, map_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) - 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 @@ -777,6 +792,9 @@ def test_func(ctx, item, idx, items): execution_count = 0 class MockExecutionState: + def register_branch_pool(self, pool): + pass + durable_execution_arn = "arn:aws:durable:us-east-1:123456789012:execution/test" def get_checkpoint_result(self, operation_id): @@ -811,7 +829,13 @@ def get_checkpoint_result(self, operation_id): # FIRST EXECUTION - should call execute execution_count = 0 map_handler( - items, test_func, config, execution_state, map_context, operation_identifier + items, + test_func, + config, + execution_state, + map_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify execute was called, replay was not @@ -825,7 +849,13 @@ def get_checkpoint_result(self, operation_id): # SECOND EXECUTION - should call replay execution_count = 1 map_handler( - items, test_func, config, execution_state, map_context, operation_identifier + items, + test_func, + config, + execution_state, + map_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify replay was called, execute was not @@ -882,7 +912,14 @@ def create_id(self, i): else f"child-{i}" ) - with patch.object(DurableContext, "_create_step_id_for_logical_step", create_id): + with ( + patch.object(DurableContext, "_create_step_id_for_logical_step", create_id), + patch.object( + OperationIdNamespace, + "create_id_for_step", + lambda self, step: f"child-{step}", + ), + ): context = create_test_context(state=mock_state) context.map( ["a", "b"], @@ -947,7 +984,14 @@ def create_id(self, i): else f"child-{i}" ) - with patch.object(DurableContext, "_create_step_id_for_logical_step", create_id): + with ( + patch.object(DurableContext, "_create_step_id_for_logical_step", create_id), + patch.object( + OperationIdNamespace, + "create_id_for_step", + lambda self, step: f"child-{step}", + ), + ): context = create_test_context(state=mock_state) context.map( ["a", "b"], @@ -974,6 +1018,9 @@ def func(ctx, item, idx, items): return {"item": item.upper(), "index": idx} class MockExecutionState: + def register_branch_pool(self, pool): + pass + durable_execution_arn = "arn:test" def get_checkpoint_result(self, operation_id): @@ -983,7 +1030,7 @@ def get_checkpoint_result(self, operation_id): execution_state = MockExecutionState() map_context = Mock() - map_context._create_step_id_for_logical_step = Mock(side_effect=["1", "2", "3"]) # noqa SLF001 + map_context._create_step_id_for_logical_step = Mock(side_effect=["1", "2", "3"]) child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func map_context.create_child_context = Mock(return_value=child_context) @@ -993,7 +1040,13 @@ def get_checkpoint_result(self, operation_id): # Execute map result = map_handler( - items, func, MapConfig(), execution_state, map_context, operation_identifier + items, + func, + MapConfig(), + execution_state, + map_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Serialize the BatchResult @@ -1058,8 +1111,15 @@ def create_id(self, i): else f"child-{i}" ) - with patch.object( - DurableContext, "_create_step_id_for_logical_step", create_id + with ( + patch.object( + DurableContext, "_create_step_id_for_logical_step", create_id + ), + patch.object( + OperationIdNamespace, + "create_id_for_step", + lambda self, step: f"child-{step}", + ), ): context = create_test_context(state=mock_state) result = context.map(["a", "b"], lambda ctx, item, idx, items: item) @@ -1117,8 +1177,15 @@ def create_id(self, i): else f"child-{i}" ) - with patch.object( - DurableContext, "_create_step_id_for_logical_step", create_id + with ( + patch.object( + DurableContext, "_create_step_id_for_logical_step", create_id + ), + patch.object( + OperationIdNamespace, + "create_id_for_step", + lambda self, step: f"child-{step}", + ), ): context = create_test_context(state=mock_state) result = context.map(["a", "b"], lambda ctx, item, idx, items: item) @@ -1183,8 +1250,15 @@ def create_id(self, i): else f"child-{i}" ) - with patch.object( - DurableContext, "_create_step_id_for_logical_step", create_id + with ( + patch.object( + DurableContext, "_create_step_id_for_logical_step", create_id + ), + patch.object( + OperationIdNamespace, + "create_id_for_step", + lambda self, step: f"child-{step}", + ), ): context = create_test_context(state=mock_state) result = context.map( @@ -1251,6 +1325,7 @@ def test_map_executor_get_iteration_name_default(): items=items, func=lambda ctx, item, idx, items: item, config=config, + operation_id_namespace=_StubNamespace(), ) assert executor.get_iteration_name(0) == "map-item-0" @@ -1270,6 +1345,7 @@ def test_map_executor_get_iteration_name_with_item_namer(): items=items, func=lambda ctx, item, idx, items: item, config=config, + operation_id_namespace=_StubNamespace(), ) assert executor.get_iteration_name(0) == "process-order-1" @@ -1292,12 +1368,12 @@ def namer(item, index): items=items, func=lambda ctx, item, idx, items: item, config=config, + operation_id_namespace=_StubNamespace(), ) - executor.get_iteration_name(0) - executor.get_iteration_name(2) - - assert received_args == [("alpha", 0), ("gamma", 2)] + assert received_args == [("alpha", 0), ("beta", 1), ("gamma", 2)] + assert executor.get_iteration_name(0) == "item-0-alpha" + assert executor.get_iteration_name(2) == "item-2-gamma" def test_map_executor_item_namer_uses_index(): @@ -1309,6 +1385,7 @@ def test_map_executor_item_namer_uses_index(): items=items, func=lambda ctx, item, idx, items: item, config=config, + operation_id_namespace=_StubNamespace(), ) assert executor.get_iteration_name(0) == "step-1" @@ -1325,6 +1402,7 @@ def test_map_executor_item_namer_none_falls_back_to_default(): items=items, func=lambda ctx, item, idx, items: item, config=config, + operation_id_namespace=_StubNamespace(), ) assert executor.get_iteration_name(0) == "map-item-0" @@ -1340,9 +1418,10 @@ def test_map_executor_from_items_passes_item_namer(): items=["a"], func=lambda ctx, item, idx, items: item, config=config, + operation_id_namespace=_StubNamespace(), ) - assert executor._item_namer is namer + assert executor.executables[0].name == "custom-0" def test_map_config_generic_with_item_namer(): @@ -1356,3 +1435,50 @@ def test_map_config_generic_with_item_namer(): # endregion + + +def test_map_handler_defaults_summary_generator_for_user_config(): + """A user config without a summary generator gets the default (JS parity). + + Without the default, a large result checkpoints an empty summary and + replay cannot reconstruct the exact live result. + """ + + class MockExecutionState: + def register_branch_pool(self, pool): + pass + + durable_execution_arn = "arn:aws:durable:us-east-1:123456789012:execution/test" + + def get_checkpoint_result(self, operation_id): + mock_result = Mock() + mock_result.is_succeeded.return_value = False + mock_result.is_failed.return_value = False + mock_result.is_existent.return_value = False + return mock_result + + captured_configs: list[MapConfig] = [] + original_from_items = MapExecutor.from_items.__func__ + + def capturing_from_items(cls, items, func, config, operation_id_namespace): + captured_configs.append(config) + return original_from_items(cls, items, func, config, operation_id_namespace) + + operation_identifier = OperationIdentifier( + "test_op", OperationSubType.MAP, "parent", "test_map" + ) + with patch.object(MapExecutor, "from_items", classmethod(capturing_from_items)): + result = map_handler( + [], + lambda ctx, item, idx, items: item, + MapConfig(max_concurrency=5), + MockExecutionState(), + Mock(), + operation_identifier, + operation_id_namespace=_StubNamespace(), + ) + + assert len(captured_configs) == 1 + assert captured_configs[0].summary_generator is None + assert captured_configs[0].max_concurrency == 5 + assert result.total_count == 0 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..a19f87e8 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 @@ -9,6 +9,8 @@ import pytest from aws_durable_execution_sdk_python.concurrency.executor import ConcurrentExecutor +from aws_durable_execution_sdk_python.identifier import OperationIdNamespace + # Mock the executor.execute method to return a BatchResult from aws_durable_execution_sdk_python.concurrency.models import ( @@ -36,6 +38,13 @@ from tests.serdes_test import CustomStrSerDes +class _StubNamespace(OperationIdNamespace): + """Test namespace producing readable ids matching checkpoint fixtures.""" + + def create_id_for_step(self, step: int) -> str: + return f"op_{step}" + + def create_test_context( state: ExecutionState | None = None, parent_id: str | None = None ) -> DurableContext: @@ -73,6 +82,7 @@ def test_parallel_executor_init(): name_prefix="test-", serdes=None, nesting_type=NestingType.FLAT, + operation_id_namespace=_StubNamespace(), ) assert executor.executables == executables @@ -96,7 +106,11 @@ def func2(ctx): callables = [func1, func2] config = ParallelConfig(max_concurrency=3, nesting_type=NestingType.FLAT) - executor = ParallelExecutor.from_callables(callables, config) + executor = ParallelExecutor.from_callables( + callables, + config, + operation_id_namespace=_StubNamespace(), + ) assert len(executor.executables) == 2 assert executor.executables[0].index == 0 @@ -119,7 +133,11 @@ def func1(ctx): callables = [func1] config = ParallelConfig() - executor = ParallelExecutor.from_callables(callables, config) + executor = ParallelExecutor.from_callables( + callables, + config, + operation_id_namespace=_StubNamespace(), + ) assert len(executor.executables) == 1 assert executor.max_concurrency is None @@ -142,6 +160,7 @@ def test_func(ctx): iteration_sub_type=OperationSubType.PARALLEL_BRANCH, name_prefix="test-", serdes=None, + operation_id_namespace=_StubNamespace(), ) child_context = "test-context" @@ -166,6 +185,7 @@ def failing_func(ctx): iteration_sub_type=OperationSubType.PARALLEL_BRANCH, name_prefix="test-", serdes=None, + operation_id_namespace=_StubNamespace(), ) child_context = "test-context" @@ -187,6 +207,9 @@ def func2(ctx): config = ParallelConfig(max_concurrency=2) class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -213,6 +236,7 @@ def mock_run_in_child_context(callable_func, name, child_config): execution_state, mock_run_in_child_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) assert result == mock_batch_result @@ -227,6 +251,9 @@ def func1(ctx): callables = [func1] class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -252,6 +279,7 @@ def mock_run_in_child_context(callable_func, name, child_config): execution_state, mock_run_in_child_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) assert result == mock_batch_result @@ -267,6 +295,9 @@ def func1(ctx): config = ParallelConfig(max_concurrency=5) class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -278,7 +309,7 @@ def get_checkpoint_result(self, operation_id): ) executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" executor_context.create_child_context = lambda *args, **kwargs: Mock() with patch.object(ParallelExecutor, "from_callables") as mock_from_callables: @@ -288,10 +319,19 @@ def get_checkpoint_result(self, operation_id): mock_from_callables.return_value = mock_executor result = parallel_handler( - callables, config, execution_state, executor_context, operation_identifier + callables, + config, + execution_state, + executor_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), ) - mock_from_callables.assert_called_once_with(callables, config) + mock_from_callables.assert_called_once() + passed_callables, passed_config = mock_from_callables.call_args.args + assert passed_callables == callables + assert passed_config.max_concurrency == 5 + assert passed_config.summary_generator is None mock_executor.execute.assert_called_once_with( execution_state, executor_context=executor_context ) @@ -307,6 +347,9 @@ def func1(ctx): callables = [func1] class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -318,7 +361,7 @@ def get_checkpoint_result(self, operation_id): ) executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" executor_context.create_child_context = lambda *args, **kwargs: Mock() with patch.object(ParallelExecutor, "from_callables") as mock_from_callables: @@ -328,7 +371,12 @@ def get_checkpoint_result(self, operation_id): mock_from_callables.return_value = mock_executor result = parallel_handler( - callables, None, execution_state, executor_context, operation_identifier + callables, + None, + execution_state, + executor_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), ) assert result == mock_batch_result @@ -351,6 +399,7 @@ def test_parallel_executor_inheritance(): iteration_sub_type=OperationSubType.PARALLEL_BRANCH, name_prefix="test-", serdes=None, + operation_id_namespace=_StubNamespace(), ) assert isinstance(executor, ConcurrentExecutor) @@ -361,7 +410,11 @@ def test_parallel_executor_from_callables_empty_list(): callables = [] config = ParallelConfig() - executor = ParallelExecutor.from_callables(callables, config) + executor = ParallelExecutor.from_callables( + callables, + config, + operation_id_namespace=_StubNamespace(), + ) assert len(executor.executables) == 0 assert executor.max_concurrency is None @@ -387,6 +440,7 @@ def dict_func(ctx): iteration_sub_type=OperationSubType.PARALLEL_BRANCH, name_prefix="test-", serdes=None, + operation_id_namespace=_StubNamespace(), ) # Test different return types @@ -408,6 +462,9 @@ def func1(ctx): callables = [func1] class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -419,7 +476,7 @@ def get_checkpoint_result(self, operation_id): ) executor_context = Mock() - executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 + executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() child_context.state.wrap_user_function = lambda func, *args, **kwargs: func executor_context.create_child_context = lambda *args, **kwargs: child_context @@ -430,6 +487,7 @@ def get_checkpoint_result(self, operation_id): execution_state, executor_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) assert result.all[0].result == "RESULT1" @@ -448,6 +506,9 @@ def mock_summary_generator(result): config = ParallelConfig(summary_generator=mock_summary_generator) class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -459,38 +520,22 @@ def get_checkpoint_result(self, operation_id): ) executor_context = Mock() - executor_context._create_step_id_for_logical_step = Mock(return_value="1") # noqa SLF001 + executor_context._create_step_id_for_logical_step = Mock(return_value="1") executor_context.create_child_context = Mock(return_value=Mock()) # Call parallel_handler parallel_handler( - callables, config, execution_state, executor_context, operation_identifier + callables, + config, + execution_state, + executor_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify that create_child_context was called once (N=1 job) assert executor_context.create_child_context.call_count == 1 - # Verify that _create_step_id_for_logical_step was called once with unique value - assert executor_context._create_step_id_for_logical_step.call_count == 1 # noqa SLF001 - - -def test_parallel_executor_from_callables_with_summary_generator(): - """Test ParallelExecutor.from_callables preserves summary_generator.""" - - def func1(ctx): - return "result1" - - def mock_summary_generator(result): - return f"Summary: {result}" - - callables = [func1] - config = ParallelConfig(summary_generator=mock_summary_generator) - - executor = ParallelExecutor.from_callables(callables, config) - - # Verify that the summary_generator is preserved in the executor - assert executor.summary_generator is mock_summary_generator - def test_parallel_handler_default_summary_generator(): """Test that parallel_handler calls executor_context methods correctly with default config.""" @@ -504,6 +549,9 @@ def func2(ctx): callables = [func1, func2] class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -515,23 +563,22 @@ def get_checkpoint_result(self, operation_id): ) executor_context = Mock() - executor_context._create_step_id_for_logical_step = Mock(side_effect=["1", "2"]) # noqa SLF001 + executor_context._create_step_id_for_logical_step = Mock(side_effect=["1", "2"]) executor_context.create_child_context = Mock(return_value=Mock()) # Call parallel_handler with None config (should use default) parallel_handler( - callables, None, execution_state, executor_context, operation_identifier + callables, + None, + execution_state, + executor_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify that create_child_context was called twice (N=2 jobs) assert executor_context.create_child_context.call_count == 2 - # Verify that _create_step_id_for_logical_step was called twice with unique values - assert executor_context._create_step_id_for_logical_step.call_count == 2 # noqa SLF001 - calls = executor_context._create_step_id_for_logical_step.call_args_list # noqa SLF001 - # Verify unique values were passed - assert calls[0] != calls[1] - def test_parallel_handler_with_explicit_none_summary_generator(): """Test that parallel_handler calls executor_context methods correctly with explicit None summary_generator.""" @@ -550,6 +597,9 @@ def func3(ctx): config = ParallelConfig(summary_generator=None) class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() mock_result.is_succeeded.return_value = False @@ -561,7 +611,7 @@ def get_checkpoint_result(self, operation_id): ) executor_context = Mock() - executor_context._create_step_id_for_logical_step = Mock( # noqa: SLF001 + executor_context._create_step_id_for_logical_step = Mock( side_effect=["1", "2", "3"] ) executor_context.create_child_context = Mock(return_value=Mock()) @@ -573,6 +623,7 @@ def get_checkpoint_result(self, operation_id): execution_state=execution_state, parallel_context=executor_context, operation_identifier=operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify that create_child_context was called 3 times (N=3 jobs) @@ -592,6 +643,9 @@ def func2(ctx): # Mock execution state that indicates operation already succeeded class MockExecutionState: + def register_branch_pool(self, pool): + pass + durable_execution_arn = "arn:aws:durable:us-east-1:123456789012:execution/test" def get_checkpoint_result(self, operation_id): @@ -610,7 +664,7 @@ def get_checkpoint_result(self, operation_id): # Mock parallel context parallel_context = Mock() - parallel_context._create_step_id_for_logical_step = Mock( # noqa: SLF001 + parallel_context._create_step_id_for_logical_step = Mock( side_effect=["child_1", "child_2"] ) @@ -634,11 +688,17 @@ def get_checkpoint_result(self, operation_id): mock_replay.return_value = expected_batch_result result = parallel_handler( - callables, config, execution_state, parallel_context, operation_identifier + callables, + config, + execution_state, + parallel_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), ) # 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 @@ -652,6 +712,9 @@ def func1(ctx): # Mock execution state that indicates operation succeeded but children need replay class MockExecutionState: + def register_branch_pool(self, pool): + pass + def get_checkpoint_result(self, operation_id): mock_result = Mock() if operation_id == "test_op": @@ -669,7 +732,7 @@ def get_checkpoint_result(self, operation_id): # Mock parallel context parallel_context = Mock() - parallel_context._create_step_id_for_logical_step = Mock(return_value="child_1") # noqa: SLF001 + parallel_context._create_step_id_for_logical_step = Mock(return_value="child_1") # Mock the executor's replay method and _execute_item_in_child_context with ( @@ -692,10 +755,16 @@ def get_checkpoint_result(self, operation_id): mock_replay.return_value = expected_batch_result result = parallel_handler( - callables, config, execution_state, parallel_context, operation_identifier + callables, + config, + execution_state, + parallel_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), ) - 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 @@ -744,6 +813,9 @@ def task2(ctx): execution_count = 0 class MockExecutionState: + def register_branch_pool(self, pool): + pass + durable_execution_arn = "arn:aws:durable:us-east-1:123456789012:execution/test" def get_checkpoint_result(self, operation_id): @@ -778,7 +850,12 @@ def get_checkpoint_result(self, operation_id): # FIRST EXECUTION - should call execute execution_count = 0 parallel_handler( - callables, config, execution_state, parallel_context, operation_identifier + callables, + config, + execution_state, + parallel_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify execute was called, replay was not @@ -792,7 +869,12 @@ def get_checkpoint_result(self, operation_id): # SECOND EXECUTION - should call replay execution_count = 1 parallel_handler( - callables, config, execution_state, parallel_context, operation_identifier + callables, + config, + execution_state, + parallel_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Verify replay was called, execute was not @@ -849,7 +931,14 @@ def create_id(self, i): else f"child-{i}" ) - with patch.object(DurableContext, "_create_step_id_for_logical_step", create_id): + with ( + patch.object(DurableContext, "_create_step_id_for_logical_step", create_id), + patch.object( + OperationIdNamespace, + "create_id_for_step", + lambda self, step: f"child-{step}", + ), + ): context = create_test_context(state=mock_state) context.parallel( [lambda ctx: "a", lambda ctx: "b"], @@ -911,7 +1000,14 @@ def create_id(self, i): else f"child-{i}" ) - with patch.object(DurableContext, "_create_step_id_for_logical_step", create_id): + with ( + patch.object(DurableContext, "_create_step_id_for_logical_step", create_id), + patch.object( + OperationIdNamespace, + "create_id_for_step", + lambda self, step: f"child-{step}", + ), + ): context = create_test_context(state=mock_state) context.parallel( [lambda ctx: "a", lambda ctx: "b"], @@ -941,6 +1037,9 @@ def func3(ctx): callables = [func1, func2, func3] class MockExecutionState: + def register_branch_pool(self, pool): + pass + durable_execution_arn = "arn:test" def get_checkpoint_result(self, operation_id): @@ -950,7 +1049,7 @@ def get_checkpoint_result(self, operation_id): execution_state = MockExecutionState() parallel_context = Mock() - parallel_context._create_step_id_for_logical_step = Mock( # noqa SLF001 + parallel_context._create_step_id_for_logical_step = Mock( side_effect=["1", "2", "3"] ) child_context = Mock() @@ -967,6 +1066,7 @@ def get_checkpoint_result(self, operation_id): execution_state, parallel_context, operation_identifier, + operation_id_namespace=_StubNamespace(), ) # Serialize the BatchResult @@ -1031,8 +1131,15 @@ def create_id(self, i): else f"child-{i}" ) - with patch.object( - DurableContext, "_create_step_id_for_logical_step", create_id + with ( + patch.object( + DurableContext, "_create_step_id_for_logical_step", create_id + ), + patch.object( + OperationIdNamespace, + "create_id_for_step", + lambda self, step: f"child-{step}", + ), ): context = create_test_context(state=mock_state) result = context.parallel([lambda ctx: "a", lambda ctx: "b"]) @@ -1090,8 +1197,15 @@ def create_id(self, i): else f"child-{i}" ) - with patch.object( - DurableContext, "_create_step_id_for_logical_step", create_id + with ( + patch.object( + DurableContext, "_create_step_id_for_logical_step", create_id + ), + patch.object( + OperationIdNamespace, + "create_id_for_step", + lambda self, step: f"child-{step}", + ), ): context = create_test_context(state=mock_state) result = context.parallel([lambda ctx: "a", lambda ctx: "b"]) @@ -1156,8 +1270,15 @@ def create_id(self, i): else f"child-{i}" ) - with patch.object( - DurableContext, "_create_step_id_for_logical_step", create_id + with ( + patch.object( + DurableContext, "_create_step_id_for_logical_step", create_id + ), + patch.object( + OperationIdNamespace, + "create_id_for_step", + lambda self, step: f"child-{step}", + ), ): context = create_test_context(state=mock_state) result = context.parallel( @@ -1216,7 +1337,11 @@ def test_parallel_executor_get_iteration_name_default(): callables = [lambda ctx: "a", lambda ctx: "b", lambda ctx: "c"] config = ParallelConfig() - executor = ParallelExecutor.from_callables(callables, config) + executor = ParallelExecutor.from_callables( + callables, + config, + operation_id_namespace=_StubNamespace(), + ) assert executor.get_iteration_name(0) == "parallel-branch-0" assert executor.get_iteration_name(1) == "parallel-branch-1" @@ -1233,7 +1358,11 @@ def test_parallel_executor_get_iteration_name_with_named_branches(): ] config = ParallelConfig() - executor = ParallelExecutor.from_callables(branches, config) + executor = ParallelExecutor.from_callables( + branches, + config, + operation_id_namespace=_StubNamespace(), + ) assert executor.get_iteration_name(0) == "fetch-user-data" assert executor.get_iteration_name(1) == "fetch-order-history" @@ -1250,7 +1379,11 @@ def test_parallel_executor_get_iteration_name_mixed(): ] config = ParallelConfig() - executor = ParallelExecutor.from_callables(branches, config) + executor = ParallelExecutor.from_callables( + branches, + config, + operation_id_namespace=_StubNamespace(), + ) assert executor.get_iteration_name(0) == "named-branch" assert executor.get_iteration_name(1) == "parallel-branch-1" @@ -1266,7 +1399,11 @@ def test_parallel_executor_get_iteration_name_none_name(): ] config = ParallelConfig() - executor = ParallelExecutor.from_callables(branches, config) + executor = ParallelExecutor.from_callables( + branches, + config, + operation_id_namespace=_StubNamespace(), + ) assert executor.get_iteration_name(0) == "parallel-branch-0" @@ -1286,6 +1423,7 @@ def test_parallel_branch_execute_item(): iteration_sub_type=OperationSubType.PARALLEL_BRANCH, name_prefix="parallel-branch-", serdes=None, + operation_id_namespace=_StubNamespace(), ) result = executor.execute_item("test-ctx", executable) @@ -1293,3 +1431,47 @@ def test_parallel_branch_execute_item(): # endregion + + +def test_parallel_handler_defaults_summary_generator_for_user_config(): + """A user config without a summary generator gets the default (JS parity).""" + + class MockExecutionState: + def register_branch_pool(self, pool): + pass + + durable_execution_arn = "arn:aws:durable:us-east-1:123456789012:execution/test" + + def get_checkpoint_result(self, operation_id): + mock_result = Mock() + mock_result.is_succeeded.return_value = False + mock_result.is_failed.return_value = False + mock_result.is_existent.return_value = False + return mock_result + + captured_configs: list[ParallelConfig] = [] + original_from_callables = ParallelExecutor.from_callables.__func__ + + def capturing_from_callables(cls, callables, config, operation_id_namespace): + captured_configs.append(config) + return original_from_callables(cls, callables, config, operation_id_namespace) + + operation_identifier = OperationIdentifier( + "test_op", OperationSubType.PARALLEL, "parent", "test_parallel" + ) + with patch.object( + ParallelExecutor, "from_callables", classmethod(capturing_from_callables) + ): + result = parallel_handler( + [], + ParallelConfig(max_concurrency=3), + MockExecutionState(), + Mock(), + operation_identifier, + operation_id_namespace=_StubNamespace(), + ) + + assert len(captured_configs) == 1 + assert captured_configs[0].summary_generator is None + assert captured_configs[0].max_concurrency == 3 + assert result.total_count == 0 diff --git a/packages/aws-durable-execution-sdk-python/tests/state_test.py b/packages/aws-durable-execution-sdk-python/tests/state_test.py index 4fb277ff..bb9a141f 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -425,6 +425,107 @@ def test_execution_state_creation(): assert state.operations == {} +def test_close_joins_registered_branch_pools(): + """close() blocks until threads in registered branch pools finish.""" + mock_lambda_client = Mock(spec=LambdaClient) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="test_token", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + branch_finished: threading.Event = threading.Event() + + def slow_branch() -> None: + time.sleep(0.2) + branch_finished.set() + + pool: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=1) + state.register_branch_pool(pool) + pool.submit(slow_branch) + # Abandon the pool the way the coordinator does on early completion. + pool.shutdown(wait=False, cancel_futures=True) + + assert not branch_finished.is_set() + state.close() + assert branch_finished.is_set() + + +def test_close_joins_branch_pools_before_stopping_checkpointing(): + """close() joins branch threads while the checkpoint batcher is alive. + + A branch blocked on an in-flight synchronous checkpoint needs the + batcher to respond before it can unwind, so the join must precede + stop_checkpointing. + """ + mock_lambda_client = Mock(spec=LambdaClient) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="test_token", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + stopped_during_branch: list[bool] = [] + + def slow_branch() -> None: + time.sleep(0.2) + stopped_during_branch.append(state._checkpointing_stopped.is_set()) # noqa: SLF001 + + pool: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=1) + state.register_branch_pool(pool) + pool.submit(slow_branch) + pool.shutdown(wait=False, cancel_futures=True) + + state.close() + assert stopped_during_branch == [False] + assert state._checkpointing_stopped.is_set() # noqa: SLF001 + + +def test_close_drains_pools_registered_while_joining(): + """A pool registered during the join (nested map/parallel started by a + branch being joined) is itself joined before checkpointing stops. + + Mock pools pin the exact interleaving: the first pool's shutdown + registers the second pool, so the second is provably absent from + close()'s first snapshot and only a drain loop joins it. + """ + mock_lambda_client: Mock = Mock(spec=LambdaClient) + state: ExecutionState = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="test_token", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + nested_pool: Mock = Mock(spec=ThreadPoolExecutor) + outer_pool: Mock = Mock(spec=ThreadPoolExecutor) + stopped_at_nested_join: list[bool] = [] + + def register_nested(*_args: object, **_kwargs: object) -> None: + state.register_branch_pool(nested_pool) + + def record_checkpointing_state(*_args: object, **_kwargs: object) -> None: + stopped_at_nested_join.append( + state._checkpointing_stopped.is_set() # noqa: SLF001 + ) + + outer_pool.shutdown.side_effect = register_nested + nested_pool.shutdown.side_effect = record_checkpointing_state + state.register_branch_pool(outer_pool) + + state.close() + + outer_pool.shutdown.assert_called_once_with(wait=True) + nested_pool.shutdown.assert_called_once_with(wait=True) + # The nested pool was joined while the checkpoint batcher was still + # alive: a nested branch blocked on a synchronous checkpoint needs the + # batcher to respond before it can unwind. + assert stopped_at_nested_join == [False] + assert state._checkpointing_stopped.is_set() # noqa: SLF001 + + def test_get_checkpoint_result_success_with_result(): """Test get_checkpoint_result with successful operation and result.""" mock_lambda_client = Mock(spec=LambdaClient)