From ed124ddc3532906e6642c8bb0cb181c7bd4e3f48 Mon Sep 17 00:00:00 2001 From: Ayushi Ahjolia Date: Thu, 30 Jul 2026 22:21:44 -0700 Subject: [PATCH] feat: add should_complete to CompletionConfig --- .../examples-catalog.json | 22 + .../src/map/map_with_should_complete.py | 50 ++ .../parallel/parallel_with_should_complete.py | 58 ++ .../test/map/test_map_with_should_complete.py | 38 + .../test_parallel_with_should_complete.py | 34 + .../__init__.py | 10 +- .../concurrency/executor.py | 108 ++- .../concurrency/models.py | 106 ++- .../config.py | 99 +++ .../tests/concurrency_test.py | 759 ++++++++++++++++++ .../tests/operation/map_test.py | 92 +++ .../tests/operation/parallel_test.py | 89 ++ 12 files changed, 1447 insertions(+), 18 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-examples/src/map/map_with_should_complete.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/src/parallel/parallel_with_should_complete.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/test/map/test_map_with_should_complete.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/test/parallel/test_parallel_with_should_complete.py diff --git a/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json b/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json index 030fce1a..d78edd28 100644 --- a/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json +++ b/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json @@ -398,6 +398,17 @@ }, "path": "./src/map/map_with_failure_tolerance.py" }, + { + "name": "Map with Should Complete", + "description": "Map operation with custom should_complete predicate for early completion", + "handler": "map_with_should_complete.handler", + "integration": true, + "durableConfig": { + "RetentionPeriodInDays": 7, + "ExecutionTimeout": 300 + }, + "path": "./src/map/map_with_should_complete.py" + }, { "name": "Map Completion Config", "description": "Reproduces issue where map with minSuccessful loses failure count", @@ -442,6 +453,17 @@ }, "path": "./src/parallel/parallel_with_failure_tolerance.py" }, + { + "name": "Parallel with Should Complete", + "description": "Parallel operation with custom quorum-based should_complete predicate", + "handler": "parallel_with_should_complete.handler", + "integration": true, + "durableConfig": { + "RetentionPeriodInDays": 7, + "ExecutionTimeout": 300 + }, + "path": "./src/parallel/parallel_with_should_complete.py" + }, { "name": "Map with Custom SerDes", "description": "Map operation with custom item-level serialization", diff --git a/packages/aws-durable-execution-sdk-python-examples/src/map/map_with_should_complete.py b/packages/aws-durable-execution-sdk-python-examples/src/map/map_with_should_complete.py new file mode 100644 index 00000000..c2a6528d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/map/map_with_should_complete.py @@ -0,0 +1,50 @@ +"""Example demonstrating map with a custom should_complete predicate. + +The predicate gives full control over when a batch completes early, +beyond the threshold-based fields (min_successful, tolerated_failure_count, +tolerated_failure_percentage). Here we stop processing as soon as we +accumulate 3 successful results, regardless of how many items remain. +""" + +from typing import Any + +from aws_durable_execution_sdk_python.config import ( + CompletionConfig, + CompletionStatus, + MapConfig, +) +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _should_complete(status: CompletionStatus) -> bool: + """Complete once 3 items have succeeded.""" + return status.success_count >= 3 + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + """Process items with a custom completion predicate.""" + items: list[int] = list(range(1, 11)) # [1, 2, ..., 10] + + config: MapConfig = MapConfig( + max_concurrency=2, + completion_config=CompletionConfig(should_complete=_should_complete), + ) + + results = context.map( + inputs=items, + func=lambda ctx, item, index, _: ctx.step( + lambda _: item * 10, name=f"process-{index}" + ), + name="map_should_complete", + config=config, + ) + + return { + "success_count": results.success_count, + "failure_count": results.failure_count, + "started_count": results.started_count, + "completion_reason": results.completion_reason.value, + "results": results.get_results(), + } diff --git a/packages/aws-durable-execution-sdk-python-examples/src/parallel/parallel_with_should_complete.py b/packages/aws-durable-execution-sdk-python-examples/src/parallel/parallel_with_should_complete.py new file mode 100644 index 00000000..0aae6e53 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/parallel/parallel_with_should_complete.py @@ -0,0 +1,58 @@ +"""Example demonstrating parallel with a custom should_complete predicate. + +Uses an index-based quorum rule: the batch completes when branch A (index 0) +succeeds OR both branches B (index 1) and C (index 2) succeed. This shows +how the items snapshot enables dependency-style completion logic. +""" + +from typing import Any + +from aws_durable_execution_sdk_python.config import ( + CompletionConfig, + CompletionStatus, + ParallelConfig, +) +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _quorum_predicate(status: CompletionStatus) -> bool: + """Complete when branch A succeeds OR both B and C succeed.""" + if not status.items: + return False + branch_a_ok: bool = status.items[0].is_succeeded + branch_b_ok: bool = len(status.items) > 1 and status.items[1].is_succeeded + branch_c_ok: bool = len(status.items) > 2 and status.items[2].is_succeeded + return branch_a_ok or (branch_b_ok and branch_c_ok) + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + """Run parallel branches with a quorum-based completion predicate.""" + config: ParallelConfig = ParallelConfig( + max_concurrency=3, + completion_config=CompletionConfig(should_complete=_quorum_predicate), + ) + + functions = [ + # Branch A - slow task + lambda ctx: ctx.step(lambda _: "Branch A done", name="branch-a"), + # Branch B - fast task + lambda ctx: ctx.step(lambda _: "Branch B done", name="branch-b"), + # Branch C - fast task + lambda ctx: ctx.step(lambda _: "Branch C done", name="branch-c"), + # Branch D - never needed if quorum met + lambda ctx: ctx.step(lambda _: "Branch D done", name="branch-d"), + ] + + results = context.parallel( + functions=functions, + name="quorum-branches", + config=config, + ) + + return { + "success_count": results.success_count, + "completion_reason": results.completion_reason.value, + "results": results.get_results(), + } diff --git a/packages/aws-durable-execution-sdk-python-examples/test/map/test_map_with_should_complete.py b/packages/aws-durable-execution-sdk-python-examples/test/map/test_map_with_should_complete.py new file mode 100644 index 00000000..6e3780e7 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/test/map/test_map_with_should_complete.py @@ -0,0 +1,38 @@ +"""Tests for map with should_complete predicate.""" + +import pytest +from aws_durable_execution_sdk_python.execution import InvocationStatus +from aws_durable_execution_sdk_python.lambda_service import OperationStatus +from src.map import map_with_should_complete +from test.conftest import deserialize_operation_payload + + +@pytest.mark.example +@pytest.mark.durable_execution( + handler=map_with_should_complete.handler, + lambda_function_name="Map with Should Complete", +) +def test_map_with_should_complete(durable_runner): + """Test map with custom should_complete predicate that stops at 3 successes.""" + with durable_runner: + result = durable_runner.run(input="test", timeout=10) + + assert result.status is InvocationStatus.SUCCEEDED + + result_data = deserialize_operation_payload(result.result) + + # Predicate completes after 3 successes; with max_concurrency=2, + # a concurrent sibling may finish before the check fires, so up to + # max_concurrency extra items can complete. + assert result_data["success_count"] >= 3 + assert result_data["success_count"] <= 4 # at most 1 extra from concurrency + assert result_data["failure_count"] == 0 + assert result_data["completion_reason"] == "CUSTOM_COMPLETION" + + # Results are item * 10 for the items processed + assert len(result_data["results"]) == result_data["success_count"] + + # Get the map operation + map_op = result.get_context("map_should_complete") + assert map_op is not None + assert map_op.status is OperationStatus.SUCCEEDED diff --git a/packages/aws-durable-execution-sdk-python-examples/test/parallel/test_parallel_with_should_complete.py b/packages/aws-durable-execution-sdk-python-examples/test/parallel/test_parallel_with_should_complete.py new file mode 100644 index 00000000..7e27f3e8 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/test/parallel/test_parallel_with_should_complete.py @@ -0,0 +1,34 @@ +"""Tests for parallel with should_complete quorum predicate.""" + +import pytest +from aws_durable_execution_sdk_python.execution import InvocationStatus +from aws_durable_execution_sdk_python.lambda_service import OperationStatus +from src.parallel import parallel_with_should_complete +from test.conftest import deserialize_operation_payload + + +@pytest.mark.example +@pytest.mark.durable_execution( + handler=parallel_with_should_complete.handler, + lambda_function_name="Parallel with Should Complete", +) +def test_parallel_with_should_complete(durable_runner): + """Test parallel with quorum predicate: branch A OR (B AND C).""" + with durable_runner: + result = durable_runner.run(input="test", timeout=10) + + assert result.status is InvocationStatus.SUCCEEDED + + result_data = deserialize_operation_payload(result.result) + + # Quorum met: either branch A succeeded (1 success), or B and C both + # succeeded (2 successes). With concurrency a raced sibling may also + # complete before the batch stops. + assert result_data["success_count"] >= 1 + assert result_data["completion_reason"] == "CUSTOM_COMPLETION" + assert len(result_data["results"]) >= 1 + + # Get the parallel operation + parallel_op = result.get_context("quorum-branches") + assert parallel_op is not None + assert parallel_op.status is OperationStatus.SUCCEEDED diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py index 6871c5d3..34e22423 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py @@ -7,7 +7,12 @@ # Helper decorators - commonly used for step functions # Concurrency from aws_durable_execution_sdk_python.concurrency.models import BatchResult -from aws_durable_execution_sdk_python.config import ParallelBranch +from aws_durable_execution_sdk_python.config import ( + BatchItemStatus, + CompletionItemStatus, + CompletionStatus, + ParallelBranch, +) from aws_durable_execution_sdk_python.context import ( DurableContext, durable_parallel_branch, @@ -42,12 +47,15 @@ __all__ = [ + "BatchItemStatus", "BatchResult", "CallbackError", "CallbackExternalError", "CallbackSubmitterError", "CallbackTimeoutError", "ChildContextError", + "CompletionItemStatus", + "CompletionStatus", "DurableContext", "DurableExecutionsError", "DurableOperationError", 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 c0e98060..7aafd6ea 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 @@ -12,7 +12,6 @@ from aws_durable_execution_sdk_python.concurrency.models import ( BatchItem, - BatchItemStatus, BatchResult, Branch, BranchEvent, @@ -24,7 +23,9 @@ Executable, ) from aws_durable_execution_sdk_python.config import ( + BatchItemStatus, ChildConfig, + CompletionItemStatus, NestingType, ) from aws_durable_execution_sdk_python.exceptions import ( @@ -151,6 +152,52 @@ def get_iteration_name(self, index: int) -> str: name: str | None = self.executables[index].name return name if name is not None else f"{self.name_prefix}{index}" + def _build_items_snapshot(self) -> tuple[CompletionItemStatus, ...]: + """Build the per-branch status snapshot for the custom predicate. + + Only called when a should_complete predicate is active. Returns a + tuple ordered by branch index so items[i] is the branch at position i. + Distinguishes None (not yet scheduled) from STARTED (running or + suspended) so predicates can reason about scheduling state. + """ + snapshot: list[CompletionItemStatus] = [] + for branch in self.branches: + if branch.status is BranchStatus.COMPLETED: + snapshot.append( + CompletionItemStatus( + index=branch.index, + status=BatchItemStatus.SUCCEEDED, + result=branch.result, + ) + ) + elif branch.status is BranchStatus.FAILED: + error_msg: str | None = ( + str(branch.error) if branch.error is not None else None + ) + snapshot.append( + CompletionItemStatus( + index=branch.index, + status=BatchItemStatus.FAILED, + error_message=error_msg, + ) + ) + elif branch.status is BranchStatus.PENDING: + snapshot.append( + CompletionItemStatus( + index=branch.index, + status=None, + ) + ) + else: + # RUNNING, SUSPENDED, SUSPENDED_WITH_TIMEOUT + snapshot.append( + CompletionItemStatus( + index=branch.index, + status=BatchItemStatus.STARTED, + ) + ) + return tuple(snapshot) + def execute( self, execution_state: ExecutionState, executor_context: DurableContext ) -> BatchResult[ResultType]: @@ -190,10 +237,42 @@ def submit(branch: Branch[CallableType, ResultType]) -> None: ) try: + # Only rebuild the items snapshot after a terminal event changes + # branch state, not after timeouts/resumes with no status change. + needs_snapshot_rebuild: bool = True + items_snapshot: tuple[CompletionItemStatus, ...] = () + # Track whether the custom predicate caused loop exit to avoid + # re-evaluating it in reason() (P2: stateful or intermittently + # failing predicates could give a different answer on re-call). + custom_predicate_fired: bool = False + # Gate custom predicate evaluation until at least one branch + # reaches a terminal state. Without this, a predicate like + # `failure_count == 0` would fire immediately before any work runs. + # Note: on resumed invocations, previously-completed branches + # replay concurrently and report terminal events quickly from + # their checkpoints. The predicate is only evaluated after these + # events arrive, so it sees accurate terminal states. + has_terminal_event: bool = False while True: + if needs_snapshot_rebuild: + items_snapshot = ( + self._build_items_snapshot() + if self.policy.should_complete is not None + else () + ) + needs_snapshot_rebuild = False if self.policy.is_complete( - succeeded, failed - ) or not self.policy.should_continue(failed): + succeeded, failed, items_snapshot, has_terminal_event + ): + # Record whether the custom predicate made this decision + # vs all-items-completed or threshold logic. + if ( + self.policy.should_complete is not None + and succeeded + failed < self.policy.total + ): + custom_predicate_fired = True + break + if not self.policy.should_continue(failed): break # Start branches in index order up to the in-flight limit. @@ -207,6 +286,7 @@ def submit(branch: Branch[CallableType, ResultType]) -> None: submit(pending.popleft()) in_flight += 1 running += 1 + needs_snapshot_rebuild = True # Resume due timed suspends in-process. One checkpoint # refresh serves the whole due wave; a failure is terminal @@ -251,18 +331,24 @@ def submit(branch: Branch[CallableType, ResultType]) -> None: succeeded += 1 running -= 1 in_flight -= 1 + needs_snapshot_rebuild = True + has_terminal_event = True case BranchEventKind.FAILED if event.error is not None: applied.fail(event.error) failed += 1 running -= 1 in_flight -= 1 + needs_snapshot_rebuild = True + has_terminal_event = True case BranchEventKind.SUSPENDED: applied.suspend() running -= 1 + needs_snapshot_rebuild = True 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 + needs_snapshot_rebuild = True case BranchEventKind.ORPHANED: # An ancestor context already checkpointed terminal, so # every further checkpoint under it is rejected. Stop @@ -291,8 +377,12 @@ def submit(branch: Branch[CallableType, ResultType]) -> None: # 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) + # decision that never fired. When the custom predicate caused the + # exit, use CUSTOM_COMPLETION directly to avoid re-evaluating it. + if custom_predicate_fired: + completion_reason: CompletionReason = CompletionReason.CUSTOM_COMPLETION + else: + completion_reason = self.policy.reason(succeeded, failed, items_snapshot) # Apply terminal events that raced the completion decision, so a # branch that finished just before the batch completed is reported @@ -370,7 +460,13 @@ def _create_result( raise InvalidStateError(msg) if completion_reason is None: - completion_reason = self.policy.reason(succeeded, failed) + # Fallback: supply items snapshot so quorum predicates work here too. + fallback_items: tuple[CompletionItemStatus, ...] = ( + self._build_items_snapshot() + if self.policy.should_complete is not None + else () + ) + completion_reason = self.policy.reason(succeeded, failed, fallback_items) return BatchResult(batch_items, completion_reason) def _branch_worker( 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 948af732..8340ec95 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 @@ -13,10 +13,17 @@ ChildContextError, InvalidStateError, ) +from aws_durable_execution_sdk_python.config import ( + BatchItemStatus, + CompletionItemStatus, + CompletionStatus, +) 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 collections.abc import Callable + from aws_durable_execution_sdk_python.config import CompletionConfig from aws_durable_execution_sdk_python.types import SummaryGenerator @@ -30,16 +37,13 @@ # region Result models -class BatchItemStatus(Enum): - SUCCEEDED = "SUCCEEDED" - FAILED = "FAILED" - STARTED = "STARTED" class CompletionReason(Enum): ALL_COMPLETED = "ALL_COMPLETED" MIN_SUCCESSFUL_REACHED = "MIN_SUCCESSFUL_REACHED" FAILURE_TOLERANCE_EXCEEDED = "FAILURE_TOLERANCE_EXCEEDED" + CUSTOM_COMPLETION = "CUSTOM_COMPLETION" @dataclass(frozen=True) @@ -48,8 +52,12 @@ class CompletionPolicy: 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. + (completion reason inference). + + When a custom predicate (``should_complete``) is provided it takes full + precedence: the threshold fields are ignored and the predicate alone + decides early completion. The batch still completes once every item + finishes regardless of the predicate's return value. Fail-fast semantics: when no criteria are configured at all, a single failure exceeds tolerance. @@ -59,6 +67,7 @@ class CompletionPolicy: min_successful: int | None = None tolerated_failure_count: int | None = None tolerated_failure_percentage: int | float | None = None + should_complete: Callable[[CompletionStatus], bool] | None = None @classmethod def from_config( @@ -72,6 +81,7 @@ def from_config( min_successful=config.min_successful, tolerated_failure_count=config.tolerated_failure_count, tolerated_failure_percentage=config.tolerated_failure_percentage, + should_complete=config.should_complete, ) @property @@ -83,8 +93,29 @@ def has_criteria(self) -> bool: or self.tolerated_failure_percentage is not None ) + def _build_status( + self, + succeeded: int, + failed: int, + items: tuple[CompletionItemStatus, ...] = (), + ) -> CompletionStatus: + """Build a CompletionStatus snapshot from current counts.""" + return CompletionStatus( + success_count=succeeded, + failure_count=failed, + completed_count=succeeded + failed, + total_count=self.total, + items=items, + ) + def is_tolerance_exceeded(self, failed: int) -> bool: - """True when failures exceed the configured tolerance.""" + """True when failures exceed the configured tolerance. + + When a custom predicate is active, tolerance checking is disabled - + the predicate owns all early-exit decisions. + """ + if self.should_complete is not None: + return False if not self.has_criteria: return failed > 0 if ( @@ -98,21 +129,53 @@ def is_tolerance_exceeded(self, failed: int) -> bool: return False def should_continue(self, failed: int) -> bool: - """True while more branches may be scheduled.""" + """True while more branches may be scheduled. + + When a custom predicate is active, scheduling is never stopped by + tolerance - the predicate controls completion via is_complete. + """ return not self.is_tolerance_exceeded(failed) - def is_complete(self, succeeded: int, failed: int) -> bool: + def is_complete( + self, + succeeded: int, + failed: int, + items: tuple[CompletionItemStatus, ...] = (), + has_terminal_event: bool = True, + ) -> bool: """True when the batch has met a completion criterion.""" if succeeded + failed >= self.total: return True + if self.should_complete is not None: + # Only evaluate the custom predicate after at least one branch + # has reached a terminal state. This prevents predicates like + # `failure_count == 0` from firing before any work runs. + if not has_terminal_event: + return False + status: CompletionStatus = self._build_status(succeeded, failed, items) + return self.should_complete(status) return self.min_successful is not None and succeeded >= self.min_successful - def reason(self, succeeded: int, failed: int) -> CompletionReason: + def reason( + self, + succeeded: int, + failed: int, + items: tuple[CompletionItemStatus, ...] = (), + ) -> 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.should_complete is not None: + status: CompletionStatus = self._build_status(succeeded, failed, items) + if self.should_complete(status): + return CompletionReason.CUSTOM_COMPLETION + # The predicate did not fire but the loop exited (all items + # completed or raced events pushed totals to total after the + # is_complete check). Report ALL_COMPLETED since every branch + # reached a terminal state. + 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 @@ -346,7 +409,28 @@ def from_items( policy: CompletionPolicy = CompletionPolicy.from_config( len(items), completion_config ) - return cls(items, policy.reason(succeeded_count, failed_count)) + + # Build items snapshot so quorum predicates receive per-item state + # even in the fallback inference path. + item_snapshots: tuple[CompletionItemStatus, ...] = tuple( + CompletionItemStatus( + index=item.index, + status=item.status, + result=item.result + if item.status is BatchItemStatus.SUCCEEDED + else None, + error_message=( + item.error.message + if item.status is BatchItemStatus.FAILED + and item.error is not None + and hasattr(item.error, "message") + else None + ), + ) + for item in items + ) + + return cls(items, policy.reason(succeeded_count, failed_count, item_snapshots)) def to_dict(self) -> dict: return { 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 e25d1b98..9da0f667 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 @@ -99,6 +99,77 @@ class NestingType(Enum): """ +class BatchItemStatus(Enum): + """The status of a batch item in map/parallel operations.""" + + SUCCEEDED = "SUCCEEDED" + FAILED = "FAILED" + STARTED = "STARTED" + + +@dataclass(frozen=True) +class CompletionItemStatus: + """Status snapshot of a single branch/item in a batch operation. + + Provided within :class:`CompletionStatus` so custom predicates can + inspect individual branch outcomes for quorum-style decisions. + + Attributes: + index: The branch/item index (stable across replay). + status: Current status as a BatchItemStatus enum value, or None if + the branch has not been scheduled yet (common when max_concurrency + limits in-flight branches). + result: The branch result when succeeded, None otherwise. + error_message: The error message when failed, None otherwise. + """ + + index: int + status: BatchItemStatus | None = None + result: object | None = None + error_message: str | None = None + + @property + def is_succeeded(self) -> bool: + """True when this branch completed successfully.""" + return self.status is BatchItemStatus.SUCCEEDED + + @property + def is_failed(self) -> bool: + """True when this branch failed.""" + return self.status is BatchItemStatus.FAILED + + @property + def is_started(self) -> bool: + """True when this branch is actively running or suspended.""" + return self.status is BatchItemStatus.STARTED + + +@dataclass(frozen=True) +class CompletionStatus: + """Progress snapshot passed to a custom completion predicate. + + Provided to the ``should_complete`` callable on each branch terminal + event so the predicate can decide whether the batch should finish early. + + Attributes: + success_count: Number of branches that have completed successfully. + failure_count: Number of branches that have failed. + completed_count: Total terminal branches (success_count + failure_count). + total_count: Total number of branches in the batch. + items: Per-branch status snapshot ordered by original index. items[i] + is always the branch defined at position i, regardless of + completion order. Enables index-based quorum rules such as + "complete when branch 0 succeeds OR branches 1 and 2 both + succeed". + """ + + success_count: int + failure_count: int + completed_count: int + total_count: int + items: tuple[CompletionItemStatus, ...] = () + + @dataclass(frozen=True) class CompletionConfig: """Configuration for determining when parallel/map operations complete. @@ -119,8 +190,22 @@ class CompletionConfig: (0.0 to 100.0). If None, no percentage limit is enforced. Use this to implement "fail if more than X% fail" semantics. + should_complete: Optional predicate for custom completion logic. When + provided, the predicate takes full precedence over the threshold + fields (min_successful, tolerated_failure_count, + tolerated_failure_percentage are ignored). The predicate receives a + CompletionStatus snapshot and returns True to complete the batch + early, or False to continue. The batch still completes once every + item finishes regardless of the predicate's return value. + + The predicate must be deterministic and side-effect-free. It + may be called multiple times per invocation and across + re-invocations (suspend/resume). It must depend only on the + CompletionStatus provided, never on external state. + Note: The operation completes when any of the completion criteria are met: + - Custom predicate returns True (when should_complete is provided) - Enough successes (min_successful reached) - Too many failures (tolerated limits exceeded) - All items/branches completed @@ -131,11 +216,17 @@ class CompletionConfig: min_successful=3, tolerated_failure_count=2 ) + + # Custom predicate: stop once 2 successes are observed + config = CompletionConfig( + should_complete=lambda status: status.success_count >= 2 + ) """ min_successful: int | None = None tolerated_failure_count: int | None = None tolerated_failure_percentage: int | float | None = None + should_complete: Callable[[CompletionStatus], bool] | None = None def __post_init__(self) -> None: if self.min_successful is not None and self.min_successful < 1: @@ -158,6 +249,9 @@ def __post_init__(self) -> None: f"{self.tolerated_failure_percentage}" ) raise ValidationError(msg) + if self.should_complete is not None and not callable(self.should_complete): + msg = "should_complete must be callable" + raise ValidationError(msg) def _validate_for_total(self, total: int) -> None: """Validate this config against the number of items it will govern. @@ -166,7 +260,12 @@ def _validate_for_total(self, total: int) -> None: 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. + + When should_complete is set the predicate takes full precedence and + threshold fields are ignored, so threshold validation is skipped. """ + if self.should_complete is not None: + return if self.min_successful is not None and self.min_successful > total: msg = ( f"min_successful cannot be greater than total items: " 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 0cf8f0ab..874e0284 100644 --- a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py @@ -32,6 +32,7 @@ from aws_durable_execution_sdk_python.config import ( ChildConfig, CompletionConfig, + CompletionStatus, MapConfig, ParallelBranch, ParallelConfig, @@ -2739,6 +2740,124 @@ def test_completion_policy_reason_defaults_to_all_completed(): assert policy.reason(succeeded=1, failed=0) is CompletionReason.ALL_COMPLETED +def test_completion_policy_from_config_copies_should_complete(): + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 2 + + config: CompletionConfig = CompletionConfig(should_complete=predicate) + policy: CompletionPolicy = CompletionPolicy.from_config(5, config) + assert policy.should_complete is predicate + assert policy.min_successful is None + + +def test_completion_policy_should_complete_disables_tolerance(): + """When a predicate is active, tolerance checking is disabled.""" + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 3 + + config: CompletionConfig = CompletionConfig( + should_complete=predicate, + tolerated_failure_count=0, + ) + policy: CompletionPolicy = CompletionPolicy.from_config(5, config) + # Even with tolerated_failure_count=0, failures should not trigger tolerance + assert not policy.is_tolerance_exceeded(failed=5) + assert policy.should_continue(failed=5) + + +def test_completion_policy_should_complete_is_complete_true(): + """Predicate returns True when success threshold is met.""" + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 2 + + policy: CompletionPolicy = CompletionPolicy.from_config( + 5, CompletionConfig(should_complete=predicate) + ) + assert not policy.is_complete(succeeded=1, failed=0) + assert policy.is_complete(succeeded=2, failed=0) + assert policy.is_complete(succeeded=2, failed=1) + + +def test_completion_policy_should_complete_is_complete_all_done(): + """Batch completes when all items finish regardless of predicate.""" + + def predicate(s: CompletionStatus) -> bool: + return False + + policy: CompletionPolicy = CompletionPolicy.from_config( + 3, CompletionConfig(should_complete=predicate) + ) + assert not policy.is_complete(succeeded=1, failed=1) + assert policy.is_complete(succeeded=2, failed=1) + + +def test_completion_policy_should_complete_reason_custom(): + """Reason is CUSTOM_COMPLETION when predicate triggers early exit.""" + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 2 + + policy: CompletionPolicy = CompletionPolicy.from_config( + 5, CompletionConfig(should_complete=predicate) + ) + assert policy.reason(succeeded=2, failed=0) is CompletionReason.CUSTOM_COMPLETION + + +def test_completion_policy_should_complete_reason_all_completed(): + """Reason is ALL_COMPLETED when all items finish without predicate firing.""" + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 10 + + policy: CompletionPolicy = CompletionPolicy.from_config( + 3, CompletionConfig(should_complete=predicate) + ) + assert policy.reason(succeeded=2, failed=1) is CompletionReason.ALL_COMPLETED + + +def test_completion_policy_should_complete_receives_correct_status(): + """Predicate receives a CompletionStatus with correct counts.""" + received: list[CompletionStatus] = [] + + def capture(status: CompletionStatus) -> bool: + received.append(status) + return status.success_count >= 2 + + policy: CompletionPolicy = CompletionPolicy.from_config( + 5, CompletionConfig(should_complete=capture) + ) + result: bool = policy.is_complete(succeeded=1, failed=1) + assert not result + assert len(received) == 1 + assert received[0].success_count == 1 + assert received[0].failure_count == 1 + assert received[0].completed_count == 2 + assert received[0].total_count == 5 + + +def test_completion_policy_should_complete_ignores_threshold_fields(): + """When predicate is present, threshold fields are ignored entirely.""" + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 3 + + config: CompletionConfig = CompletionConfig( + min_successful=1, + tolerated_failure_count=0, + tolerated_failure_percentage=0, + should_complete=predicate, + ) + policy: CompletionPolicy = CompletionPolicy.from_config(5, config) + # min_successful=1 would normally trigger early completion at 1 success + assert not policy.is_complete(succeeded=1, failed=0) + # tolerated_failure_count=0 would normally fail-fast + assert policy.should_continue(failed=3) + # Only the predicate decides + assert policy.is_complete(succeeded=3, failed=0) + + # endregion CompletionPolicy @@ -3731,3 +3850,643 @@ def test_operation_id_namespace_derivation_is_stable(): expected_no_prefix: str = hashlib.blake2b(b"7").hexdigest()[:64] assert OperationIdNamespace(None).create_id_for_step(7) == expected_no_prefix + + +# region Custom completion predicate (should_complete) integration tests + + +def test_executor_should_complete_early_exit_on_success_count(): + """Executor completes early when should_complete predicate returns True.""" + + class TestExecutor(ConcurrentExecutor): + def execute_item(self, child_context, executable): + return f"result_{executable.index}" + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 2 + + completion_config: CompletionConfig = CompletionConfig(should_complete=predicate) + + executables: list[Executable] = [ + Executable(0, lambda: None), + Executable(1, lambda: None), + Executable(2, lambda: None), + Executable(3, lambda: None), + Executable(4, lambda: None), + ] + + executor: TestExecutor = TestExecutor( + executables=executables, + max_concurrency=1, + 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() + + 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: BatchResult = executor.execute(execution_state, executor_context) + + assert result.completion_reason is CompletionReason.CUSTOM_COMPLETION + # With max_concurrency=1 and sequential execution, exactly 2 should succeed + assert result.success_count == 2 + # Remaining items were never started (or still STARTED if raced) + assert result.started_count + result.success_count + result.failure_count <= 5 + + +def test_executor_should_complete_all_finish_when_predicate_never_fires(): + """Batch completes normally when predicate never returns True.""" + + class TestExecutor(ConcurrentExecutor): + def execute_item(self, child_context, executable): + return f"result_{executable.index}" + + # Predicate requires 100 successes - never reachable with 3 items + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 100 + + completion_config: CompletionConfig = CompletionConfig(should_complete=predicate) + + executables: list[Executable] = [ + Executable(0, lambda: None), + Executable(1, lambda: None), + Executable(2, lambda: None), + ] + + executor: TestExecutor = TestExecutor( + executables=executables, + max_concurrency=3, + 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() + + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda *args: "1" + executor_context.create_child_context = lambda *args, **kwargs: Mock() + + result: BatchResult = executor.execute(execution_state, executor_context) + + assert result.completion_reason is CompletionReason.ALL_COMPLETED + assert result.success_count == 3 + assert result.total_count == 3 + + +def test_executor_should_complete_with_mixed_success_and_failure(): + """Predicate can inspect both success and failure counts.""" + + class TestExecutor(ConcurrentExecutor): + def execute_item(self, child_context, executable): + if executable.index % 2 == 0: + return f"result_{executable.index}" + msg: str = f"error_{executable.index}" + raise ValueError(msg) + + # Complete when we have at least 1 success AND 1 failure + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 1 and s.failure_count >= 1 + + completion_config: CompletionConfig = CompletionConfig(should_complete=predicate) + + executables: list[Executable] = [ + Executable(0, lambda: None), + Executable(1, lambda: None), + Executable(2, lambda: None), + Executable(3, lambda: None), + ] + + executor: TestExecutor = TestExecutor( + executables=executables, + max_concurrency=1, + 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() + + 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: BatchResult = executor.execute(execution_state, executor_context) + + assert result.completion_reason is CompletionReason.CUSTOM_COMPLETION + assert result.success_count >= 1 + assert result.failure_count >= 1 + + +def test_executor_should_complete_predicate_ignores_threshold_fields(): + """Threshold fields are ignored when predicate is active in live execution.""" + + class TestExecutor(ConcurrentExecutor): + def execute_item(self, child_context, executable): + if executable.index == 0: + msg: str = "fail" + raise ValueError(msg) + return f"result_{executable.index}" + + # min_successful=1 and tolerated_failure_count=0 would normally cause + # early exit at the first failure, but predicate takes precedence + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 2 + + completion_config: CompletionConfig = CompletionConfig( + min_successful=1, + tolerated_failure_count=0, + should_complete=predicate, + ) + + executables: list[Executable] = [ + Executable(0, lambda: None), + Executable(1, lambda: None), + Executable(2, lambda: None), + Executable(3, lambda: None), + Executable(4, lambda: None), + ] + + executor: TestExecutor = TestExecutor( + executables=executables, + max_concurrency=1, + 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() + + 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: BatchResult = executor.execute(execution_state, executor_context) + + # Predicate waited for 2 successes despite tolerated_failure_count=0 + assert result.completion_reason is CompletionReason.CUSTOM_COMPLETION + assert result.success_count == 2 + assert result.failure_count == 1 + + +def test_batch_result_from_dict_custom_completion_reason(): + """BatchResult deserializes CUSTOM_COMPLETION reason correctly.""" + data: dict = { + "all": [ + {"index": 0, "status": "SUCCEEDED", "result": "ok", "error": None}, + {"index": 1, "status": "SUCCEEDED", "result": "ok", "error": None}, + {"index": 2, "status": "STARTED", "result": None, "error": None}, + ], + "completionReason": "CUSTOM_COMPLETION", + } + result: BatchResult = BatchResult.from_dict(data) + assert result.completion_reason is CompletionReason.CUSTOM_COMPLETION + assert result.success_count == 2 + assert result.started_count == 1 + + +def test_completion_record_custom_completion_round_trip(): + """CompletionRecord correctly parses CUSTOM_COMPLETION from a summary payload.""" + payload: str = json.dumps( + { + "type": "MAP", + "totalCount": 5, + "completionReason": "CUSTOM_COMPLETION", + "startedIndexes": [3, 4], + "startedCount": 2, + "successCount": 2, + "failureCount": 1, + "status": "SUCCEEDED", + } + ) + record: CompletionRecord | None = CompletionRecord.from_summary_payload(payload) + assert record is not None + assert record.completion_reason is CompletionReason.CUSTOM_COMPLETION + assert record.started_total == 5 + assert record.started_indexes == frozenset({3, 4}) + + +def test_executor_should_complete_no_hang_guard(): + """Predicate is not evaluated before any terminal event occurs. + + A predicate like `failure_count == 0` must not fire before work starts. + The predicate only fires after the first COMPLETED/FAILED event, so at + least one branch always runs. An always-True predicate fires on the + first terminal event and stops the batch with exactly 1 completion. + """ + + class TestExecutor(ConcurrentExecutor): + def execute_item(self, child_context, executable): + return f"result_{executable.index}" + + # Predicate always fires — but only after first terminal event + def predicate(s: CompletionStatus) -> bool: + return True + + completion_config: CompletionConfig = CompletionConfig(should_complete=predicate) + + executables: list[Executable] = [ + Executable(0, lambda: None), + Executable(1, lambda: None), + Executable(2, lambda: None), + ] + + executor: TestExecutor = TestExecutor( + executables=executables, + max_concurrency=1, + 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() + + 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: BatchResult = executor.execute(execution_state, executor_context) + + # Predicate fired after first branch completed — exactly 1 success + assert result.completion_reason is CompletionReason.CUSTOM_COMPLETION + assert result.success_count == 1 + assert result.failure_count == 0 + + +def test_executor_should_complete_flat_nesting_type(): + """Predicate works correctly under NestingType.FLAT.""" + + class TestExecutor(ConcurrentExecutor): + def execute_item(self, child_context, executable): + return f"result_{executable.index}" + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 2 + + completion_config: CompletionConfig = CompletionConfig(should_complete=predicate) + + executables: list[Executable] = [ + Executable(0, lambda: None), + Executable(1, lambda: None), + Executable(2, lambda: None), + Executable(3, lambda: None), + ] + + executor: TestExecutor = TestExecutor( + executables=executables, + max_concurrency=1, + completion_config=completion_config, + sub_type_top="TOP", + sub_type_iteration="ITER", + name_prefix="test_", + serdes=None, + operation_id_namespace=_StubNamespace(), + nesting_type=NestingType.FLAT, + ) + + 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: BatchResult = executor.execute(execution_state, executor_context) + + assert result.completion_reason is CompletionReason.CUSTOM_COMPLETION + assert result.success_count == 2 + + +def test_executor_should_complete_index_based_quorum(): + """Predicate can inspect per-item status for quorum rules. + + Quorum rule: complete when branch 0 succeeds OR (branches 1 AND 2 both succeed). + Branch 0 fails, branches 1 and 2 succeed, triggering the quorum. + """ + + class TestExecutor(ConcurrentExecutor): + def execute_item(self, child_context, executable): + if executable.index == 0: + msg: str = "branch 0 fails" + raise ValueError(msg) + return f"result_{executable.index}" + + def quorum_predicate(s: CompletionStatus) -> bool: + # Branch 0 succeeds OR (branches 1 AND 2 both succeed) + if len(s.items) == 0: + return False + branch_0_ok: bool = len(s.items) > 0 and s.items[0].is_succeeded + branch_1_ok: bool = len(s.items) > 1 and s.items[1].is_succeeded + branch_2_ok: bool = len(s.items) > 2 and s.items[2].is_succeeded + return branch_0_ok or (branch_1_ok and branch_2_ok) + + completion_config: CompletionConfig = CompletionConfig( + should_complete=quorum_predicate + ) + + executables: list[Executable] = [ + Executable(0, lambda: None), + Executable(1, lambda: None), + Executable(2, lambda: None), + Executable(3, lambda: None), + ] + + executor: TestExecutor = TestExecutor( + executables=executables, + max_concurrency=1, + 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() + + 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: BatchResult = executor.execute(execution_state, executor_context) + + # Branch 0 failed, branches 1+2 succeeded -> quorum met + assert result.completion_reason is CompletionReason.CUSTOM_COMPLETION + assert result.failure_count == 1 + assert result.success_count == 2 + + +def test_executor_should_complete_predicate_raises_propagates(): + """A predicate that raises propagates the exception to the caller. + + Documents the current behavior: the predicate must not raise. If it + does, the exception escapes the coordinator loop unmodified. + """ + + class TestExecutor(ConcurrentExecutor): + def execute_item(self, child_context, executable): + return f"result_{executable.index}" + + def bad_predicate(s: CompletionStatus) -> bool: + msg: str = "predicate bug" + raise RuntimeError(msg) + + completion_config: CompletionConfig = CompletionConfig( + should_complete=bad_predicate + ) + + executables: list[Executable] = [ + Executable(0, lambda: None), + Executable(1, lambda: None), + ] + + executor: TestExecutor = TestExecutor( + executables=executables, + max_concurrency=1, + 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() + + 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 + + with pytest.raises(RuntimeError, match="predicate bug"): + executor.execute(execution_state, executor_context) + + +def test_validate_for_total_skipped_when_should_complete_set(): + """_validate_for_total does not raise when should_complete is set.""" + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 1 + + config: CompletionConfig = CompletionConfig( + min_successful=100, + should_complete=predicate, + ) + # Would raise without should_complete since min_successful > total + config._validate_for_total(5) + + +def test_validate_for_total_raises_without_should_complete(): + """_validate_for_total raises when min_successful > total without predicate.""" + config: CompletionConfig = CompletionConfig(min_successful=10) + with pytest.raises(ValidationError): + config._validate_for_total(5) + + +def test_replay_round_trip_custom_completion_without_reinvoking_predicate(): + """Replay reconstructs a CUSTOM_COMPLETION batch from the checkpointed record. + + The predicate must NOT be re-invoked during replay. This test proves + that by using a predicate that would give a different answer if called + during replay (it counts invocations), yet the replayed result matches + the live result exactly. + """ + call_count: list[int] = [0] + + def counting_predicate(s: CompletionStatus) -> bool: + call_count[0] += 1 + return s.success_count >= 2 + + executables: list[Executable] = [ + Executable(0, partial(lambda i: f"live_{i}", 0)), + Executable(1, partial(lambda i: f"live_{i}", 1)), + Executable(2, partial(lambda i: f"live_{i}", 2)), + Executable(3, partial(lambda i: f"live_{i}", 3)), + ] + executor = _make_executor( + executables, + max_concurrency=1, + completion_config=CompletionConfig(should_complete=counting_predicate), + ) + execution_state, executor_context = _make_executor_mocks() + live: BatchResult = executor.execute(execution_state, executor_context) + + # Live run should have completed early at 2 successes + assert live.completion_reason is CompletionReason.CUSTOM_COMPLETION + assert live.success_count == 2 + live_calls: int = call_count[0] + assert live_calls > 0 + + # Build the summary envelope (as would be checkpointed) + summary: str = envelope_summary_generator("MapResult", None)(live) + top_checkpoint = Mock() + top_checkpoint.result = summary + + # Replay: set up branch checkpoints for the terminal items + replay_state, replay_context = _make_replay_mocks( + { + "op_0": _succeeded_checkpoint('"replayed_0"'), + "op_1": _succeeded_checkpoint('"replayed_1"'), + } + ) + replay_executor = _make_executor( + executables, + max_concurrency=1, + completion_config=CompletionConfig(should_complete=counting_predicate), + ) + + # Reset call count to detect any predicate calls during replay + call_count[0] = 0 + replayed: BatchResult = replay_executor.replay( + replay_state, replay_context, top_checkpoint + ) + + # Predicate was NOT called during replay + assert call_count[0] == 0 + + # Replayed result matches live result structure + assert replayed.completion_reason is CompletionReason.CUSTOM_COMPLETION + assert replayed.success_count == live.success_count + assert [(i.index, i.status) for i in replayed.all] == [ + (i.index, i.status) for i in live.all + ] + + +def test_suspend_resume_mid_batch_with_should_complete(): + """Predicate fires on a later invocation after a branch suspends and resumes. + + Branch 0 suspends (timed), branch 1 succeeds, and after the timed + resume branch 0 succeeds - predicate fires at 2 successes. + """ + call_counts: dict[int, int] = {0: 0} + + def timed_then_ok(): + call_counts[0] += 1 + if call_counts[0] == 1: + msg: str = "retry shortly" + raise TimedSuspendExecution(msg, time.time() + 0.1) + return "branch_0_done" + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 2 + + executables: list[Executable] = [ + Executable(0, timed_then_ok), + Executable(1, lambda: "branch_1_done"), + Executable(2, lambda: "branch_2_done"), + Executable(3, lambda: "branch_3_done"), + ] + executor = _make_executor( + executables, + max_concurrency=2, + completion_config=CompletionConfig(should_complete=predicate), + ) + execution_state, executor_context = _make_executor_mocks() + + result: BatchResult = executor.execute(execution_state, executor_context) + + # Branch 0 suspended then resumed; the predicate fired at 2 successes + assert result.completion_reason is CompletionReason.CUSTOM_COMPLETION + assert result.success_count >= 2 + + +def test_executor_should_complete_items_reflect_started_status(): + """Predicate sees STARTED (not None) for submitted branches and SUCCEEDED + for completed branches after a terminal event triggers a snapshot rebuild. + """ + observed_items: list[tuple] = [] + + def inspecting_predicate(s: CompletionStatus) -> bool: + # Only capture after at least one branch has completed + if not observed_items and s.success_count >= 1 and s.items: + for item in s.items: + observed_items.append((item.index, item.status)) + return s.success_count >= 2 + + class TestExecutor(ConcurrentExecutor): + def execute_item(self, child_context, executable): + return f"result_{executable.index}" + + executables: list[Executable] = [ + Executable(0, lambda: None), + Executable(1, lambda: None), + Executable(2, lambda: None), + Executable(3, lambda: None), + ] + + executor: TestExecutor = TestExecutor( + executables=executables, + max_concurrency=2, + completion_config=CompletionConfig(should_complete=inspecting_predicate), + 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" + 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: BatchResult = executor.execute(execution_state, executor_context) + + assert result.completion_reason is CompletionReason.CUSTOM_COMPLETION + + # After a terminal event the snapshot was rebuilt from live branch state. + # Verify predicate observed correct statuses: + assert len(observed_items) == 4 + statuses: set = {status for _, status in observed_items} + from aws_durable_execution_sdk_python.config import BatchItemStatus + + # At least one SUCCEEDED branch must be visible + assert BatchItemStatus.SUCCEEDED in statuses + # Branches not yet submitted should be None (PENDING internally) + assert None in statuses + + +# endregion Custom completion predicate (should_complete) integration tests 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 02f69a7f..fa35d3e1 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 @@ -17,6 +17,7 @@ from aws_durable_execution_sdk_python.identifier import OperationIdNamespace from aws_durable_execution_sdk_python.config import ( CompletionConfig, + CompletionStatus, MapConfig, NestingType, ) @@ -1482,3 +1483,94 @@ def capturing_from_items(cls, items, func, config, operation_id_namespace): assert captured_configs[0].summary_generator is None assert captured_configs[0].max_concurrency == 5 assert result.total_count == 0 + + +# region Custom completion predicate (should_complete) tests + + +def test_map_handler_with_should_complete_predicate(): + """Test map_handler passes should_complete through to executor.""" + items: list[str] = ["a", "b", "c", "d", "e"] + + def callable_func(ctx, item, idx, items): + return f"result_{item}" + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 2 + + config: MapConfig = MapConfig( + max_concurrency=1, + completion_config=CompletionConfig(should_complete=predicate), + ) + + mock_batch_result: BatchResult = BatchResult( + all=[ + BatchItem(index=0, status=BatchItemStatus.SUCCEEDED, result="result_a"), + BatchItem(index=1, status=BatchItemStatus.SUCCEEDED, result="result_b"), + BatchItem(index=2, status=BatchItemStatus.STARTED), + BatchItem(index=3, status=BatchItemStatus.STARTED), + BatchItem(index=4, status=BatchItemStatus.STARTED), + ], + completion_reason=CompletionReason.CUSTOM_COMPLETION, + ) + + executor_context = Mock() + + with patch.object( + MapExecutor, "execute", return_value=mock_batch_result + ) 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 + return mock_result + + execution_state = MockExecutionState() + operation_identifier: OperationIdentifier = OperationIdentifier( + "test_op", OperationSubType.MAP, "parent", "test_map" + ) + + result: BatchResult = map_handler( + items, + callable_func, + config, + execution_state, + executor_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), + ) + + mock_execute.assert_called_once() + assert result.completion_reason is CompletionReason.CUSTOM_COMPLETION + assert result.success_count == 2 + + +def test_map_executor_from_items_preserves_should_complete(): + """Test MapExecutor.from_items passes should_complete to the policy.""" + items: list[str] = ["a", "b", "c"] + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 1 + + def callable_func(ctx, item, idx, items): + return item.upper() + + config: MapConfig = MapConfig( + completion_config=CompletionConfig(should_complete=predicate), + ) + + executor: MapExecutor = MapExecutor.from_items( + items, + callable_func, + config, + operation_id_namespace=_StubNamespace(), + ) + + assert executor.policy.should_complete is predicate + + +# endregion Custom completion predicate (should_complete) tests 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 a19f87e8..52bfd811 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 @@ -22,6 +22,7 @@ ) from aws_durable_execution_sdk_python.config import ( CompletionConfig, + CompletionStatus, NestingType, ParallelConfig, ) @@ -1475,3 +1476,91 @@ def capturing_from_callables(cls, callables, config, operation_id_namespace): assert captured_configs[0].summary_generator is None assert captured_configs[0].max_concurrency == 3 assert result.total_count == 0 + + +# region Custom completion predicate (should_complete) tests + + +def test_parallel_handler_with_should_complete_predicate(): + """Test parallel_handler passes should_complete through to executor.""" + callables = [ + lambda ctx: "result_a", + lambda ctx: "result_b", + lambda ctx: "result_c", + lambda ctx: "result_d", + ] + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 2 + + config: ParallelConfig = ParallelConfig( + max_concurrency=1, + completion_config=CompletionConfig(should_complete=predicate), + ) + + mock_batch_result: BatchResult = BatchResult( + all=[ + BatchItem(index=0, status=BatchItemStatus.SUCCEEDED, result="result_a"), + BatchItem(index=1, status=BatchItemStatus.SUCCEEDED, result="result_b"), + BatchItem(index=2, status=BatchItemStatus.STARTED), + BatchItem(index=3, status=BatchItemStatus.STARTED), + ], + completion_reason=CompletionReason.CUSTOM_COMPLETION, + ) + + executor_context = Mock() + + with patch.object( + ParallelExecutor, "execute", return_value=mock_batch_result + ) 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 + return mock_result + + execution_state = MockExecutionState() + operation_identifier: OperationIdentifier = OperationIdentifier( + "test_op", OperationSubType.PARALLEL, "parent", "test_parallel" + ) + + result: BatchResult = parallel_handler( + callables, + config, + execution_state, + executor_context, + operation_identifier, + operation_id_namespace=_StubNamespace(), + ) + + mock_execute.assert_called_once() + assert result.completion_reason is CompletionReason.CUSTOM_COMPLETION + assert result.success_count == 2 + + +def test_parallel_executor_from_callables_preserves_should_complete(): + """Test ParallelExecutor.from_callables passes should_complete to the policy.""" + + def predicate(s: CompletionStatus) -> bool: + return s.success_count >= 1 + + config: ParallelConfig = ParallelConfig( + completion_config=CompletionConfig(should_complete=predicate), + ) + + callables = [lambda ctx: "a", lambda ctx: "b"] + + executor: ParallelExecutor = ParallelExecutor.from_callables( + callables, + config, + operation_id_namespace=_StubNamespace(), + ) + + assert executor.policy.should_complete is predicate + + +# endregion Custom completion predicate (should_complete) tests