Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
@@ -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(),
}
Original file line number Diff line number Diff line change
@@ -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(),
}
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -42,12 +47,15 @@


__all__ = [
"BatchItemStatus",
"BatchResult",
"CallbackError",
"CallbackExternalError",
"CallbackSubmitterError",
"CallbackTimeoutError",
"ChildContextError",
"CompletionItemStatus",
"CompletionStatus",
"DurableContext",
"DurableExecutionsError",
"DurableOperationError",
Expand Down
Loading
Loading