feat: add distributed map operation

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

@@ -6,8 +6,27 @@
 # 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 (
+    BatchResult,
+    DistributedMapCompletionReason,
+    DistributedMapItemError,
+    DistributedMapResult,
+    DistributedMapResultItem,
+    DistributedMapStatus,
+    DistributedMapSummary,
+)
+from aws_durable_execution_sdk_python.config import (
+    FailureDestination,
+    DistributedMapCompletionConfig,
+    DistributedMapConfig,
+    DistributedMapDestination,
+    DistributedMapDestinationConfig,
+    DistributedMapProcessor,
+    DistributedMapSource,
+    ParallelBranch,
+    ProcessorRetryConfig,
+    SuccessDestination,
+)
 from aws_durable_execution_sdk_python.context import (
     DurableContext,
     durable_parallel_branch,
@@ -28,6 +47,7 @@
     ExecutionError,
     InvocationError,
     InvokeError,
+    DistributedMapError,
     PluginLoadError,
     RetryableSerDesError,
     SerDesError,
@@ -55,14 +75,30 @@
     "DurableExecutionsError",
     "DurableOperationError",
     "ExecutionError",
+    "FailureDestination",
     "InvocationError",
     "InvokeError",
+    "DistributedMapCompletionConfig",
+    "DistributedMapCompletionReason",
+    "DistributedMapConfig",
+    "DistributedMapDestination",
+    "DistributedMapDestinationConfig",
+    "DistributedMapError",
+    "DistributedMapItemError",
+    "DistributedMapProcessor",
+    "DistributedMapResult",
+    "DistributedMapResultItem",
+    "DistributedMapSource",
+    "DistributedMapStatus",
+    "DistributedMapSummary",
     "ParallelBranch",
     "PluginLoadError",
+    "ProcessorRetryConfig",
     "RetryableSerDesError",
     "SerDesError",
     "StepContext",
     "StepError",
+    "SuccessDestination",
     "ValidationError",
     "WaitForConditionError",
     "WithRetryConfig",

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

@@ -5,13 +5,14 @@
 import json
 import logging
 from collections import Counter
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum
-from typing import TYPE_CHECKING, Generic, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 from aws_durable_execution_sdk_python.exceptions import (
     ChildContextError,
     InvalidStateError,
+    DistributedMapError,
 )
 from aws_durable_execution_sdk_python.lambda_service import ErrorObject
 from aws_durable_execution_sdk_python.types import BatchResult as BatchResultProtocol
@@ -550,3 +551,137 @@ def fatal(cls, index: int, error: BaseException) -> BranchEvent[ResultType]:
 
 
 # endregion concurrency models
+
+
+# region map run result models
+class DistributedMapStatus(Enum):
+    SUCCEEDED = "SUCCEEDED"
+    FAILED = "FAILED"
+    STOPPED = "STOPPED"
+    TIMED_OUT = "TIMED_OUT"
+
+
+class DistributedMapCompletionReason(Enum):
+    ALL_COMPLETED = "ALL_COMPLETED"
+    ITEM_LIMIT_REACHED = "ITEM_LIMIT_REACHED"
+    STOPPED = "STOPPED"
+    TIMED_OUT = "TIMED_OUT"
+    FAILURE_TOLERANCE_EXCEEDED = "FAILURE_TOLERANCE_EXCEEDED"
+    SOURCE_FAILED = "SOURCE_FAILED"
+    DESTINATION_FAILED = "DESTINATION_FAILED"
+    INLINE_RESULT_LIMIT_EXCEEDED = "INLINE_RESULT_LIMIT_EXCEEDED"
+    INVALID_CONFIGURATION = "INVALID_CONFIGURATION"
+    QUOTA_EXCEEDED = "QUOTA_EXCEEDED"
+    KMS_ACCESS_DENIED = "KMS_ACCESS_DENIED"
+    INTERNAL_ERROR = "INTERNAL_ERROR"
+
+
+@dataclass(frozen=True)
+class DistributedMapSummary:
+    """Outcome of a map run without per-item results.
+
+    Resolved by ``ctx.distributed_map`` for every terminal state. A non-``SUCCEEDED``
+    run resolves with this summary rather than raising. Use
+    :meth:`throw_if_error` to opt into raising.
+    """
+
+    status: DistributedMapStatus
+    completion_reason: DistributedMapCompletionReason
+    success_count: int
+    failure_count: int
+    unprocessed_count: int
+    distributed_map_run_arn: str | None = None
+    completion_details: str | None = None
+    total_count: int | None = None
+
+    @property
+    def distributed_map_id(self) -> str | None:
+        """Map-run id derived from the ARN, or ``None`` when the ARN is absent."""
+        if not self.distributed_map_run_arn:
+            return None
+        return self.distributed_map_run_arn.rsplit(":", 1)[-1]
+
+    @property
+    def has_failure(self) -> bool:
+        """``True`` when any item permanently failed."""
+        return self.failure_count > 0
+
+    def throw_if_error(self) -> None:
+        """Raise :class:`DistributedMapError` on any non-success outcome."""
+        if self.status is not DistributedMapStatus.SUCCEEDED:
+            detail = f", {self.completion_details}" if self.completion_details else ""
+            msg = (
+                f"Map run ended {self.status.value} "
+                f"(reason: {self.completion_reason.value}{detail})"
+            )
+            raise DistributedMapError(msg)
+        if self.failure_count > 0:
+            msg = (
+                f"Map run succeeded but {self.failure_count} item(s) permanently failed"
+            )
+            raise DistributedMapError(msg)
+
+
+@dataclass(frozen=True)
+class DistributedMapItemError:
+    """Error for a single failed map run item."""
+
+    error_type: str
+    error_message: str
+
+
+@dataclass(frozen=True)
+class DistributedMapResultItem:
+    """Outcome of a single map run item."""
+
+    item_id: str
+    status: str
+    output: Any | None = None
+    error: DistributedMapItemError | None = None
+
+
+@dataclass(frozen=True)
+class DistributedMapResult(DistributedMapSummary):
+    """Outcome of a map run with per-item results.
+
+    Resolved by ``ctx.distributed_map`` when inline result collection is enabled.
+    Extends :class:`DistributedMapSummary` with the retained per-item results.
+    """
+
+    all: list[DistributedMapResultItem] = field(default_factory=list)
+
+    def succeeded(self) -> list[DistributedMapResultItem]:
+        """Return the items that succeeded."""
+        return [item for item in self.all if item.status == "SUCCEEDED"]
+
+    def failed(self) -> list[DistributedMapResultItem]:
+        """Return the items that permanently failed."""
+        return [item for item in self.all if item.status == "FAILED"]
+
+    def get_results(self) -> list[Any]:
+        """Return the outputs of the succeeded items."""
+        return [item.output for item in self.all if item.status == "SUCCEEDED"]
+
+    def get_errors(self) -> list[DistributedMapItemError]:
+        """Return the errors of the failed items."""
+        return [
+            item.error
+            for item in self.all
+            if item.status == "FAILED" and item.error is not None
+        ]
+
+    def throw_if_error(self) -> None:
+        """Raise the first failed item's error, otherwise defer to the summary rule."""
+        if self.status is DistributedMapStatus.SUCCEEDED and self.failure_count > 0:
+            failed = self.failed()
+            if failed:
+                first = failed[0]
+                if first.error is not None:
+                    msg = f"{first.error.error_type}: {first.error.error_message}"
+                    raise DistributedMapError(msg)
+                msg = f"item {first.item_id} failed"
+                raise DistributedMapError(msg)
+        super().throw_if_error()
+
+
+# endregion map run result models