feat: add should_complete to CompletionConfig (#605)

Author: Ayushi Ahjolia

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",

packages/aws-durable-execution-sdk-python-examples/src/map/map_with_should_complete.py

@@ -0,0 +1,53 @@
+"""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,
+    CompletionDecision,
+    CompletionStatus,
+    MapConfig,
+    complete_batch,
+    continue_batch,
+)
+from aws_durable_execution_sdk_python.context import DurableContext
+from aws_durable_execution_sdk_python.execution import durable_execution
+
+
+def _should_complete(status: CompletionStatus) -> CompletionDecision:
+    """Complete once 3 items have succeeded."""
+    return complete_batch() if status.success_count >= 3 else continue_batch()
+
+
+@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(),
+    }

packages/aws-durable-execution-sdk-python-examples/src/parallel/parallel_with_should_complete.py

@@ -0,0 +1,70 @@
+"""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 (
+    BatchItemStatus,
+    CompletionConfig,
+    CompletionDecision,
+    CompletionStatus,
+    ParallelConfig,
+    complete_batch,
+    continue_batch,
+)
+from aws_durable_execution_sdk_python.context import DurableContext
+from aws_durable_execution_sdk_python.execution import durable_execution
+
+
+def _quorum_predicate(status: CompletionStatus) -> CompletionDecision:
+    """Complete when branch A succeeds OR both B and C succeed."""
+    if not status.items:
+        return continue_batch()
+    branch_a_ok: bool = status.items[0].status is BatchItemStatus.SUCCEEDED
+    branch_b_ok: bool = (
+        len(status.items) > 1 and status.items[1].status is BatchItemStatus.SUCCEEDED
+    )
+    branch_c_ok: bool = (
+        len(status.items) > 2 and status.items[2].status is BatchItemStatus.SUCCEEDED
+    )
+    return (
+        complete_batch()
+        if branch_a_ok or (branch_b_ok and branch_c_ok)
+        else continue_batch()
+    )
+
+
+@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(),
+    }

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_SUCCEEDED"
+
+    # 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

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_SUCCEEDED"
+    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

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py

@@ -6,8 +6,20 @@
 # Main context - used in every durable function
 # 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.concurrency.models import (
+    BatchCompletionError,
+    BatchResult,
+)
+from aws_durable_execution_sdk_python.config import (
+    BatchItemStatus,
+    CompletionDecision,
+    CompletionItemStatus,
+    CompletionOutcome,
+    CompletionStatus,
+    ParallelBranch,
+    complete_batch,
+    continue_batch,
+)
 from aws_durable_execution_sdk_python.context import (
     DurableContext,
     durable_parallel_branch,
@@ -45,12 +57,18 @@
 
 
 __all__ = [
+    "BatchCompletionError",
+    "BatchItemStatus",
     "BatchResult",
     "CallbackError",
     "CallbackExternalError",
     "CallbackSubmitterError",
     "CallbackTimeoutError",
     "ChildContextError",
+    "CompletionDecision",
+    "CompletionItemStatus",
+    "CompletionOutcome",
+    "CompletionStatus",
     "DurableContext",
     "DurableExecutionsError",
     "DurableOperationError",
@@ -67,6 +85,8 @@
     "WaitForConditionError",
     "WithRetryConfig",
     "__version__",
+    "complete_batch",
+    "continue_batch",
     "durable_execution",
     "durable_parallel_branch",
     "durable_step",