From 535d35c856580d4e3b6649936b33272a2a383f65 Mon Sep 17 00:00:00 2001 From: nvasiu Date: Thu, 20 Aug 2026 18:53:50 +0000 Subject: [PATCH] feat: add distributed map operation - Add ctx.distributed_map with inline, S3, and reader sources - Add DistributedMapConfig, processor, completion, and destination config types - Add DistributedMapResult/Summary result types and DistributedMapError - Add function-authoring helpers for item and batch handlers - Serialize the DISTRIBUTED_MAP operation and add its executor --- .../__init__.py | 40 +- .../concurrency/models.py | 139 +- .../config.py | 532 +++++- .../context.py | 53 + .../distributed_map_helpers.py | 241 +++ .../exceptions.py | 4 + .../lambda_service.py | 343 ++++ .../operation/distributed_map.py | 502 +++++ .../plugin.py | 1 + .../aws_durable_execution_sdk_python/state.py | 3 + .../tests/context_test.py | 36 + .../tests/distributed_map_helpers_test.py | 114 ++ .../e2e/distributed_map_helpers_int_test.py | 148 ++ .../tests/e2e/distributed_map_int_test.py | 246 +++ .../tests/operation/distributed_map_test.py | 1699 +++++++++++++++++ 15 files changed, 4095 insertions(+), 6 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/distributed_map_helpers.py create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/distributed_map.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/distributed_map_helpers_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/distributed_map_helpers_int_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/distributed_map_int_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/distributed_map_test.py 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 a464c419..8394481e 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 @@ -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", 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..15fd9804 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 @@ -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 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..c04f03f9 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 @@ -6,7 +6,7 @@ import random from dataclasses import dataclass, field from enum import Enum, StrEnum -from typing import TYPE_CHECKING, Generic, TypeVar +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypeVar from aws_durable_execution_sdk_python.exceptions import ValidationError @@ -16,7 +16,7 @@ T = TypeVar("T") if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Sequence from aws_durable_execution_sdk_python.lambda_service import OperationSubType from aws_durable_execution_sdk_python.retries import RetryDecision @@ -327,6 +327,534 @@ class StepConfig: serdes: SerDes | None = None +# region map run configuration + + +def _parse_s3_uri(uri: str) -> tuple[str, str | None]: + """Split an ``s3://bucket/path`` URI into its bucket and path parts.""" + rest = uri.removeprefix("s3://") + bucket, _, path = rest.partition("/") + if not bucket: + msg = f"Invalid S3 URI: {uri}" + raise ValidationError(msg) + return bucket, path or None + + +def _validate_bucket_owner(value: str | None) -> None: + """Validate an expected bucket owner is a 12-digit account id.""" + if value is not None and (len(value) != 12 or not value.isdigit()): # noqa: PLR2004 + msg = f"expected_bucket_owner must be a 12-digit account id, got: {value}" + raise ValidationError(msg) + + +def _validate_columns(name: str, columns: tuple[str, ...] | None) -> None: + """Validate a CSV columns/headers tuple is non-empty with no duplicates.""" + if columns is None: + return + if len(columns) == 0: + msg = f"{name} must be non-empty" + raise ValidationError(msg) + if len(set(columns)) != len(columns): + msg = f"{name} must not contain duplicates" + raise ValidationError(msg) + + +def _validate_function_name(name: str) -> None: + """Validate a Lambda function reference is present and correctly formatted.""" + if not name: + msg = "function name must be non-empty" + raise ValidationError(msg) + from aws_durable_execution_sdk_python.lambda_service import ( + lambda_function_name_matcher, + ) + + pattern, max_len = lambda_function_name_matcher() + if len(name) > max_len or pattern.fullmatch(name) is None: + msg = f"invalid Lambda function reference: {name!r}" + raise ValidationError(msg) + + +@dataclass(frozen=True) +class DistributedMapCompletionConfig: + """Failure-tolerance configuration for a map run.""" + + tolerated_failure_count: int | None = None + tolerated_failure_percentage: float | None = None + minimum_sample_size: int | None = None + + def __post_init__(self) -> None: + if ( + self.tolerated_failure_count is not None + and self.tolerated_failure_percentage is not None + ): + msg = ( + "tolerated_failure_count and tolerated_failure_percentage " + "are mutually exclusive" + ) + raise ValidationError(msg) + if ( + self.minimum_sample_size is not None + and self.tolerated_failure_percentage is None + ): + msg = "minimum_sample_size is only valid with tolerated_failure_percentage" + 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) + if self.minimum_sample_size is not None and self.minimum_sample_size < 1: + msg = ( + "minimum_sample_size must be at least 1, got: " + f"{self.minimum_sample_size}" + ) + raise ValidationError(msg) + + @staticmethod + def failure_count(count: int) -> DistributedMapCompletionConfig: + """Abort once this many items have permanently failed.""" + return DistributedMapCompletionConfig(tolerated_failure_count=count) + + @staticmethod + def failure_percentage( + percentage: float, *, minimum_sample_size: int | None = None + ) -> DistributedMapCompletionConfig: + """Abort once the failure rate exceeds this percentage.""" + return DistributedMapCompletionConfig( + tolerated_failure_percentage=percentage, + minimum_sample_size=minimum_sample_size, + ) + + +@dataclass(frozen=True) +class ProcessorRetryConfig: + """Retry configuration for a map run processor.""" + + UNLIMITED: ClassVar[str] = "unlimited" + + max_retry_attempts: int | Literal["unlimited"] | None = None + max_retry_duration: Duration | None = None + + def __post_init__(self) -> None: + if isinstance(self.max_retry_attempts, int) and self.max_retry_attempts < 0: + msg = ( + "max_retry_attempts must be non-negative or " + "ProcessorRetryConfig.UNLIMITED, " + f"got: {self.max_retry_attempts}" + ) + raise ValidationError(msg) + if self.max_retry_duration is not None and not ( + 60 <= self.max_retry_duration.to_seconds() <= 21600 # noqa: PLR2004 + ): + msg = ( + "max_retry_duration must be between 1 minute and 6 hours, got: " + f"{self.max_retry_duration.to_seconds()}s" + ) + raise ValidationError(msg) + + +@dataclass(frozen=True) +class DistributedMapProcessor: + """Processor configuration for a map run.""" + + function_name: str + response_mode: str | None = None # None = batch mode + batch_size: int | None = None + retry: ProcessorRetryConfig | None = None + durable_execution_name_prefix: str | None = None + + def __post_init__(self) -> None: + _validate_function_name(self.function_name) + if self.batch_size is not None and not ( + 1 <= self.batch_size <= 10000 # noqa: PLR2004 + ): + msg = f"batch_size must be between 1 and 10000, got: {self.batch_size}" + raise ValidationError(msg) + + @classmethod + def report_batch_outcome( + cls, + name: str, + *, + batch_size: int | None = None, + retry: ProcessorRetryConfig | None = None, + durable_execution_name_prefix: str | None = None, + ) -> DistributedMapProcessor: + """Processor that reports a single pass/fail outcome for the whole batch, with no per-item results.""" + return cls( + function_name=name, + response_mode=None, + batch_size=batch_size, + retry=retry, + durable_execution_name_prefix=durable_execution_name_prefix, + ) + + @classmethod + def report_failed_items( + cls, + name: str, + *, + batch_size: int | None = None, + retry: ProcessorRetryConfig | None = None, + durable_execution_name_prefix: str | None = None, + ) -> DistributedMapProcessor: + """Processor that reports the ids of failed items, with all others marked succeeded.""" + return cls( + function_name=name, + response_mode="ReportBatchItemFailures", + batch_size=batch_size, + retry=retry, + durable_execution_name_prefix=durable_execution_name_prefix, + ) + + @classmethod + def report_item_results( + cls, + name: str, + *, + batch_size: int | None = None, + retry: ProcessorRetryConfig | None = None, + durable_execution_name_prefix: str | None = None, + ) -> DistributedMapProcessor: + """Processor that reports the results (output or error) for every item.""" + return cls( + function_name=name, + response_mode="ReportBatchItemResults", + batch_size=batch_size, + retry=retry, + durable_execution_name_prefix=durable_execution_name_prefix, + ) + + +@dataclass(frozen=True) +class S3SourceConfig: + """Resolved S3 source configuration.""" + + bucket: str + key: str | None = None + prefix: str | None = None + transform: str | None = None + fmt: str | None = None + delimiter: str | None = None + headers: tuple[str, ...] | None = None + expected_bucket_owner: str | None = None + + def __post_init__(self) -> None: + _validate_bucket_owner(self.expected_bucket_owner) + _validate_columns("headers", self.headers) + + +@dataclass(frozen=True) +class ReaderSourceConfig: + """Resolved reader-function source configuration.""" + + function_name: str + initial_state: Any = None + state_serdes: SerDes | None = None # None = DEFAULT_JSON_SERDES + + def __post_init__(self) -> None: + _validate_function_name(self.function_name) + + +@dataclass(frozen=True) +class DistributedMapSource: + """Source configuration for a map run.""" + + source_type: str + max_items: int | None = None + inline_items: tuple[Any, ...] | None = None + inline_serdes: SerDes | None = None # None = DEFAULT_JSON_SERDES + s3: S3SourceConfig | None = None + reader: ReaderSourceConfig | None = None + + def __post_init__(self) -> None: + if self.max_items is not None and self.max_items < 1: + msg = f"max_items must be at least 1, got: {self.max_items}" + raise ValidationError(msg) + + @classmethod + def inline( + cls, + items: Sequence[Any], + *, + serdes: SerDes | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """An in-memory list of items embedded in the start checkpoint.""" + return cls( + source_type="INLINE", + inline_items=tuple(items), + inline_serdes=serdes, + max_items=max_items, + ) + + class S3: + """S3 source factories.""" + + @staticmethod + def json_lines( + uri: str, + *, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read a single object, treating each line as an item.""" + bucket, key = _parse_s3_uri(uri) + if key is None: + msg = "json_lines requires an S3 object key" + raise ValidationError(msg) + return DistributedMapSource( + source_type="S3", + max_items=max_items, + s3=S3SourceConfig( + bucket=bucket, + key=key, + fmt="JSON_LINES", + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @staticmethod + def csv( + uri: str, + *, + headers: Sequence[str] | None = None, + delimiter: str = "COMMA", + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read a single object, treating each record as an item.""" + bucket, key = _parse_s3_uri(uri) + if key is None: + msg = "csv requires an S3 object key" + raise ValidationError(msg) + return DistributedMapSource( + source_type="S3", + max_items=max_items, + s3=S3SourceConfig( + bucket=bucket, + key=key, + fmt="CSV", + delimiter=delimiter, + headers=tuple(headers) if headers is not None else None, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @staticmethod + def objects( + prefix_uri: str, + *, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read each object under a prefix as one item.""" + bucket, prefix = _parse_s3_uri(prefix_uri) + return DistributedMapSource( + source_type="S3", + max_items=max_items, + s3=S3SourceConfig( + bucket=bucket, + prefix=prefix or "", + transform="NONE", + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @staticmethod + def flattened_json_lines( + prefix_uri: str, + *, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read a prefix, flattening each object's lines into items.""" + bucket, prefix = _parse_s3_uri(prefix_uri) + return DistributedMapSource( + source_type="S3", + max_items=max_items, + s3=S3SourceConfig( + bucket=bucket, + prefix=prefix or "", + transform="LOAD_AND_FLATTEN", + fmt="JSON_LINES", + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @staticmethod + def flattened_csv( + prefix_uri: str, + *, + headers: Sequence[str] | None = None, + delimiter: str = "COMMA", + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read a prefix, flattening each object's records into items.""" + bucket, prefix = _parse_s3_uri(prefix_uri) + return DistributedMapSource( + source_type="S3", + max_items=max_items, + s3=S3SourceConfig( + bucket=bucket, + prefix=prefix or "", + transform="LOAD_AND_FLATTEN", + fmt="CSV", + delimiter=delimiter, + headers=tuple(headers) if headers is not None else None, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + class Reader: + """Reader-function source factories.""" + + @staticmethod + def from_function( + name: str, + *, + initial_state: Any = None, + state_serdes: SerDes | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Page items from a customer-supplied reader Lambda function.""" + return DistributedMapSource( + source_type="READER_FUNCTION", + max_items=max_items, + reader=ReaderSourceConfig( + function_name=name, + initial_state=initial_state, + state_serdes=state_serdes, + ), + ) + + +@dataclass(frozen=True) +class SuccessDestination: + """S3 destination for succeeded items.""" + + bucket: str + prefix: str + include_input: bool = False + include_output: bool = True + expected_bucket_owner: str | None = None + + def __post_init__(self) -> None: + _validate_bucket_owner(self.expected_bucket_owner) + if not (self.include_input or self.include_output): + msg = "success destination must include input or output" + raise ValidationError(msg) + + +@dataclass(frozen=True) +class FailureDestination: + """S3 destination for permanently-failed items.""" + + bucket: str + prefix: str + include_input: bool = True + include_error: bool = True + expected_bucket_owner: str | None = None + + def __post_init__(self) -> None: + _validate_bucket_owner(self.expected_bucket_owner) + if not (self.include_input or self.include_error): + msg = "failure destination must include input or error" + raise ValidationError(msg) + + +@dataclass(frozen=True) +class DistributedMapDestinationConfig: + """Destination routing for map run results.""" + + on_success: SuccessDestination | None = None + on_failure: FailureDestination | None = None + + +@dataclass(frozen=True) +class DistributedMapDestination: + """Destination factories for a map run.""" + + class S3: + """S3 destination factories.""" + + @staticmethod + def successes( + prefix_uri: str, + *, + include_input: bool = False, + include_output: bool = True, + expected_bucket_owner: str | None = None, + ) -> SuccessDestination: + """Route succeeded item records to an S3 prefix.""" + bucket, prefix = _parse_s3_uri(prefix_uri) + return SuccessDestination( + bucket=bucket, + prefix=prefix or "", + include_input=include_input, + include_output=include_output, + expected_bucket_owner=expected_bucket_owner, + ) + + @staticmethod + def failures( + prefix_uri: str, + *, + include_input: bool = True, + include_error: bool = True, + expected_bucket_owner: str | None = None, + ) -> FailureDestination: + """Route permanently-failed item records to an S3 prefix.""" + bucket, prefix = _parse_s3_uri(prefix_uri) + return FailureDestination( + bucket=bucket, + prefix=prefix or "", + include_input=include_input, + include_error=include_error, + expected_bucket_owner=expected_bucket_owner, + ) + + +@dataclass(frozen=True) +class DistributedMapConfig: + """Configuration for map run operations.""" + + destination: DistributedMapDestinationConfig | None = None + completion_config: DistributedMapCompletionConfig | None = None + timeout: Duration | None = None + collect_results: bool = False + item_serdes: SerDes | None = None # None = DEFAULT_JSON_SERDES + + def __post_init__(self) -> None: + if self.timeout is not None and not ( + 0 < self.timeout.to_seconds() <= 7776000 # noqa: PLR2004 + ): + msg = ( + "timeout must be positive and at most 90 days, got: " + f"{self.timeout.to_seconds()}s" + ) + raise ValidationError(msg) + if self.item_serdes is not None and not self.collect_results: + msg = "item_serdes requires collect_results=True" + raise ValidationError(msg) + + +# endregion map run configuration + + @dataclass(frozen=True) class ChildConfig(Generic[T]): """Configuration options for child context operations. 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 ff1ab6bc..0f2b9032 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 @@ -20,12 +20,17 @@ Duration, InvokeConfig, MapConfig, + DistributedMapConfig, + DistributedMapProcessor, + DistributedMapSource, ParallelBranch, ParallelConfig, StepConfig, WaitForCallbackConfig, ) from aws_durable_execution_sdk_python.concurrency.models import ( + DistributedMapResult, + DistributedMapSummary, envelope_summary_generator, ) from aws_durable_execution_sdk_python.exceptions import ( @@ -54,6 +59,9 @@ from aws_durable_execution_sdk_python.operation.child import child_handler from aws_durable_execution_sdk_python.operation.invoke import InvokeOperationExecutor from aws_durable_execution_sdk_python.operation.map import map_handler +from aws_durable_execution_sdk_python.operation.distributed_map import ( + DistributedMapOperationExecutor, +) from aws_durable_execution_sdk_python.operation.parallel import parallel_handler from aws_durable_execution_sdk_python.operation.step import StepOperationExecutor from aws_durable_execution_sdk_python.operation.wait import WaitOperationExecutor @@ -658,6 +666,51 @@ def invoke( ) return executor.process() + def distributed_map( + self, + source: DistributedMapSource | Sequence[Any], + processor: DistributedMapProcessor, + max_concurrency: int, + name: str | None = None, + config: DistributedMapConfig | None = None, + ) -> DistributedMapSummary | DistributedMapResult: + """Start a distributed map run and resolve with its summary. + + Args: + source: The items to process (a typed source or a plain-list shorthand) + processor: The processor configuration built via a DistributedMapProcessor factory + max_concurrency: Maximum concurrent processor invocations + name: Optional name for the operation + config: Optional run-level configuration + + Returns: + The map run's summary, or a DistributedMapResult when config.collect_results is set + """ + if not isinstance(source, (DistributedMapSource, list, tuple)): + msg = "source must be a DistributedMapSource or a list/tuple of items" + raise ValidationError(msg) + if max_concurrency <= 0: + msg = "max_concurrency must be greater than zero" + raise ValidationError(msg) + if not config: + config = DistributedMapConfig() + with self._replay_aware(): + operation_id = self._create_step_id() + executor: DistributedMapOperationExecutor = DistributedMapOperationExecutor( + source=source, + processor=processor, + max_concurrency=max_concurrency, + state=self.state, + operation_identifier=OperationIdentifier( + operation_id=operation_id, + sub_type=OperationSubType.DISTRIBUTED_MAP, + parent_id=self._parent_id, + name=name, + ), + config=config, + ) + return executor.process() + def map( self, inputs: Sequence[U], diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/distributed_map_helpers.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/distributed_map_helpers.py new file mode 100644 index 00000000..93f3daf1 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/distributed_map_helpers.py @@ -0,0 +1,241 @@ +"""Authoring helpers for distributed map processor and reader Lambda functions. + +Wrappers that own the request/response format so customers can write a plain +function to process items/batches or read pages. Imported explicitly from this +module, not the main package. +""" + +from __future__ import annotations + +import json +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Any, Callable, Literal + +from aws_durable_execution_sdk_python.concurrency.models import BatchItemStatus +from aws_durable_execution_sdk_python.config import CompletionConfig, MapConfig +from aws_durable_execution_sdk_python.exceptions import ValidationError +from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python.serdes import ( + DEFAULT_JSON_SERDES, + SerDes, + SerDesContext, +) + +_CTX = SerDesContext() +_READER_STATE_LIMIT = 32 * 1024 + + +@dataclass(frozen=True) +class ReaderPage: + """A page returned by a reader function: its items and the next state.""" + + items: list[Any] = field(default_factory=list) + next_state: Any | None = None + + +def _to_item(serdes: SerDes[Any], body: Any) -> Any: + """Recover a typed item from a record body (a JSON value).""" + return serdes.deserialize(json.dumps(body), _CTX) + + +def _to_json_value(serdes: SerDes[Any], value: Any) -> Any: + """Serialize a value and return it as a JSON value for the wire.""" + return json.loads(serdes.serialize(value, _CTX)) + + +def _error_entry(item_id: str, exc: BaseException) -> dict[str, Any]: + return { + "itemIdentifier": item_id, + "error": {"errorType": type(exc).__name__, "errorMessage": str(exc)}, + } + + +def _validate_report(report: str) -> None: + if report not in ("results", "failures"): + msg = f"report must be 'results' or 'failures', got: {report!r}" + raise ValidationError(msg) + + +def create_distributed_map_item_handler( + func: Callable[[Any], Any], + *, + item_serdes: SerDes[Any] | None = None, + concurrency: int | None = None, + report: Literal["results", "failures"] = "results", +) -> Callable[..., dict[str, Any]]: + """Wrap a Lambda to be used as a report_item_results processor. + + Pass ``report="failures"`` for a report_failed_items processor. ``func`` + takes one item and returns its output or raises. + """ + _validate_report(report) + serdes = item_serdes or DEFAULT_JSON_SERDES + + def handler(event: dict[str, Any], _context: Any = None) -> dict[str, Any]: + records: list[dict[str, Any]] = event["records"] + + def run(record: dict[str, Any]) -> Any: + return func(_to_item(serdes, record["body"])) + + outputs: list[Any] = [None] * len(records) + errors: list[BaseException | None] = [None] * len(records) + with ThreadPoolExecutor(max_workers=concurrency or len(records) or 1) as pool: + futures = {pool.submit(run, r): i for i, r in enumerate(records)} + for future, i in futures.items(): + try: + outputs[i] = future.result() + except Exception as exc: # noqa: BLE001 + errors[i] = exc + + results: list[dict[str, Any]] = [] + failures: list[dict[str, Any]] = [] + for i, record in enumerate(records): + item_id = record["itemId"] + err = errors[i] + if err is not None: + failures.append(_error_entry(item_id, err)) + elif report == "results": + results.append( + { + "itemIdentifier": item_id, + "output": _to_json_value(serdes, outputs[i]), + } + ) + + if report == "failures": + return {"batchItemFailures": failures} + return {"batchItemResults": results, "batchItemFailures": failures} + + return handler + + +def create_distributed_map_batch_handler( + func: Callable[[list[Any]], Any], + *, + item_serdes: SerDes[Any] | None = None, +) -> Callable[..., Any]: + """Wrap a Lambda to be used as a report_batch_outcome processor. + + ``func`` takes the whole batch of items. Returning succeeds every item; + raising fails every item. + """ + serdes = item_serdes or DEFAULT_JSON_SERDES + + def handler(event: dict[str, Any], _context: Any = None) -> Any: + items = [_to_item(serdes, record["body"]) for record in event["records"]] + return func(items) + + return handler + + +def create_distributed_map_reader( + func: Callable[[Any], ReaderPage], + *, + state_serdes: SerDes[Any] | None = None, +) -> Callable[..., dict[str, Any]]: + """Wrap a Lambda to be used as a reader source. + + ``func`` takes the current state and returns a ReaderPage. A ``next_state`` + of ``None`` signals the source is exhausted. + """ + serdes = state_serdes or DEFAULT_JSON_SERDES + + def handler(event: dict[str, Any], _context: Any = None) -> dict[str, Any]: + raw_state = event.get("state") + state = serdes.deserialize(raw_state, _CTX) if raw_state is not None else None + max_items: int = event["maxItems"] + + page = func(state) + if len(page.items) > max_items: + msg = f"reader returned {len(page.items)} items, exceeding maxItems {max_items}" + raise ValueError(msg) + + response: dict[str, Any] = {"items": list(page.items)} + if page.next_state is not None: + next_state = serdes.serialize(page.next_state, _CTX) + if len(next_state.encode("utf-8")) > _READER_STATE_LIMIT: + msg = f"reader next_state exceeds the {_READER_STATE_LIMIT // 1024} KB limit" + raise ValueError(msg) + response["nextState"] = next_state + return response + + return handler + + +def _item_response( + records: list[dict[str, Any]], + outputs: dict[int, Any], + failures: list[dict[str, Any]], + serdes: SerDes[Any], + report: str, +) -> dict[str, Any]: + """Build the item-handler response from per-index outputs and failures.""" + if report == "failures": + return {"batchItemFailures": failures} + results = [ + {"itemIdentifier": records[i]["itemId"], "output": _to_json_value(serdes, out)} + for i, out in outputs.items() + ] + return {"batchItemResults": results, "batchItemFailures": failures} + + +def create_distributed_map_item_handler_with_durable_execution( + func: Callable[..., Any], + *, + item_serdes: SerDes[Any] | None = None, + report: Literal["results", "failures"] = "results", +) -> Callable[..., dict[str, Any]]: + """Durable variant of the item handler; ``func`` receives (context, item).""" + _validate_report(report) + serdes = item_serdes or DEFAULT_JSON_SERDES + + @durable_execution + def handler(event: dict[str, Any], context: Any) -> dict[str, Any]: + records: list[dict[str, Any]] = event["records"] + + def per_item(ctx: Any, body: Any, _index: int, _inputs: Any) -> Any: + return func(ctx, _to_item(serdes, body)) + + batch = context.map( + [record["body"] for record in records], + per_item, + config=MapConfig(completion_config=CompletionConfig.all_completed()), + ) + + outputs: dict[int, Any] = {} + failures: list[dict[str, Any]] = [] + for bi in batch.all: + item_id = records[bi.index]["itemId"] + if bi.status is BatchItemStatus.SUCCEEDED: + outputs[bi.index] = bi.result + else: + err = bi.error + failures.append( + { + "itemIdentifier": item_id, + "error": { + "errorType": (err.type or "") if err else "", + "errorMessage": (err.message or "") if err else "", + }, + } + ) + return _item_response(records, outputs, failures, serdes, report) + + return handler + + +def create_distributed_map_batch_handler_with_durable_execution( + func: Callable[..., Any], + *, + item_serdes: SerDes[Any] | None = None, +) -> Callable[..., Any]: + """Durable variant of the batch handler; ``func`` receives (context, items).""" + serdes = item_serdes or DEFAULT_JSON_SERDES + + @durable_execution + def handler(event: dict[str, Any], context: Any) -> Any: + items = [_to_item(serdes, record["body"]) for record in event["records"]] + return func(context, items) + + return handler diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py index c9c1402d..fa807adc 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py @@ -327,6 +327,10 @@ class InvokeError(DurableOperationError): """Raised when a durable invoke operation fails.""" +class DistributedMapError(DurableOperationError): + """Raised when a durable map run operation fails.""" + + class ChildContextError(DurableOperationError): """Raised when a child context (run_in_child_context, map, parallel) fails.""" diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py index eb5ee78b..7be4fb41 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py @@ -3,13 +3,16 @@ import builtins import copy import datetime +import functools import logging +import re from collections.abc import MutableMapping from dataclasses import dataclass, field from enum import Enum from typing import TYPE_CHECKING, Any, NoReturn, Protocol, TypeAlias, cast import boto3 +import botocore.session from botocore.config import Config from aws_durable_execution_sdk_python.__about__ import __version__ @@ -38,6 +41,17 @@ logger = logging.getLogger(__name__) +@functools.lru_cache(maxsize=1) +def lambda_function_name_matcher() -> tuple[re.Pattern[str], int]: + """The allowed shape and length for a Lambda function name, read from botocore.""" + shape = ( + botocore.session.get_session() + .get_service_model("lambda") + .shape_for("NamespacedFunctionName") + ) + return re.compile(shape.metadata["pattern"]), shape.metadata["max"] + + def _is_in_var_dir(module_file: str = __file__) -> bool: """Return True if this SDK is installed under /var/lang/. @@ -74,6 +88,7 @@ class OperationType(Enum): WAIT = "WAIT" CALLBACK = "CALLBACK" CHAINED_INVOKE = "CHAINED_INVOKE" + DISTRIBUTED_MAP = "DISTRIBUTED_MAP" @classmethod def from_sub_type(cls, sub_type: OperationSubType) -> OperationType: @@ -86,6 +101,8 @@ def from_sub_type(cls, sub_type: OperationSubType) -> OperationType: return OperationType.CHAINED_INVOKE case OperationSubType.CALLBACK: return OperationType.CALLBACK + case OperationSubType.DISTRIBUTED_MAP: + return OperationType.DISTRIBUTED_MAP case ( OperationSubType.WAIT_FOR_CALLBACK | OperationSubType.RUN_IN_CHILD_CONTEXT @@ -116,6 +133,7 @@ class OperationSubType(Enum): WAIT_FOR_CALLBACK = "WaitForCallback" WAIT_FOR_CONDITION = "WaitForCondition" CHAINED_INVOKE = "ChainedInvoke" + DISTRIBUTED_MAP = "DistributedMap" class InvocationStatus(Enum): @@ -383,6 +401,69 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> ChainedInvokeDetails: ) +@dataclass(frozen=True) +class DistributedMapResultItemWire: + """Wire representation of a single map run item's outcome.""" + + item_id: str + status: str + output: Any | None = None + error: ErrorObject | None = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapResultItemWire: + error_raw = data.get("Error") + return cls( + item_id=data.get("ItemId", ""), + status=data.get("Status", ""), + output=data.get("Output"), + error=ErrorObject.from_dict(error_raw) if error_raw else None, + ) + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = { + "ItemId": self.item_id, + "Status": self.status, + } + if self.output is not None: + result["Output"] = self.output + if self.error is not None: + result["Error"] = self.error.to_dict() + return result + + +@dataclass(frozen=True) +class DistributedMapDetails: + status: str + completion_reason: str + distributed_map_run_arn: str | None = None + completion_details: str | None = None + total_count: int | None = None + success_count: int = 0 + failure_count: int = 0 + unprocessed_count: int = 0 + results: tuple[DistributedMapResultItemWire, ...] | None = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapDetails: + results_raw = data.get("Results") + return cls( + status=data["Status"], + completion_reason=data["CompletionReason"], + distributed_map_run_arn=data.get("DistributedMapRunArn"), + completion_details=data.get("CompletionDetails"), + total_count=data.get("TotalCount"), + success_count=data.get("SuccessCount", 0), + failure_count=data.get("FailureCount", 0), + unprocessed_count=data.get("UnprocessedCount", 0), + results=tuple( + DistributedMapResultItemWire.from_dict(item) for item in results_raw + ) + if results_raw is not None + else None, + ) + + @dataclass(frozen=True) class StepOptions: next_attempt_delay_seconds: int = 0 @@ -478,6 +559,205 @@ def to_dict(self) -> MutableMapping[str, Any]: return result +@dataclass(frozen=True) +class DistributedMapSourceWire: + """Wire representation of a map run source.""" + + source_type: str + max_items: int | None = None + inline_items: tuple[Any, ...] | None = None + s3_config: MutableMapping[str, Any] | None = None + reader_config: MutableMapping[str, Any] | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {"Type": self.source_type} + if self.source_type == "INLINE": + result["InlineSourceConfig"] = {"Items": list(self.inline_items or ())} + elif self.source_type == "S3" and self.s3_config is not None: + result["S3SourceConfig"] = dict(self.s3_config) + elif self.source_type == "READER_FUNCTION" and self.reader_config is not None: + result["ReaderFunctionSourceConfig"] = dict(self.reader_config) + if self.max_items is not None: + result["MaxItemsToRead"] = self.max_items + return result + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapSourceWire: + source_type = data.get("Type", "INLINE") + inline_cfg = data.get("InlineSourceConfig") or {} + return cls( + source_type=source_type, + max_items=data.get("MaxItemsToRead"), + inline_items=tuple(inline_cfg.get("Items", ())) + if source_type == "INLINE" + else None, + s3_config=data.get("S3SourceConfig"), + reader_config=data.get("ReaderFunctionSourceConfig"), + ) + + +@dataclass(frozen=True) +class DistributedMapProcessorWire: + """Wire representation of a map run processor.""" + + function_name: str + function_response_types: tuple[str, ...] | None = None + batch_size: int | None = None + max_retry_attempts: int | None = None + max_retry_duration_seconds: int | None = None + durable_execution_name_prefix: str | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {"FunctionName": self.function_name} + if self.function_response_types: + result["FunctionResponseTypes"] = list(self.function_response_types) + if self.batch_size is not None: + result["BatchSize"] = self.batch_size + if self.max_retry_attempts is not None: + result["MaxRetryAttempts"] = self.max_retry_attempts + if self.max_retry_duration_seconds is not None: + result["MaxRetryDurationSeconds"] = self.max_retry_duration_seconds + if self.durable_execution_name_prefix is not None: + result["DurableExecutionNamePrefix"] = self.durable_execution_name_prefix + return result + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapProcessorWire: + response_types = data.get("FunctionResponseTypes") + return cls( + function_name=data.get("FunctionName", ""), + function_response_types=tuple(response_types) if response_types else None, + batch_size=data.get("BatchSize"), + max_retry_attempts=data.get("MaxRetryAttempts"), + max_retry_duration_seconds=data.get("MaxRetryDurationSeconds"), + durable_execution_name_prefix=data.get("DurableExecutionNamePrefix"), + ) + + +@dataclass(frozen=True) +class DistributedMapCompletionConfigWire: + """Wire representation of a map run completion (failure-tolerance) config.""" + + tolerated_failure_count: int | None = None + tolerated_failure_percentage: float | None = None + minimum_sample_size: int | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {} + if self.tolerated_failure_count is not None: + result["ToleratedFailureCount"] = self.tolerated_failure_count + if self.tolerated_failure_percentage is not None: + result["ToleratedFailurePercentage"] = self.tolerated_failure_percentage + if self.minimum_sample_size is not None: + result["MinimumSampleSize"] = self.minimum_sample_size + return result + + @classmethod + def from_dict( + cls, data: MutableMapping[str, Any] + ) -> DistributedMapCompletionConfigWire: + return cls( + tolerated_failure_count=data.get("ToleratedFailureCount"), + tolerated_failure_percentage=data.get("ToleratedFailurePercentage"), + minimum_sample_size=data.get("MinimumSampleSize"), + ) + + +@dataclass(frozen=True) +class DistributedMapDestinationWire: + """Wire representation of a map run destination config.""" + + on_success: MutableMapping[str, Any] | None = None + on_failure: MutableMapping[str, Any] | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {} + if self.on_success is not None: + result["OnSuccess"] = dict(self.on_success) + if self.on_failure is not None: + result["OnFailure"] = dict(self.on_failure) + return result + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapDestinationWire: + return cls( + on_success=data.get("OnSuccess"), + on_failure=data.get("OnFailure"), + ) + + +@dataclass(frozen=True) +class DistributedMapResultCollectionWire: + """Wire representation of the map run result-collection setting.""" + + mode: str + + def to_dict(self) -> MutableMapping[str, Any]: + return {"Mode": self.mode} + + @classmethod + def from_dict( + cls, data: MutableMapping[str, Any] + ) -> DistributedMapResultCollectionWire: + return cls(mode=data.get("Mode", "NONE")) + + +@dataclass(frozen=True) +class DistributedMapOptions: + """Configuration options for starting a map run.""" + + max_concurrency: int + source: DistributedMapSourceWire + processor: DistributedMapProcessorWire + destination: DistributedMapDestinationWire | None = None + completion_config: DistributedMapCompletionConfigWire | None = None + result_collection: DistributedMapResultCollectionWire | None = None + timeout_seconds: int | None = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapOptions: + source_raw = data.get("Source") or {} + processor_raw = data.get("Processor") or {} + destination_raw = data.get("Destination") + completion_raw = data.get("CompletionConfig") + result_collection_raw = data.get("ResultCollection") + return cls( + max_concurrency=data["MaxConcurrency"], + source=DistributedMapSourceWire.from_dict(source_raw), + processor=DistributedMapProcessorWire.from_dict(processor_raw), + destination=DistributedMapDestinationWire.from_dict(destination_raw) + if destination_raw is not None + else None, + completion_config=DistributedMapCompletionConfigWire.from_dict( + completion_raw + ) + if completion_raw is not None + else None, + result_collection=DistributedMapResultCollectionWire.from_dict( + result_collection_raw + ) + if result_collection_raw is not None + else None, + timeout_seconds=data.get("TimeoutSeconds"), + ) + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = { + "MaxConcurrency": self.max_concurrency, + "Source": self.source.to_dict(), + "Processor": self.processor.to_dict(), + } + if self.destination is not None: + result["Destination"] = self.destination.to_dict() + if self.completion_config is not None: + result["CompletionConfig"] = self.completion_config.to_dict() + if self.result_collection is not None: + result["ResultCollection"] = self.result_collection.to_dict() + if self.timeout_seconds is not None: + result["TimeoutSeconds"] = self.timeout_seconds + return result + + @dataclass(frozen=True) class ContextOptions: replay_children: ReplayChildren = False @@ -510,6 +790,7 @@ class OperationUpdate: wait_options: WaitOptions | None = None callback_options: CallbackOptions | None = None chained_invoke_options: ChainedInvokeOptions | None = None + distributed_map_options: DistributedMapOptions | None = None def to_dict(self) -> MutableMapping[str, Any]: result: MutableMapping[str, Any] = { @@ -538,6 +819,8 @@ def to_dict(self) -> MutableMapping[str, Any]: result["CallbackOptions"] = self.callback_options.to_dict() if self.chained_invoke_options: result["ChainedInvokeOptions"] = self.chained_invoke_options.to_dict() + if self.distributed_map_options: + result["DistributedMapOptions"] = self.distributed_map_options.to_dict() return result @@ -566,6 +849,12 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> OperationUpdate: if invoke_data := data.get("ChainedInvokeOptions"): chained_invoke_options = ChainedInvokeOptions.from_dict(invoke_data) + distributed_map_options = None + if distributed_map_options_data := data.get("DistributedMapOptions"): + distributed_map_options = DistributedMapOptions.from_dict( + distributed_map_options_data + ) + return cls( operation_id=data["Id"], operation_type=OperationType(data["Type"]), @@ -580,6 +869,7 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> OperationUpdate: wait_options=wait_options, callback_options=callback_options, chained_invoke_options=chained_invoke_options, + distributed_map_options=distributed_map_options, ) @classmethod @@ -763,6 +1053,26 @@ def create_invoke_start( # endregion invoke + # region map run + @classmethod + def create_distributed_map_start( + cls, + identifier: OperationIdentifier, + distributed_map_options: DistributedMapOptions, + ) -> OperationUpdate: + """Create an instance of OperationUpdate for type: DISTRIBUTED_MAP, action: START.""" + return cls( + operation_id=identifier.operation_id, + parent_id=identifier.parent_id, + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + action=OperationAction.START, + name=identifier.name, + distributed_map_options=distributed_map_options, + ) + + # endregion map run + # region wait for condition @classmethod def create_wait_for_condition_start( @@ -886,6 +1196,7 @@ class Operation: wait_details: WaitDetails | None = None callback_details: CallbackDetails | None = None chained_invoke_details: ChainedInvokeDetails | None = None + distributed_map_details: DistributedMapDetails | None = None @classmethod def from_dict(cls, data: MutableMapping[str, Any]) -> Operation: @@ -930,6 +1241,12 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> Operation: chained_invoke_details ) + distributed_map_details = None + if distributed_map_details_input := data.get("DistributedMapDetails"): + distributed_map_details = DistributedMapDetails.from_dict( + distributed_map_details_input + ) + return cls( operation_id=data["Id"], operation_type=operation_type, @@ -945,6 +1262,7 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> Operation: wait_details=wait_details, callback_details=callback_details, chained_invoke_details=chained_invoke_details, + distributed_map_details=distributed_map_details, ) def to_dict(self) -> MutableMapping[str, Any]: @@ -1009,6 +1327,31 @@ def to_dict(self) -> MutableMapping[str, Any]: if self.chained_invoke_details.error: invoke_dict["Error"] = self.chained_invoke_details.error.to_dict() result["ChainedInvokeDetails"] = invoke_dict + if self.distributed_map_details: + distributed_map_details_dict: MutableMapping[str, Any] = { + "Status": self.distributed_map_details.status, + "CompletionReason": self.distributed_map_details.completion_reason, + "SuccessCount": self.distributed_map_details.success_count, + "FailureCount": self.distributed_map_details.failure_count, + "UnprocessedCount": self.distributed_map_details.unprocessed_count, + } + if self.distributed_map_details.distributed_map_run_arn: + distributed_map_details_dict["DistributedMapRunArn"] = ( + self.distributed_map_details.distributed_map_run_arn + ) + if self.distributed_map_details.completion_details: + distributed_map_details_dict["CompletionDetails"] = ( + self.distributed_map_details.completion_details + ) + if self.distributed_map_details.total_count is not None: + distributed_map_details_dict["TotalCount"] = ( + self.distributed_map_details.total_count + ) + if self.distributed_map_details.results is not None: + distributed_map_details_dict["Results"] = [ + item.to_dict() for item in self.distributed_map_details.results + ] + result["DistributedMapDetails"] = distributed_map_details_dict return result def to_json_dict(self) -> MutableMapping[str, Any]: diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/distributed_map.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/distributed_map.py new file mode 100644 index 00000000..a7737833 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/distributed_map.py @@ -0,0 +1,502 @@ +"""Implement the Durable map run operation.""" + +from __future__ import annotations + +import json +import logging +from typing import TYPE_CHECKING, Any + +from aws_durable_execution_sdk_python.concurrency.models import ( + DistributedMapCompletionReason, + DistributedMapItemError, + DistributedMapResult, + DistributedMapResultItem, + DistributedMapStatus, + DistributedMapSummary, +) +from aws_durable_execution_sdk_python.config import ( + DistributedMapProcessor, + DistributedMapSource, + ProcessorRetryConfig, +) +from aws_durable_execution_sdk_python.exceptions import ( + DistributedMapError, + ExecutionError, + ValidationError, +) +from aws_durable_execution_sdk_python.lambda_service import ( + DistributedMapCompletionConfigWire, + DistributedMapDestinationWire, + DistributedMapOptions, + DistributedMapProcessorWire, + DistributedMapResultCollectionWire, + DistributedMapSourceWire, + OperationUpdate, +) +from aws_durable_execution_sdk_python.operation.base import ( + CheckResult, + OperationExecutor, +) +from aws_durable_execution_sdk_python.serdes import ( + DEFAULT_JSON_SERDES, + deserialize, + serialize, +) +from aws_durable_execution_sdk_python.suspend import suspend_with_optional_resume_delay + +if TYPE_CHECKING: + from collections.abc import Sequence + + from aws_durable_execution_sdk_python.config import ( + DistributedMapConfig, + DistributedMapDestinationConfig, + S3SourceConfig, + ) + from aws_durable_execution_sdk_python.identifier import OperationIdentifier + from aws_durable_execution_sdk_python.lambda_service import DistributedMapDetails + from aws_durable_execution_sdk_python.state import ( + CheckpointedResult, + ExecutionState, + ) + +logger = logging.getLogger(__name__) + +# The backend uses -1 for unlimited retries, set when a customer passes ProcessorRetryConfig.UNLIMITED. +_UNLIMITED_RETRY_WIRE = -1 + +# Size limits for the inline item list (1 MB) and the reader's saved state (32 KB). +_INLINE_SIZE_LIMIT = 1024 * 1024 +_READER_STATE_LIMIT = 32 * 1024 + + +def _s3_config_to_wire(s3: S3SourceConfig) -> dict[str, Any]: + """Translate a resolved S3 source config into its wire dict.""" + result: dict[str, Any] = {"Bucket": s3.bucket} + if s3.key is not None: + result["Key"] = s3.key + if s3.prefix is not None: + result["KeyPrefix"] = s3.prefix + if s3.transform is not None: + result["Transform"] = s3.transform + if s3.expected_bucket_owner is not None: + result["ExpectedBucketOwner"] = s3.expected_bucket_owner + if s3.fmt is not None: + result["Format"] = s3.fmt + if s3.fmt == "CSV": + csv_options: dict[str, Any] = { + "HeaderLocation": "GIVEN" if s3.headers is not None else "FIRST_ROW", + } + if s3.headers is not None: + csv_options["Headers"] = list(s3.headers) + if s3.delimiter is not None: + csv_options["Delimiter"] = s3.delimiter + result["CsvFormatOptions"] = csv_options + return result + + +def _inline_items_to_wire( + items: tuple[Any, ...], + serdes: Any, + operation_id: str, + durable_execution_arn: str, +) -> tuple[Any, ...]: + """Serialize each inline item to its wire JSON value, enforcing the 1 MB cap.""" + wire_items: list[Any] = [] + for item in items: + serialized = serialize( + serdes=serdes, + value=item, + operation_id=operation_id, + durable_execution_arn=durable_execution_arn, + ) + try: + wire_items.append(json.loads(serialized)) + except json.JSONDecodeError as e: + msg = "inline source serdes must produce a JSON value for each item" + raise ValidationError(msg) from e + total = len(json.dumps(wire_items, separators=(",", ":")).encode("utf-8")) + if total > _INLINE_SIZE_LIMIT: + msg = ( + f"inline source exceeds the {_INLINE_SIZE_LIMIT // 1024 // 1024} MB limit " + f"(serialized size: {total} bytes)" + ) + raise ValidationError(msg) + return tuple(wire_items) + + +def _source_to_wire( + source: DistributedMapSource | Sequence[Any], + operation_id: str, + durable_execution_arn: str, +) -> DistributedMapSourceWire: + """Translate a source (typed or plain-list shorthand) into its wire form.""" + if not isinstance(source, DistributedMapSource): + # A plain list is treated as an inline source with the default serializer. + wire_items = _inline_items_to_wire( + tuple(source), DEFAULT_JSON_SERDES, operation_id, durable_execution_arn + ) + return DistributedMapSourceWire(source_type="INLINE", inline_items=wire_items) + + if source.source_type == "INLINE": + wire_items = _inline_items_to_wire( + source.inline_items or (), + source.inline_serdes or DEFAULT_JSON_SERDES, + operation_id, + durable_execution_arn, + ) + return DistributedMapSourceWire( + source_type="INLINE", + inline_items=wire_items, + max_items=source.max_items, + ) + if source.source_type == "S3" and source.s3 is not None: + return DistributedMapSourceWire( + source_type="S3", + max_items=source.max_items, + s3_config=_s3_config_to_wire(source.s3), + ) + if source.source_type == "READER_FUNCTION" and source.reader is not None: + reader_config: dict[str, Any] = {"FunctionName": source.reader.function_name} + if source.reader.initial_state is not None: + state = serialize( + serdes=source.reader.state_serdes or DEFAULT_JSON_SERDES, + value=source.reader.initial_state, + operation_id=operation_id, + durable_execution_arn=durable_execution_arn, + ) + if len(state.encode("utf-8")) > _READER_STATE_LIMIT: + msg = ( + f"reader initial_state exceeds the " + f"{_READER_STATE_LIMIT // 1024} KB limit" + ) + raise ValidationError(msg) + reader_config["InitialState"] = state + return DistributedMapSourceWire( + source_type="READER_FUNCTION", + max_items=source.max_items, + reader_config=reader_config, + ) + msg = f"Unsupported map run source type: {source.source_type}" + raise ExecutionError(msg) + + +def _processor_to_wire( + processor: DistributedMapProcessor, +) -> DistributedMapProcessorWire: + """Translate a processor config into its wire form, mapping unlimited to -1.""" + response_types = ( + (processor.response_mode,) if processor.response_mode is not None else None + ) + + max_retry_attempts: int | None = None + max_retry_duration_seconds: int | None = None + if processor.retry is not None: + attempts = processor.retry.max_retry_attempts + if attempts == ProcessorRetryConfig.UNLIMITED: + max_retry_attempts = _UNLIMITED_RETRY_WIRE + elif isinstance(attempts, int): + max_retry_attempts = attempts + if processor.retry.max_retry_duration is not None: + max_retry_duration_seconds = processor.retry.max_retry_duration.to_seconds() + + return DistributedMapProcessorWire( + function_name=processor.function_name, + function_response_types=response_types, + batch_size=processor.batch_size, + max_retry_attempts=max_retry_attempts, + max_retry_duration_seconds=max_retry_duration_seconds, + durable_execution_name_prefix=processor.durable_execution_name_prefix, + ) + + +def _completion_to_wire( + config: DistributedMapConfig, +) -> DistributedMapCompletionConfigWire | None: + completion = config.completion_config + if completion is None or ( + completion.tolerated_failure_count is None + and completion.tolerated_failure_percentage is None + and completion.minimum_sample_size is None + ): + return None + return DistributedMapCompletionConfigWire( + tolerated_failure_count=completion.tolerated_failure_count, + tolerated_failure_percentage=completion.tolerated_failure_percentage, + minimum_sample_size=completion.minimum_sample_size, + ) + + +def _destination_to_wire( + destination: DistributedMapDestinationConfig | None, +) -> DistributedMapDestinationWire | None: + if destination is None: + return None + on_success: dict[str, Any] | None = None + on_failure: dict[str, Any] | None = None + if destination.on_success is not None: + s = destination.on_success + include: list[str] = [] + if s.include_input: + include.append("INPUT") + if s.include_output: + include.append("OUTPUT") + s3_config: dict[str, Any] = {"Bucket": s.bucket, "KeyPrefix": s.prefix} + if s.expected_bucket_owner is not None: + s3_config["ExpectedBucketOwner"] = s.expected_bucket_owner + on_success = { + "Type": "S3", + "Include": include, + "S3DestinationConfig": s3_config, + } + if destination.on_failure is not None: + f = destination.on_failure + f_include: list[str] = [] + if f.include_input: + f_include.append("INPUT") + if f.include_error: + f_include.append("ERROR") + f_s3_config: dict[str, Any] = {"Bucket": f.bucket, "KeyPrefix": f.prefix} + if f.expected_bucket_owner is not None: + f_s3_config["ExpectedBucketOwner"] = f.expected_bucket_owner + on_failure = { + "Type": "S3", + "Include": f_include, + "S3DestinationConfig": f_s3_config, + } + if on_success is None and on_failure is None: + return None + return DistributedMapDestinationWire(on_success=on_success, on_failure=on_failure) + + +def _build_distributed_map_options( + source: DistributedMapSource | Sequence[Any], + processor: DistributedMapProcessor, + max_concurrency: int, + config: DistributedMapConfig, + operation_id: str, + durable_execution_arn: str, +) -> DistributedMapOptions: + """Assemble the wire options payload from the operands and config.""" + result_collection = ( + DistributedMapResultCollectionWire(mode="INLINE") + if config.collect_results + else None + ) + return DistributedMapOptions( + max_concurrency=max_concurrency, + source=_source_to_wire(source, operation_id, durable_execution_arn), + processor=_processor_to_wire(processor), + destination=_destination_to_wire(config.destination), + completion_config=_completion_to_wire(config), + result_collection=result_collection, + timeout_seconds=config.timeout.to_seconds() + if config.timeout is not None + else None, + ) + + +def _summary_fields(details: DistributedMapDetails) -> dict[str, Any]: + """Shared summary fields extracted from the terminal details block.""" + try: + status = DistributedMapStatus(details.status) + completion_reason = DistributedMapCompletionReason(details.completion_reason) + except ValueError as e: + msg = ( + f"Unknown map run status or completion reason from the backend " + f"({details.status!r}, {details.completion_reason!r})" + ) + raise ExecutionError(msg) from e + return { + "status": status, + "completion_reason": completion_reason, + "success_count": details.success_count, + "failure_count": details.failure_count, + "unprocessed_count": details.unprocessed_count, + "distributed_map_run_arn": details.distributed_map_run_arn, + "completion_details": details.completion_details, + "total_count": details.total_count, + } + + +class DistributedMapOperationExecutor(OperationExecutor[DistributedMapSummary]): + """Executor for map run operations. + + Creates the START checkpoint if none exists, then suspends until the + backend completes the run and re-invokes the parent. On resume, a + ``DistributedMapSummary`` (or ``DistributedMapResult`` when result collection is enabled) + is built from the checkpointed ``DistributedMapDetails``. + """ + + def __init__( + self, + source: DistributedMapSource | Sequence[Any], + processor: DistributedMapProcessor, + max_concurrency: int, + state: ExecutionState, + operation_identifier: OperationIdentifier, + config: DistributedMapConfig, + ): + """Initialize the map run operation executor. + + Args: + source: The items to process (typed source or plain-list shorthand) + processor: The processor configuration + max_concurrency: Maximum concurrent processor invocations + state: The execution state + operation_identifier: The operation identifier + config: Configuration for the map run operation + """ + self.source = source + self.processor = processor + self.max_concurrency = max_concurrency + self.state = state + self.operation_identifier = operation_identifier + self.config = config + + def _resolve_summary( + self, details: DistributedMapDetails | None + ) -> DistributedMapSummary: + """Reconstruct the resolved summary/result from the terminal details.""" + if details is None: + msg = "DISTRIBUTED_MAP operation succeeded but carried no DistributedMapDetails" + raise ExecutionError(msg) + fields = _summary_fields(details) + if not self.config.collect_results: + return DistributedMapSummary(**fields) + return DistributedMapResult(**fields, all=self._deserialize_items(details)) + + def _deserialize_items( + self, details: DistributedMapDetails + ) -> list[DistributedMapResultItem]: + """Deserialize the per-item wire results into customer result items.""" + items: list[DistributedMapResultItem] = [] + for wire in details.results or (): + output: Any | None = None + if wire.output is not None: + # Output is already a JSON value; re-dump to text for the serdes. + output = deserialize( + serdes=self.config.item_serdes or DEFAULT_JSON_SERDES, + data=json.dumps(wire.output), + operation_id=self.operation_identifier.operation_id, + durable_execution_arn=self.state.durable_execution_arn, + ) + error = ( + DistributedMapItemError( + error_type=wire.error.type or "", + error_message=wire.error.message or "", + ) + if wire.error is not None + else None + ) + items.append( + DistributedMapResultItem( + item_id=wire.item_id, + status=wire.status, + output=output, + error=error, + ) + ) + return items + + def check_result_status(self) -> CheckResult[DistributedMapSummary]: + """Check operation status and create the START checkpoint if needed. + + Called twice by process() when creating synchronous checkpoints: once before + and once after, to detect if the operation completed immediately. + + Returns: + CheckResult indicating the next action to take + + Raises: + SuspendExecution: For STARTED operations waiting for completion + """ + checkpointed_result: CheckpointedResult = self.state.get_checkpoint_result( + self.operation_identifier.operation_id + ) + + # Terminal success - build the summary/result from DistributedMapDetails + if checkpointed_result.is_succeeded(): + operation = checkpointed_result.operation + summary = self._resolve_summary( + operation.distributed_map_details if operation else None + ) + return CheckResult.create_completed(summary) + + # Operation-level terminal failure + if ( + checkpointed_result.is_failed() + or checkpointed_result.is_timed_out() + or checkpointed_result.is_stopped() + ): + msg = ( + f"Distributed map operation " + f"'{self.operation_identifier.name or self.operation_identifier.operation_id}' " + f"ended with status " + f"{checkpointed_result.status.value if checkpointed_result.status else 'UNKNOWN'}" + ) + checkpointed_result.raise_operation_error(DistributedMapError, msg=msg) + + # Started - ready to suspend + if checkpointed_result.is_started(): + logger.debug( + "⏳ Map run %s still in progress, will suspend", + self.operation_identifier.name + or self.operation_identifier.operation_id, + ) + return CheckResult.create_is_ready_to_execute(checkpointed_result) + + # Create START checkpoint if not exists + if not checkpointed_result.is_existent(): + start_operation: OperationUpdate = ( + OperationUpdate.create_distributed_map_start( + identifier=self.operation_identifier, + distributed_map_options=_build_distributed_map_options( + source=self.source, + processor=self.processor, + max_concurrency=self.max_concurrency, + config=self.config, + operation_id=self.operation_identifier.operation_id, + durable_execution_arn=self.state.durable_execution_arn, + ), + ) + ) + # Checkpoint map run START with blocking (is_sync=True). + # Must ensure the map run is recorded before suspending execution. + self.state.create_checkpoint(operation_update=start_operation, is_sync=True) + + logger.debug( + "🚀 Map run %s started, will check for immediate completion", + self.operation_identifier.name + or self.operation_identifier.operation_id, + ) + + # Signal to process() that checkpoint was created - to recheck status + # for immediate completion before proceeding. + return CheckResult.create_started() + + # Ready to suspend (checkpoint exists but not in a terminal or started state) + return CheckResult.create_is_ready_to_execute(checkpointed_result) + + def execute( + self, _checkpointed_result: CheckpointedResult + ) -> DistributedMapSummary: + """Execute map run operation by suspending to wait for async completion. + + The map run operation doesn't execute synchronously - it suspends and + the backend runs the map run asynchronously. + + Args: + checkpointed_result: The checkpoint data (unused, but required by interface) + + Returns: + Never returns - always suspends + + Raises: + Always suspends via suspend_with_optional_resume_delay + ExecutionError: If suspend doesn't raise (should never happen) + """ + msg: str = f"Map run {self.operation_identifier.operation_id} started, suspending for completion" + suspend_with_optional_resume_delay(msg) + # This line should never be reached since suspend_with_optional_resume_delay always raises + error_msg: str = "suspend_with_optional_resume_delay should have raised an exception, but did not." + raise ExecutionError(error_msg) from None diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py index daab80cc..cc6a9116 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py @@ -49,6 +49,7 @@ class OperationType(Enum): WAIT = "WAIT" CALLBACK = "CALLBACK" CHAINED_INVOKE = "CHAINED_INVOKE" + DISTRIBUTED_MAP = "DISTRIBUTED_MAP" def _to_invocation_status(status: ServiceInvocationStatus) -> InvocationStatus: 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 26aefbe3..6b6013a8 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 @@ -134,6 +134,9 @@ def create_from_operation(cls, operation: Operation) -> CheckpointedResult: result = context_details.result if context_details else None error = context_details.error if context_details else None + # OperationType.DISTRIBUTED_MAP has no operation-level result/error. + # The operation itself carries the map run's outcome. + return cls( operation=operation, status=operation.status, result=result, error=error ) 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 c5b8ceeb..4a397030 100644 --- a/packages/aws-durable-execution-sdk-python/tests/context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/context_test.py @@ -15,6 +15,7 @@ Duration, InvokeConfig, MapConfig, + DistributedMapProcessor, ParallelBranch, ParallelConfig, StepConfig, @@ -1021,6 +1022,41 @@ def test_wait_with_time_less_than_one(mock_executor_class): # endregion wait +# region distributed_map +@pytest.mark.parametrize("max_concurrency", [0, -1]) +def test_map_run_rejects_non_positive_max_concurrency(max_concurrency: int): + """Test distributed_map raises ValidationError when max_concurrency is not positive.""" + context = create_test_context() + + with pytest.raises( + ValidationError, match="max_concurrency must be greater than zero" + ): + context.distributed_map( + ["a"], + DistributedMapProcessor.report_batch_outcome("test_processor"), + max_concurrency=max_concurrency, + ) + + +@pytest.mark.parametrize( + "source", + ["s3://bucket", b"bytes", bytearray(b"ba"), {"a": 1}, {"a", "b"}], +) +def test_distributed_map_rejects_non_list_source(source): + """Test distributed_map rejects sources that are not a DistributedMapSource or list/tuple.""" + context = create_test_context() + + with pytest.raises(ValidationError, match="list/tuple"): + context.distributed_map( + source, + DistributedMapProcessor.report_batch_outcome("test_processor"), + max_concurrency=1, + ) + + +# endregion distributed_map + + # region run_in_child_context @patch("aws_durable_execution_sdk_python.context.child_handler") def test_run_in_child_context_basic(mock_handler): diff --git a/packages/aws-durable-execution-sdk-python/tests/distributed_map_helpers_test.py b/packages/aws-durable-execution-sdk-python/tests/distributed_map_helpers_test.py new file mode 100644 index 00000000..adc0ba06 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/distributed_map_helpers_test.py @@ -0,0 +1,114 @@ +"""Unit tests for the non-durable distributed map authoring wrappers.""" + +from __future__ import annotations + +import pytest + +from aws_durable_execution_sdk_python.distributed_map_helpers import ( + ReaderPage, + create_distributed_map_batch_handler, + create_distributed_map_item_handler, + create_distributed_map_item_handler_with_durable_execution, + create_distributed_map_reader, +) +from aws_durable_execution_sdk_python.exceptions import ValidationError + + +def test_item_handler_invalid_report_rejected(): + with pytest.raises(ValidationError, match="report must be"): + create_distributed_map_item_handler(lambda x: x, report="bogus") + + +def test_durable_item_handler_invalid_report_rejected(): + with pytest.raises(ValidationError, match="report must be"): + create_distributed_map_item_handler_with_durable_execution( + lambda _ctx, item: item, report="bogus" + ) + + +def test_item_handler_reports_results_in_order(): + handler = create_distributed_map_item_handler(lambda x: x * 2) + resp = handler( + {"records": [{"itemId": "0", "body": 2}, {"itemId": "1", "body": 3}]} + ) + assert resp["batchItemResults"] == [ + {"itemIdentifier": "0", "output": 4}, + {"itemIdentifier": "1", "output": 6}, + ] + assert resp["batchItemFailures"] == [] + + +def test_item_handler_captures_failures(): + def process(x): + if x == "bad": + msg = "boom" + raise ValueError(msg) + return x + + handler = create_distributed_map_item_handler(process) + resp = handler( + {"records": [{"itemId": "0", "body": "ok"}, {"itemId": "1", "body": "bad"}]} + ) + assert resp["batchItemResults"] == [{"itemIdentifier": "0", "output": "ok"}] + assert resp["batchItemFailures"] == [ + { + "itemIdentifier": "1", + "error": {"errorType": "ValueError", "errorMessage": "boom"}, + } + ] + + +def test_item_handler_failures_form_reports_only_failures(): + handler = create_distributed_map_item_handler(lambda x: x, report="failures") + resp = handler({"records": [{"itemId": "0", "body": 1}]}) + assert resp == {"batchItemFailures": []} + + +def test_batch_handler_success_and_propagates_error(): + seen: list = [] + handler = create_distributed_map_batch_handler( + lambda items: seen.extend(items) or "done" + ) + assert ( + handler({"records": [{"itemId": "0", "body": 1}, {"itemId": "1", "body": 2}]}) + == "done" + ) + assert seen == [1, 2] + + def boom(_items): + msg = "batch failed" + raise RuntimeError(msg) + + failing = create_distributed_map_batch_handler(boom) + with pytest.raises(RuntimeError, match="batch failed"): + failing({"records": [{"itemId": "0", "body": 1}]}) + + +def test_reader_returns_items_and_next_state_then_exhausts(): + def read(state): + if state is None: + return ReaderPage(items=[1, 2], next_state={"page": 1}) + return ReaderPage(items=[3]) + + handler = create_distributed_map_reader(read) + first = handler({"state": None, "maxItems": 10}) + assert first["items"] == [1, 2] + assert first["nextState"] == '{"page": 1}' + + second = handler({"state": '{"page": 1}', "maxItems": 10}) + assert second["items"] == [3] + assert "nextState" not in second + + +def test_reader_rejects_page_over_max_items(): + handler = create_distributed_map_reader(lambda _s: ReaderPage(items=[1, 2, 3])) + with pytest.raises(ValueError, match="exceeding maxItems"): + handler({"state": None, "maxItems": 2}) + + +def test_reader_rejects_oversized_next_state(): + handler = create_distributed_map_reader( + lambda _s: ReaderPage(items=[1], next_state="x" * 40_000) + ) + with pytest.raises(ValueError, match="32 KB limit"): + handler({"state": None, "maxItems": 10}) diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/distributed_map_helpers_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/distributed_map_helpers_int_test.py new file mode 100644 index 00000000..95a0ff38 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/distributed_map_helpers_int_test.py @@ -0,0 +1,148 @@ +"""Integration tests for the durable distributed map authoring wrappers. + +The durable variants are ``@durable_execution`` handlers; they run through the +full invocation harness with a mocked checkpoint backend. +""" + +from __future__ import annotations + +import json +from unittest.mock import Mock, patch + +from aws_durable_execution_sdk_python.distributed_map_helpers import ( + create_distributed_map_batch_handler_with_durable_execution, + create_distributed_map_item_handler_with_durable_execution, +) +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, # noqa: F401 (ensures decorator import path is valid) +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + Operation, + OperationStatus, + OperationType, +) + +_ARN = "arn:aws:lambda:us-east-1:123456789012:function:proc:1/durable-execution/execution-1" + + +def _lambda_context(): + ctx = Mock() + ctx.aws_request_id = "test-request-id" + ctx.client_context = None + ctx.identity = None + ctx._epoch_deadline_time_in_ms = 0 # noqa: SLF001 + ctx.invoked_function_arn = "test-arn" + ctx.tenant_id = None + return ctx + + +def _event(records: list[dict]): + return { + "DurableExecutionArn": _ARN, + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [ + { + "Id": "execution-1", + "Type": "EXECUTION", + "Status": "STARTED", + "ExecutionDetails": { + "InputPayload": json.dumps({"records": records}) + }, + } + ], + "NextMarker": "", + }, + "LocalRunner": True, + } + + +def _run(handler, records: list[dict]): + operations = [ + Operation( + operation_id="execution-1", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + ) + ] + + def mock_checkpoint( + durable_execution_arn, checkpoint_token, updates, client_token="token" + ): # noqa: S107 + for update in updates: + operations.append( + Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + ) + ) + return CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState( + operations=operations.copy() + ), + ) + + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client_class.initialize_client.return_value = mock_client + mock_client.checkpoint = mock_checkpoint + return handler(_event(records), _lambda_context()) + + +def test_durable_item_handler_reports_results(): + handler = create_distributed_map_item_handler_with_durable_execution( + lambda _ctx, item: item * 2 + ) + result = _run(handler, [{"itemId": "0", "body": 2}, {"itemId": "1", "body": 3}]) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + data = json.loads(result["Result"]) + assert data["batchItemResults"] == [ + {"itemIdentifier": "0", "output": 4}, + {"itemIdentifier": "1", "output": 6}, + ] + assert data["batchItemFailures"] == [] + + +def test_durable_item_handler_captures_failure(): + def process(_ctx, item): + if item == "bad": + msg = "boom" + raise ValueError(msg) + return item + + handler = create_distributed_map_item_handler_with_durable_execution(process) + result = _run( + handler, [{"itemId": "0", "body": "ok"}, {"itemId": "1", "body": "bad"}] + ) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + data = json.loads(result["Result"]) + assert data["batchItemResults"] == [{"itemIdentifier": "0", "output": "ok"}] + assert data["batchItemFailures"][0]["itemIdentifier"] == "1" + assert data["batchItemFailures"][0]["error"]["errorType"] == "ValueError" + + +def test_durable_batch_handler_returns_value(): + handler = create_distributed_map_batch_handler_with_durable_execution( + lambda _ctx, items: {"count": len(items)} + ) + result = _run(handler, [{"itemId": "0", "body": 1}, {"itemId": "1", "body": 2}]) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + assert json.loads(result["Result"]) == {"count": 2} + + +def test_durable_item_handler_failures_form(): + handler = create_distributed_map_item_handler_with_durable_execution( + lambda _ctx, item: item, report="failures" + ) + result = _run(handler, [{"itemId": "0", "body": 1}]) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + data = json.loads(result["Result"]) + assert data == {"batchItemFailures": []} diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/distributed_map_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/distributed_map_int_test.py new file mode 100644 index 00000000..6e7c572a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/distributed_map_int_test.py @@ -0,0 +1,246 @@ +"""Integration tests for ctx.distributed_map through a full durable_execution invocation. + +Drives the real suspend/resume flow: the first invocation starts the DISTRIBUTED_MAP +operation and suspends (PENDING); a replay invocation with the operation +completed (carrying DistributedMapDetails) resolves to the summary/result. The backend +is mocked; no emulator or real service is involved. +""" + +from __future__ import annotations + +import json +from typing import Any +from unittest.mock import Mock, patch + +from aws_durable_execution_sdk_python.concurrency.models import DistributedMapResult +from aws_durable_execution_sdk_python.config import ( + DistributedMapConfig, + DistributedMapProcessor, +) +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + Operation, + OperationStatus, + OperationType, +) +from tests.test_helpers import operation_id_sequence + + +def _lambda_context(): + ctx = Mock() + ctx.aws_request_id = "test-request-id" + ctx.client_context = None + ctx.identity = None + ctx._epoch_deadline_time_in_ms = 0 # noqa: SLF001 + ctx.invoked_function_arn = "test-arn" + ctx.tenant_id = None + return ctx + + +_ARN = ( + "arn:aws:lambda:us-east-1:123456789012:function:test-func:1" + "/durable-execution/exec-001/inv-001" +) + + +def _initial_event(): + return { + "DurableExecutionArn": _ARN, + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [ + { + "Id": "execution-1", + "Type": "EXECUTION", + "Status": "STARTED", + "ExecutionDetails": {"InputPayload": "{}"}, + } + ], + "NextMarker": "", + }, + "LocalRunner": True, + } + + +def _replay_event(distributed_map_details: dict): + distributed_map_id = next(operation_id_sequence()) + return distributed_map_id, { + "DurableExecutionArn": _ARN, + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [ + { + "Id": "execution-1", + "Type": "EXECUTION", + "Status": "STARTED", + "ExecutionDetails": {"InputPayload": "{}"}, + }, + { + "Id": distributed_map_id, + "Type": "DISTRIBUTED_MAP", + "Status": "SUCCEEDED", + "ParentId": "execution-1", + "DistributedMapDetails": distributed_map_details, + }, + ], + "NextMarker": "", + }, + "LocalRunner": True, + } + + +def _tracking_checkpoint(): + """Checkpoint mock that records created operations as STARTED.""" + calls: list = [] + operations = [ + Operation( + operation_id="execution-1", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + ) + ] + + def mock_checkpoint( + durable_execution_arn, checkpoint_token, updates, client_token="token" + ): # noqa: S107 + calls.append(updates) + for update in updates: + operations.append( + Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + ) + ) + return CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState( + operations=operations.copy() + ), + ) + + return calls, mock_checkpoint + + +def _run(handler, event): + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client_class.initialize_client.return_value = mock_client + _calls, mock_checkpoint = _tracking_checkpoint() + mock_client.checkpoint = mock_checkpoint + return handler(event, _lambda_context()) + + +def test_map_run_suspends_then_resumes_with_summary(): + @durable_execution + def handler(event, context: DurableContext) -> dict[str, Any]: + summary = context.distributed_map( + ["a", "b"], + DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=2, + ) + return { + "status": summary.status.value, + "success": summary.success_count, + "failure": summary.failure_count, + "distributed_map_id": summary.distributed_map_id, + } + + # First invocation suspends. + first = _run(handler, _initial_event()) + assert first["Status"] == InvocationStatus.PENDING.value + + # Replay with the run completed. + _map_run_id, replay_event = _replay_event( + { + "Status": "SUCCEEDED", + "CompletionReason": "ALL_COMPLETED", + "SuccessCount": 2, + "FailureCount": 0, + "UnprocessedCount": 0, + "TotalCount": 2, + "DistributedMapRunArn": "arn:aws:lambda:us-east-1:123456789012:map-run:abc", + } + ) + replay = _run(handler, replay_event) + assert replay["Status"] == InvocationStatus.SUCCEEDED.value + data = json.loads(replay["Result"]) + assert data == { + "status": "SUCCEEDED", + "success": 2, + "failure": 0, + "distributed_map_id": "abc", + } + + +def test_map_run_collect_results_returns_items(): + @durable_execution + def handler(event, context: DurableContext) -> dict[str, Any]: + result = context.distributed_map( + ["a", "b"], + DistributedMapProcessor.report_item_results("proc"), + max_concurrency=2, + config=DistributedMapConfig(collect_results=True), + ) + assert isinstance(result, DistributedMapResult) + return { + "results": result.get_results(), + "errors": [e.error_message for e in result.get_errors()], + } + + _map_run_id, replay_event = _replay_event( + { + "Status": "SUCCEEDED", + "CompletionReason": "ALL_COMPLETED", + "SuccessCount": 1, + "FailureCount": 1, + "UnprocessedCount": 0, + "TotalCount": 2, + "Results": [ + {"ItemId": "0", "Status": "SUCCEEDED", "Output": 5}, + { + "ItemId": "1", + "Status": "FAILED", + "Error": {"ErrorType": "E", "ErrorMessage": "boom"}, + }, + ], + } + ) + replay = _run(handler, replay_event) + assert replay["Status"] == InvocationStatus.SUCCEEDED.value + data = json.loads(replay["Result"]) + assert data == {"results": [5], "errors": ["boom"]} + + +def test_map_run_throw_if_error_fails_execution(): + @durable_execution + def handler(event, context: DurableContext) -> dict[str, Any]: + summary = context.distributed_map( + ["a"], + DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=1, + ) + summary.throw_if_error() + return {"status": summary.status.value} + + _map_run_id, replay_event = _replay_event( + { + "Status": "FAILED", + "CompletionReason": "FAILURE_TOLERANCE_EXCEEDED", + "SuccessCount": 0, + "FailureCount": 1, + "UnprocessedCount": 0, + "TotalCount": 1, + } + ) + replay = _run(handler, replay_event) + assert replay["Status"] == InvocationStatus.FAILED.value diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/distributed_map_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/distributed_map_test.py new file mode 100644 index 00000000..d11a46bf --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/distributed_map_test.py @@ -0,0 +1,1699 @@ +"""Unit tests for map run handler.""" + +from __future__ import annotations + +import json +from unittest.mock import Mock, patch + +import pytest + +from aws_durable_execution_sdk_python.concurrency.models import ( + DistributedMapCompletionReason, + DistributedMapItemError, + DistributedMapResult, + DistributedMapResultItem, + DistributedMapStatus, + DistributedMapSummary, +) +from aws_durable_execution_sdk_python.config import ( + Duration, + DistributedMapCompletionConfig, + DistributedMapConfig, + DistributedMapDestination, + DistributedMapDestinationConfig, + DistributedMapProcessor, + DistributedMapSource, + ProcessorRetryConfig, +) +from aws_durable_execution_sdk_python.exceptions import ( + ExecutionError, + DistributedMapError, + SuspendExecution, + ValidationError, +) +from aws_durable_execution_sdk_python.identifier import OperationIdentifier +from aws_durable_execution_sdk_python.lambda_service import ( + ErrorObject, + DistributedMapDetails, + DistributedMapOptions, + DistributedMapResultItemWire, + Operation, + OperationAction, + OperationStatus, + OperationSubType, + OperationType, +) +from aws_durable_execution_sdk_python.operation.distributed_map import ( + DistributedMapOperationExecutor, +) +from aws_durable_execution_sdk_python.serdes import DEFAULT_JSON_SERDES +from aws_durable_execution_sdk_python.state import CheckpointedResult, ExecutionState + + +# Test helper - wraps DistributedMapOperationExecutor with a simple handler signature. +def distributed_map_handler( + source, processor, max_concurrency, state, operation_identifier, config=None +): + """Test helper that wraps DistributedMapOperationExecutor and runs it. + + ``processor`` may be a function-name string (wrapped as a batch-outcome + processor) or an already-built DistributedMapProcessor. + """ + if not config: + config = DistributedMapConfig() + if isinstance(processor, str): + processor = DistributedMapProcessor.report_batch_outcome(processor) + executor = DistributedMapOperationExecutor( + source=source, + processor=processor, + max_concurrency=max_concurrency, + state=state, + operation_identifier=operation_identifier, + config=config, + ) + return executor.process() + + +def _identifier( + operation_id: str, name: str | None = "test_map_run" +) -> OperationIdentifier: + return OperationIdentifier( + operation_id, OperationSubType.DISTRIBUTED_MAP, None, name + ) + + +def test_map_run_handler_already_succeeded(): + """Test distributed_map_handler returns a summary when the operation already succeeded.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + operation = Operation( + operation_id="mr1", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + status="SUCCEEDED", + completion_reason="ALL_COMPLETED", + success_count=5, + failure_count=0, + unprocessed_count=0, + total_count=5, + distributed_map_run_arn="arn:aws:lambda:us-east-1:123456789012:map-run:abc", + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + result = distributed_map_handler( + source=["a", "b"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr1"), + ) + + assert result.status is DistributedMapStatus.SUCCEEDED + assert result.completion_reason is DistributedMapCompletionReason.ALL_COMPLETED + assert result.success_count == 5 + assert result.failure_count == 0 + assert result.total_count == 5 + assert result.distributed_map_id == "abc" + mock_state.create_checkpoint.assert_not_called() + + +def test_map_run_handler_resolves_non_success_without_raising(): + """Test a non-SUCCEEDED run resolves with a summary rather than raising. + + The durable operation succeeded (it delivered a result), but the run's own + status is FAILED. distributed_map must return the summary, not raise. + """ + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + operation = Operation( + operation_id="mr2", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + status="FAILED", + completion_reason="FAILURE_TOLERANCE_EXCEEDED", + success_count=3, + failure_count=2, + unprocessed_count=1, + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + result = distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr2"), + ) + + assert result.status is DistributedMapStatus.FAILED + assert ( + result.completion_reason + is DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED + ) + assert result.failure_count == 2 + assert result.has_failure is True + + +def test_map_run_handler_succeeded_no_details_raises(): + """Test a succeeded operation carrying no DistributedMapDetails raises ExecutionError.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + operation = Operation( + operation_id="mr3", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=None, + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + with pytest.raises(ExecutionError): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr3"), + ) + + +def test_map_run_handler_already_started(): + """Test distributed_map_handler suspends when the operation is already started.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + operation = Operation( + operation_id="mr5", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + with pytest.raises( + SuspendExecution, match="Map run mr5 started, suspending for completion" + ): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr5"), + ) + + +def test_map_run_handler_new_operation(): + """Test distributed_map_handler creates a START checkpoint for a new operation.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + started_op = Operation( + operation_id="mr6", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + started = CheckpointedResult.create_from_operation(started_op) + mock_state.get_checkpoint_result.side_effect = [not_found, started] + + with pytest.raises(SuspendExecution): + distributed_map_handler( + source=["a", "b", "c"], + processor="test_processor", + max_concurrency=42, + state=mock_state, + operation_identifier=_identifier("mr6"), + ) + + mock_state.create_checkpoint.assert_called_once() + operation_update = mock_state.create_checkpoint.call_args[1]["operation_update"] + assert operation_update.operation_id == "mr6" + assert operation_update.operation_type == OperationType.DISTRIBUTED_MAP + assert operation_update.action == OperationAction.START + assert operation_update.name == "test_map_run" + + distributed_map_options = operation_update.to_dict()["DistributedMapOptions"] + assert distributed_map_options["MaxConcurrency"] == 42 + assert distributed_map_options["Processor"]["FunctionName"] == "test_processor" + assert distributed_map_options["Source"]["InlineSourceConfig"]["Items"] == [ + "a", + "b", + "c", + ] + + +def test_map_run_handler_no_config(): + """Test distributed_map_handler uses a default config when none is provided.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + started_op = Operation( + operation_id="mr7", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + started = CheckpointedResult.create_from_operation(started_op) + mock_state.get_checkpoint_result.side_effect = [not_found, started] + + with pytest.raises(SuspendExecution): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr7"), + config=None, + ) + + mock_state.create_checkpoint.assert_called_once() + + +# ============================================================================ +# Immediate Response Handling Tests +# ============================================================================ + + +def test_map_run_immediate_response_get_checkpoint_result_called_twice(): + """Test get_checkpoint_result is called twice when a checkpoint is created.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + started_op = Operation( + operation_id="mr8", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + started = CheckpointedResult.create_from_operation(started_op) + mock_state.get_checkpoint_result.side_effect = [not_found, started] + + with pytest.raises(SuspendExecution): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr8"), + ) + + assert mock_state.get_checkpoint_result.call_count == 2 + + +def test_map_run_immediate_response_create_checkpoint_is_sync_true(): + """Test create_checkpoint is called with is_sync=True.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + started_op = Operation( + operation_id="mr9", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + started = CheckpointedResult.create_from_operation(started_op) + mock_state.get_checkpoint_result.side_effect = [not_found, started] + + with pytest.raises(SuspendExecution): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr9"), + ) + + mock_state.create_checkpoint.assert_called_once() + assert mock_state.create_checkpoint.call_args[1]["is_sync"] is True + + +def test_map_run_immediate_response_immediate_success(): + """Test immediate success: second check returns SUCCEEDED, summary returned.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + succeeded_op = Operation( + operation_id="mr10", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + status="SUCCEEDED", + completion_reason="ALL_COMPLETED", + success_count=1, + ), + ) + succeeded = CheckpointedResult.create_from_operation(succeeded_op) + mock_state.get_checkpoint_result.side_effect = [not_found, succeeded] + + result = distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr10"), + ) + + assert result.status is DistributedMapStatus.SUCCEEDED + assert result.success_count == 1 + mock_state.create_checkpoint.assert_called_once() + assert mock_state.get_checkpoint_result.call_count == 2 + + +def test_map_run_immediate_response_no_immediate_response(): + """Test no immediate response: second check returns STARTED, suspends.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + started_op = Operation( + operation_id="mr12", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + started = CheckpointedResult.create_from_operation(started_op) + mock_state.get_checkpoint_result.side_effect = [not_found, started] + + with pytest.raises(SuspendExecution): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr12"), + ) + + mock_state.create_checkpoint.assert_called_once() + assert mock_state.get_checkpoint_result.call_count == 2 + + +def test_map_run_immediate_response_already_completed(): + """Test already completed: first check is SUCCEEDED, no checkpoint created.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + succeeded_op = Operation( + operation_id="mr13", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + status="SUCCEEDED", completion_reason="ALL_COMPLETED" + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(succeeded_op) + ) + + result = distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr13"), + ) + + assert result.status is DistributedMapStatus.SUCCEEDED + mock_state.create_checkpoint.assert_not_called() + assert mock_state.get_checkpoint_result.call_count == 1 + + +@patch( + "aws_durable_execution_sdk_python.operation.distributed_map.suspend_with_optional_resume_delay" +) +def test_map_run_handler_suspend_does_not_raise(mock_suspend): + """Test distributed_map_handler raises ExecutionError if suspend does not raise.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + started_op = Operation( + operation_id="mr14", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + started = CheckpointedResult.create_from_operation(started_op) + mock_state.get_checkpoint_result.side_effect = [not_found, started] + + mock_suspend.return_value = None + + with pytest.raises( + ExecutionError, + match="suspend_with_optional_resume_delay should have raised an exception, but did not.", + ): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr14"), + ) + + mock_suspend.assert_called_once() + + +# ============================================================================ +# Wire serialization and result-collection tests (slices 3-4) +# ============================================================================ + + +def _start_options(state_calls, source, processor, max_concurrency, config): + """Run the executor through the new-operation path and return the sent options dict.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + mock_state.get_checkpoint_result.side_effect = state_calls + + executor = DistributedMapOperationExecutor( + source=source, + processor=processor, + max_concurrency=max_concurrency, + state=mock_state, + operation_identifier=_identifier("mrw"), + config=config, + ) + with pytest.raises(SuspendExecution): + executor.process() + update = mock_state.create_checkpoint.call_args[1]["operation_update"] + return update.to_dict()["DistributedMapOptions"] + + +def _new_op_state_calls(): + not_found = CheckpointedResult.create_not_found() + started = CheckpointedResult.create_from_operation( + Operation( + operation_id="mrw", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + ) + return [not_found, started] + + +def test_processor_report_failed_items_sets_response_types(): + """report_failed_items serializes FunctionResponseTypes=ReportBatchItemFailures.""" + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_failed_items("proc", batch_size=25), + max_concurrency=4, + config=DistributedMapConfig(), + ) + processor = options["Processor"] + assert processor["FunctionName"] == "proc" + assert processor["FunctionResponseTypes"] == ["ReportBatchItemFailures"] + assert processor["BatchSize"] == 25 + + +def test_processor_unlimited_retries_maps_to_negative_one(): + """ProcessorRetryConfig.UNLIMITED serializes to MaxRetryAttempts=-1.""" + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_item_results( + "proc", + retry=ProcessorRetryConfig( + max_retry_attempts=ProcessorRetryConfig.UNLIMITED, + max_retry_duration=Duration.from_hours(1), + ), + ), + max_concurrency=1, + config=DistributedMapConfig(), + ) + processor = options["Processor"] + assert processor["FunctionResponseTypes"] == ["ReportBatchItemResults"] + assert processor["MaxRetryAttempts"] == -1 + assert processor["MaxRetryDurationSeconds"] == 3600 + + +def test_processor_explicit_retry_attempts_pass_through(): + """A plain int retry count passes through unchanged (no sentinel mapping).""" + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_batch_outcome( + "proc", retry=ProcessorRetryConfig(max_retry_attempts=0) + ), + max_concurrency=1, + config=DistributedMapConfig(), + ) + assert options["Processor"]["MaxRetryAttempts"] == 0 + # batch mode reports no per-item response types + assert "FunctionResponseTypes" not in options["Processor"] + + +def test_s3_source_serializes_config(): + """An S3 json_lines source serializes to an S3SourceConfig block.""" + options = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.S3.json_lines( + "s3://bucket/data.jsonl", max_items=500 + ), + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=2, + config=DistributedMapConfig(), + ) + source = options["Source"] + assert source["Type"] == "S3" + assert source["MaxItemsToRead"] == 500 + assert source["S3SourceConfig"]["Bucket"] == "bucket" + assert source["S3SourceConfig"]["Key"] == "data.jsonl" + assert source["S3SourceConfig"]["Format"] == "JSON_LINES" + + +def test_full_config_serializes_all_blocks(): + """Completion, destination, timeout, and result-collection blocks all serialize.""" + config = DistributedMapConfig( + completion_config=DistributedMapCompletionConfig.failure_percentage( + 5, minimum_sample_size=200 + ), + destination=DistributedMapDestinationConfig( + on_success=DistributedMapDestination.S3.successes("s3://out/ok"), + on_failure=DistributedMapDestination.S3.failures("s3://out/bad"), + ), + timeout=Duration.from_minutes(30), + collect_results=True, + ) + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_item_results("proc"), + max_concurrency=8, + config=config, + ) + assert options["CompletionConfig"] == { + "ToleratedFailurePercentage": 5, + "MinimumSampleSize": 200, + } + on_success = options["Destination"]["OnSuccess"] + assert on_success["Type"] == "S3" + assert on_success["S3DestinationConfig"]["Bucket"] == "out" + assert on_success["S3DestinationConfig"]["KeyPrefix"] == "ok" + on_failure = options["Destination"]["OnFailure"] + assert on_failure["Include"] == ["INPUT", "ERROR"] + assert on_failure["S3DestinationConfig"]["Bucket"] == "out" + assert options["TimeoutSeconds"] == 1800 + assert options["ResultCollection"] == {"Mode": "INLINE"} + + +def test_collect_results_returns_map_run_result_with_items(): + """When collect_results is set, a DistributedMapResult with per-item results is built.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + error = ErrorObject(message="boom", type="ItemError", data=None, stack_trace=None) + operation = Operation( + operation_id="mrr", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + status="SUCCEEDED", + completion_reason="ALL_COMPLETED", + success_count=1, + failure_count=1, + unprocessed_count=0, + total_count=2, + results=( + DistributedMapResultItemWire( + item_id="0", status="SUCCEEDED", output=42 + ), + DistributedMapResultItemWire(item_id="1", status="FAILED", error=error), + ), + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + result = distributed_map_handler( + source=["a", "b"], + processor=DistributedMapProcessor.report_item_results("proc"), + max_concurrency=2, + state=mock_state, + operation_identifier=_identifier("mrr"), + config=DistributedMapConfig(collect_results=True), + ) + + assert isinstance(result, DistributedMapResult) + assert len(result.all) == 2 + assert result.get_results() == [42] + errors = result.get_errors() + assert len(errors) == 1 + assert errors[0].error_type == "ItemError" + assert errors[0].error_message == "boom" + + +def test_collect_results_disabled_returns_plain_summary(): + """Without collect_results, a plain DistributedMapSummary is returned (not DistributedMapResult).""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="mrs", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + status="SUCCEEDED", + completion_reason="ALL_COMPLETED", + success_count=2, + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + result = distributed_map_handler( + source=["a", "b"], + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=2, + state=mock_state, + operation_identifier=_identifier("mrs"), + ) + + assert isinstance(result, DistributedMapSummary) + assert not isinstance(result, DistributedMapResult) + assert result.success_count == 2 + + +def test_csv_source_header_location_wire(): + """CSV headers map to GIVEN; expected_columns stays client-side (FIRST_ROW, not sent).""" + # headers -> HeaderLocation GIVEN, headers sent + given = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.S3.csv("s3://b/data.csv", headers=["a", "b"]), + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + csv_opts = given["Source"]["S3SourceConfig"]["CsvFormatOptions"] + assert csv_opts["HeaderLocation"] == "GIVEN" + assert csv_opts["Headers"] == ["a", "b"] + assert csv_opts["Delimiter"] == "COMMA" + + # no headers -> HeaderLocation FIRST_ROW + first_row = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.S3.csv("s3://b/data.csv"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + s3_cfg = first_row["Source"]["S3SourceConfig"] + assert s3_cfg["CsvFormatOptions"]["HeaderLocation"] == "FIRST_ROW" + assert "Headers" not in s3_cfg["CsvFormatOptions"] + + +# ============================================================================ +# Call-site validation and wire-serdes tests +# ============================================================================ + + +def test_retry_duration_out_of_range_rejected(): + with pytest.raises(ValidationError, match="between 1 minute and 6 hours"): + ProcessorRetryConfig(max_retry_duration=Duration.from_seconds(30)) + with pytest.raises(ValidationError, match="between 1 minute and 6 hours"): + ProcessorRetryConfig(max_retry_duration=Duration.from_hours(7)) + + +def test_expected_bucket_owner_must_be_12_digits(): + with pytest.raises(ValidationError, match="12-digit"): + DistributedMapSource.S3.json_lines( + "s3://b/k.jsonl", expected_bucket_owner="123" + ) + + +def test_csv_headers_duplicates_rejected(): + with pytest.raises(ValidationError, match="duplicates"): + DistributedMapSource.S3.csv("s3://b/k.csv", headers=["a", "a"]) + + +def test_json_lines_requires_key(): + with pytest.raises(ValidationError, match="object key"): + DistributedMapSource.S3.json_lines("s3://bucket-only") + + +def test_timeout_out_of_range_rejected(): + with pytest.raises(ValidationError, match="at most 90 days"): + DistributedMapConfig(timeout=Duration.from_days(91)) + + +def test_item_serdes_requires_collect_results(): + with pytest.raises(ValidationError, match="requires collect_results"): + DistributedMapConfig(item_serdes=DEFAULT_JSON_SERDES) + + +def test_empty_function_name_rejected(): + with pytest.raises(ValidationError, match="non-empty"): + DistributedMapProcessor.report_batch_outcome("") + + +def test_inline_source_over_1mb_rejected(): + big = ["x" * 100_000] * 12 # ~1.2 MB serialized + with pytest.raises(ValidationError, match="1 MB limit"): + _start_options( + _new_op_state_calls(), + source=big, + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + + +def test_reader_state_serialized_into_wire_and_capped(): + # typed initial_state is serialized into the opaque wire string + options = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.Reader.from_function( + "reader", initial_state={"page": 0} + ), + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + reader_cfg = options["Source"]["ReaderFunctionSourceConfig"] + assert reader_cfg["FunctionName"] == "reader" + assert reader_cfg["InitialState"] == '{"page": 0}' + + # oversize state is rejected + with pytest.raises(ValidationError, match="32 KB limit"): + _start_options( + _new_op_state_calls(), + source=DistributedMapSource.Reader.from_function( + "reader", initial_state="x" * 40_000 + ), + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + + +def test_invalid_function_reference_rejected(): + with pytest.raises(ValidationError, match="Lambda function reference"): + DistributedMapProcessor.report_batch_outcome("bad name!") + + +def test_valid_function_references_accepted(): + for ref in ( + "my-func", + "my-func:PROD", + "123456789012:function:my-func", + "arn:aws:lambda:us-east-1:123456789012:function:my-func", + "arn:aws:lambda:us-east-1:123456789012:function:my-func:1", + ): + # Should not raise. + DistributedMapProcessor.report_batch_outcome(ref) + + +def test_inline_custom_serdes_applied_to_wire(): + """A custom inline serdes transforms each item's wire value.""" + from aws_durable_execution_sdk_python.serdes import SerDes + + class _UpperSerDes(SerDes): + def serialize(self, value, _serdes_context): # noqa: ANN001, ANN201 + return json.dumps(value.upper()) + + def deserialize(self, data, _serdes_context): # noqa: ANN001, ANN201 + return json.loads(data).lower() + + options = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.inline(["a", "b"], serdes=_UpperSerDes()), + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + assert options["Source"]["InlineSourceConfig"]["Items"] == ["A", "B"] + + +# ============================================================================ +# Coverage-gap tests: result helpers, from_dict, destinations, source variants +# ============================================================================ + + +def test_summary_throw_if_error(): + ok = DistributedMapSummary( + status=DistributedMapStatus.SUCCEEDED, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=2, + failure_count=0, + unprocessed_count=0, + ) + ok.throw_if_error() # no raise + + failed = DistributedMapSummary( + status=DistributedMapStatus.FAILED, + completion_reason=DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED, + success_count=0, + failure_count=1, + unprocessed_count=0, + ) + with pytest.raises(DistributedMapError): + failed.throw_if_error() + + succeeded_with_failures = DistributedMapSummary( + status=DistributedMapStatus.SUCCEEDED, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=1, + failure_count=1, + unprocessed_count=0, + ) + with pytest.raises(DistributedMapError): + succeeded_with_failures.throw_if_error() + + +def test_map_run_result_succeeded_failed_filters(): + items = [ + DistributedMapResultItem(item_id="0", status="SUCCEEDED", output=1), + DistributedMapResultItem( + item_id="1", + status="FAILED", + error=DistributedMapItemError(error_type="E", error_message="boom"), + ), + ] + result = DistributedMapResult( + status=DistributedMapStatus.SUCCEEDED, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=1, + failure_count=1, + unprocessed_count=0, + all=items, + ) + assert [i.item_id for i in result.succeeded()] == ["0"] + assert [i.item_id for i in result.failed()] == ["1"] + assert result.get_results() == [1] + assert result.get_errors()[0].error_message == "boom" + + +def test_map_run_details_from_dict_parses_results(): + data = { + "Status": "SUCCEEDED", + "CompletionReason": "ALL_COMPLETED", + "SuccessCount": 1, + "FailureCount": 1, + "UnprocessedCount": 0, + "TotalCount": 2, + "DistributedMapRunArn": "arn:aws:lambda:us-east-1:123456789012:map-run:x", + "Results": [ + {"ItemId": "0", "Status": "SUCCEEDED", "Output": 5}, + { + "ItemId": "1", + "Status": "FAILED", + "Error": {"ErrorType": "E", "ErrorMessage": "boom"}, + }, + ], + } + details = DistributedMapDetails.from_dict(data) + assert details.status == "SUCCEEDED" + assert details.total_count == 2 + assert details.results is not None + assert details.results[0].item_id == "0" + assert details.results[1].error is not None + assert details.results[1].error.type == "E" + + +def test_map_run_options_from_dict_round_trip(): + sent = _start_options( + _new_op_state_calls(), + source=["a", "b"], + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=7, + config=DistributedMapConfig(), + ) + parsed = DistributedMapOptions.from_dict(sent) + assert parsed.max_concurrency == 7 + assert parsed.source.source_type == "INLINE" + assert parsed.processor.function_name == "proc" + + +def test_destination_only_success(): + opts = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_success=DistributedMapDestination.S3.successes("s3://out/ok") + ) + ), + ) + dest = opts["Destination"] + assert "OnSuccess" in dest + assert "OnFailure" not in dest + + +def test_destination_only_failure(): + opts = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_failure=DistributedMapDestination.S3.failures("s3://out/bad") + ) + ), + ) + dest = opts["Destination"] + assert "OnFailure" in dest + assert "OnSuccess" not in dest + + +def test_s3_objects_source_wire(): + opts = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.S3.objects("s3://b/prefix/"), + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + s3 = opts["Source"]["S3SourceConfig"] + assert s3["Transform"] == "NONE" + assert s3["KeyPrefix"] == "prefix/" + assert "Format" not in s3 + + +def test_s3_flattened_json_lines_source_wire(): + opts = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.S3.flattened_json_lines("s3://b/prefix/"), + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + s3 = opts["Source"]["S3SourceConfig"] + assert s3["Transform"] == "LOAD_AND_FLATTEN" + assert s3["Format"] == "JSON_LINES" + + +# ============================================================================ +# Coverage-gap tests (batch 2): validations, wire branches, round-trips +# ============================================================================ + + +def _start_executor(source, processor, config, max_concurrency=1): + """Build an executor on the new-operation path (for error-path tests).""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + mock_state.get_checkpoint_result.side_effect = _new_op_state_calls() + return DistributedMapOperationExecutor( + source=source, + processor=processor, + max_concurrency=max_concurrency, + state=mock_state, + operation_identifier=_identifier("mrw"), + config=config, + ) + + +# --- Completion config validations --- + + +def test_completion_count_and_percentage_mutually_exclusive(): + with pytest.raises(ValidationError, match="mutually exclusive"): + DistributedMapCompletionConfig( + tolerated_failure_count=1, tolerated_failure_percentage=5 + ) + + +def test_completion_sample_size_requires_percentage(): + with pytest.raises(ValidationError, match="minimum_sample_size"): + DistributedMapCompletionConfig(minimum_sample_size=10) + + +def test_completion_negative_count_rejected(): + with pytest.raises(ValidationError, match="non-negative"): + DistributedMapCompletionConfig(tolerated_failure_count=-1) + + +def test_completion_percentage_out_of_range_rejected(): + with pytest.raises(ValidationError, match="between 0 and 100"): + DistributedMapCompletionConfig(tolerated_failure_percentage=150) + + +def test_completion_sample_size_below_one_rejected(): + with pytest.raises(ValidationError, match="at least 1"): + DistributedMapCompletionConfig( + tolerated_failure_percentage=5, minimum_sample_size=0 + ) + + +def test_completion_failure_count_factory(): + assert DistributedMapCompletionConfig.failure_count(3).tolerated_failure_count == 3 + + +def test_empty_completion_config_omitted_from_wire(): + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig(completion_config=DistributedMapCompletionConfig()), + ) + assert "CompletionConfig" not in options + + +# --- Retry duration lower bound --- + + +def test_retry_duration_below_minimum_rejected(): + with pytest.raises(ValidationError, match="1 minute and 6 hours"): + ProcessorRetryConfig(max_retry_duration=Duration.from_seconds(30)) + + +# --- Source validations / variants --- + + +def test_max_items_below_one_rejected(): + with pytest.raises(ValidationError, match="at least 1"): + DistributedMapSource.S3.json_lines("s3://b/k.jsonl", max_items=0) + + +def test_csv_requires_object_key(): + with pytest.raises(ValidationError, match="csv requires an S3 object key"): + DistributedMapSource.S3.csv("s3://bucket") + + +def test_objects_whole_bucket_prefix_wire(): + options = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.S3.objects("s3://bucket"), + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + s3 = options["Source"]["S3SourceConfig"] + assert s3["KeyPrefix"] == "" + assert s3["Transform"] == "NONE" + assert "Key" not in s3 + + +def test_flattened_csv_source_wire(): + options = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.S3.flattened_csv( + "s3://b/prefix", headers=["a", "b"] + ), + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + s3 = options["Source"]["S3SourceConfig"] + assert s3["Transform"] == "LOAD_AND_FLATTEN" + assert s3["Format"] == "CSV" + assert s3["CsvFormatOptions"]["HeaderLocation"] == "GIVEN" + assert s3["CsvFormatOptions"]["Headers"] == ["a", "b"] + + +def test_reader_source_without_initial_state_omits_state(): + options = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.Reader.from_function("reader"), + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + reader = options["Source"]["ReaderFunctionSourceConfig"] + assert reader["FunctionName"] == "reader" + assert "InitialState" not in reader + + +def test_inline_serdes_non_json_rejected(): + from aws_durable_execution_sdk_python.serdes import SerDes + + class _BadSerDes(SerDes): + def serialize(self, value, _serdes_context): # noqa: ANN001, ANN201, ARG002 + return "{not json" + + def deserialize(self, data, _serdes_context): # noqa: ANN001, ANN201, ARG002 + return data + + executor = _start_executor( + DistributedMapSource.inline([1], serdes=_BadSerDes()), + DistributedMapProcessor.report_batch_outcome("p"), + DistributedMapConfig(), + ) + with pytest.raises(ValidationError, match="must produce a JSON value"): + executor.process() + + +def test_unsupported_source_type_raises(): + executor = _start_executor( + DistributedMapSource(source_type="BOGUS"), + DistributedMapProcessor.report_batch_outcome("p"), + DistributedMapConfig(), + ) + with pytest.raises(ExecutionError, match="Unsupported map run source type"): + executor.process() + + +# --- Destination include permutations + validation --- + + +def test_success_destination_include_input_and_owner(): + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_success=DistributedMapDestination.S3.successes( + "s3://out/ok", + include_input=True, + include_output=True, + expected_bucket_owner="123456789012", + ) + ) + ), + ) + on_success = options["Destination"]["OnSuccess"] + assert on_success["Include"] == ["INPUT", "OUTPUT"] + assert on_success["S3DestinationConfig"]["ExpectedBucketOwner"] == "123456789012" + assert "OnFailure" not in options["Destination"] + + +def test_failure_destination_error_only_and_owner(): + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_failure=DistributedMapDestination.S3.failures( + "s3://out/bad", + include_input=False, + include_error=True, + expected_bucket_owner="123456789012", + ) + ) + ), + ) + on_failure = options["Destination"]["OnFailure"] + assert on_failure["Include"] == ["ERROR"] + assert on_failure["S3DestinationConfig"]["ExpectedBucketOwner"] == "123456789012" + + +def test_success_destination_all_false_rejected(): + with pytest.raises(ValidationError, match="success destination must include"): + DistributedMapDestination.S3.successes( + "s3://out/ok", include_input=False, include_output=False + ) + + +def test_failure_destination_all_false_rejected(): + with pytest.raises(ValidationError, match="failure destination must include"): + DistributedMapDestination.S3.failures( + "s3://out/bad", include_input=False, include_error=False + ) + + +# --- Unknown backend enum on resume --- + + +def test_unknown_status_raises_execution_error(): + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="mrx", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + status="BOGUS", + completion_reason="ALL_COMPLETED", + success_count=0, + failure_count=0, + unprocessed_count=0, + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + with pytest.raises(ExecutionError, match="Unknown map run status"): + distributed_map_handler( + source=["a"], + processor="p", + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("mrx"), + ) + + +# --- Summary / result helpers --- + + +def test_summary_distributed_map_id_none_without_arn(): + summary = DistributedMapSummary( + status=DistributedMapStatus.SUCCEEDED, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=0, + failure_count=0, + unprocessed_count=0, + ) + assert summary.distributed_map_id is None + assert summary.has_failure is False + + +def _result(status, completion_reason, failure_count, items): + return DistributedMapResult( + status=status, + completion_reason=completion_reason, + success_count=len(items) - failure_count, + failure_count=failure_count, + unprocessed_count=0, + all=items, + ) + + +def test_result_throw_if_error_raises_first_item_error(): + result = _result( + DistributedMapStatus.SUCCEEDED, + DistributedMapCompletionReason.ALL_COMPLETED, + 1, + [ + DistributedMapResultItem(item_id="0", status="SUCCEEDED", output=1), + DistributedMapResultItem( + item_id="1", + status="FAILED", + error=DistributedMapItemError(error_type="E", error_message="boom"), + ), + ], + ) + with pytest.raises(DistributedMapError, match="E: boom"): + result.throw_if_error() + + +def test_result_throw_if_error_names_item_without_detail(): + result = _result( + DistributedMapStatus.SUCCEEDED, + DistributedMapCompletionReason.ALL_COMPLETED, + 1, + [DistributedMapResultItem(item_id="7", status="FAILED", error=None)], + ) + with pytest.raises(DistributedMapError, match="item 7 failed"): + result.throw_if_error() + + +def test_result_throw_if_error_falls_back_to_summary_when_no_items(): + result = DistributedMapResult( + status=DistributedMapStatus.SUCCEEDED, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=0, + failure_count=2, + unprocessed_count=0, + all=[], + ) + with pytest.raises(DistributedMapError, match="2 item"): + result.throw_if_error() + + +def test_result_throw_if_error_run_level_failure(): + result = DistributedMapResult( + status=DistributedMapStatus.FAILED, + completion_reason=DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED, + success_count=0, + failure_count=1, + unprocessed_count=0, + all=[], + ) + with pytest.raises(DistributedMapError, match="Map run ended FAILED"): + result.throw_if_error() + + +def test_result_throw_if_error_clean_success_does_not_raise(): + result = DistributedMapResult( + status=DistributedMapStatus.SUCCEEDED, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=1, + failure_count=0, + unprocessed_count=0, + all=[DistributedMapResultItem(item_id="0", status="SUCCEEDED", output=1)], + ) + result.throw_if_error() + + +# --- Wire round-trips (lambda_service) --- + + +def test_result_item_wire_round_trip(): + item = DistributedMapResultItemWire.from_dict( + {"ItemId": "0", "Status": "SUCCEEDED", "Output": {"x": 1}} + ) + assert item.output == {"x": 1} + assert item.to_dict() == {"ItemId": "0", "Status": "SUCCEEDED", "Output": {"x": 1}} + + failed = DistributedMapResultItemWire.from_dict( + { + "ItemId": "1", + "Status": "FAILED", + "Error": {"ErrorType": "E", "ErrorMessage": "boom"}, + } + ) + dumped = failed.to_dict() + assert dumped["Error"]["ErrorType"] == "E" + assert "Output" not in dumped + + +def test_operation_to_dict_round_trip_preserves_results(): + op = Operation( + operation_id="opx", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + status="SUCCEEDED", + completion_reason="ALL_COMPLETED", + success_count=1, + failure_count=1, + unprocessed_count=0, + total_count=2, + distributed_map_run_arn="arn:aws:lambda:us-east-1:123456789012:map-run:z", + completion_details="done", + results=( + DistributedMapResultItemWire(item_id="0", status="SUCCEEDED", output=5), + DistributedMapResultItemWire( + item_id="1", + status="FAILED", + error=ErrorObject( + message="boom", type="E", data=None, stack_trace=None + ), + ), + ), + ), + ) + block = op.to_dict()["DistributedMapDetails"] + assert block["Results"][0] == {"ItemId": "0", "Status": "SUCCEEDED", "Output": 5} + assert block["DistributedMapRunArn"].endswith("map-run:z") + assert block["CompletionDetails"] == "done" + assert block["TotalCount"] == 2 + + parsed = Operation.from_dict(op.to_dict()) + assert parsed.distributed_map_details is not None + assert parsed.distributed_map_details.results[0].output == 5 + assert parsed.distributed_map_details.results[1].error.type == "E" + + +def test_options_full_round_trip(): + options_dict = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.S3.csv("s3://b/f.csv", headers=["a"]), + processor=DistributedMapProcessor.report_item_results( + "proc", + batch_size=5, + retry=ProcessorRetryConfig( + max_retry_attempts=ProcessorRetryConfig.UNLIMITED, + max_retry_duration=Duration.from_minutes(10), + ), + durable_execution_name_prefix="pfx", + ), + max_concurrency=3, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_success=DistributedMapDestination.S3.successes("s3://o/ok"), + on_failure=DistributedMapDestination.S3.failures("s3://o/bad"), + ), + completion_config=DistributedMapCompletionConfig.failure_count(2), + timeout=Duration.from_minutes(5), + collect_results=True, + ), + ) + parsed = DistributedMapOptions.from_dict(options_dict) + assert parsed.max_concurrency == 3 + assert parsed.source.source_type == "S3" + assert parsed.processor.function_response_types == ("ReportBatchItemResults",) + assert parsed.processor.max_retry_attempts == -1 + assert parsed.processor.durable_execution_name_prefix == "pfx" + assert parsed.destination is not None + assert parsed.completion_config.tolerated_failure_count == 2 + assert parsed.result_collection.mode == "INLINE" + assert parsed.timeout_seconds == 300 + assert parsed.to_dict()["MaxConcurrency"] == 3 + + +# ============================================================================ +# Coverage-gap tests (batch 3): remaining validation and wire branches +# ============================================================================ + + +def test_invalid_s3_uri_rejected(): + with pytest.raises(ValidationError, match="Invalid S3 URI"): + DistributedMapSource.S3.json_lines("s3:///key.jsonl") + + +def test_csv_empty_headers_rejected(): + with pytest.raises(ValidationError, match="must be non-empty"): + DistributedMapSource.S3.csv("s3://b/f.csv", headers=[]) + + +def test_negative_retry_attempts_rejected(): + with pytest.raises(ValidationError, match="non-negative"): + ProcessorRetryConfig(max_retry_attempts=-5) + + +def test_batch_size_out_of_range_rejected(): + with pytest.raises(ValidationError, match="between 1 and 10000"): + DistributedMapProcessor.report_batch_outcome("p", batch_size=0) + + +def test_csv_first_row_no_headers_wire(): + options = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.S3.csv("s3://b/f.csv"), + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + csv_opts = options["Source"]["S3SourceConfig"]["CsvFormatOptions"] + assert csv_opts["HeaderLocation"] == "FIRST_ROW" + assert "Headers" not in csv_opts + assert csv_opts["Delimiter"] == "COMMA" + + +def test_s3_source_expected_bucket_owner_wire(): + options = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.S3.json_lines( + "s3://b/k.jsonl", expected_bucket_owner="123456789012" + ), + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + assert options["Source"]["S3SourceConfig"]["ExpectedBucketOwner"] == "123456789012" + + +def test_empty_destination_config_omitted(): + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig(destination=DistributedMapDestinationConfig()), + ) + assert "Destination" not in options + + +def test_success_destination_input_only(): + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_success=DistributedMapDestination.S3.successes( + "s3://out/ok", include_input=True, include_output=False + ) + ) + ), + ) + assert options["Destination"]["OnSuccess"]["Include"] == ["INPUT"] + + +def test_failure_destination_input_only(): + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_failure=DistributedMapDestination.S3.failures( + "s3://out/bad", include_input=True, include_error=False + ) + ) + ), + ) + assert options["Destination"]["OnFailure"]["Include"] == ["INPUT"] + + +def test_reader_source_options_round_trip(): + options_dict = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.Reader.from_function( + "reader", initial_state={"page": 0} + ), + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + parsed = DistributedMapOptions.from_dict(options_dict) + assert parsed.source.source_type == "READER_FUNCTION" + assert parsed.source.reader_config["FunctionName"] == "reader" + + +def test_operation_to_dict_minimal_details_omits_optionals(): + op = Operation( + operation_id="opm", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + status="SUCCEEDED", + completion_reason="ALL_COMPLETED", + success_count=1, + failure_count=0, + unprocessed_count=0, + ), + ) + block = op.to_dict()["DistributedMapDetails"] + assert block["Status"] == "SUCCEEDED" + assert "DistributedMapRunArn" not in block + assert "CompletionDetails" not in block + assert "TotalCount" not in block + assert "Results" not in block + + +@pytest.mark.parametrize( + "status", + [OperationStatus.FAILED, OperationStatus.STOPPED, OperationStatus.TIMED_OUT], +) +def test_operation_level_terminal_failure_raises(status): + """An operation-level terminal failure raises rather than hanging.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="mrf", + operation_type=OperationType.DISTRIBUTED_MAP, + status=status, + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + with pytest.raises(DistributedMapError): + distributed_map_handler( + source=["a"], + processor="p", + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("mrf"), + ) + + +# ============================================================================ +# Coverage-gap tests (batch 4): serdes decode, falsy output, duration-only retry +# ============================================================================ + + +def test_custom_item_serdes_applied_on_decode(): + """A custom item_serdes transforms each per-item output on decode.""" + from aws_durable_execution_sdk_python.serdes import SerDes + + class _UpperSerDes(SerDes): + def serialize(self, value, _serdes_context): # noqa: ANN001, ANN201 + return json.dumps(value) + + def deserialize(self, data, _serdes_context): # noqa: ANN001, ANN201 + return json.loads(data).upper() + + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="cs", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + status="SUCCEEDED", + completion_reason="ALL_COMPLETED", + success_count=1, + failure_count=0, + unprocessed_count=0, + results=( + DistributedMapResultItemWire( + item_id="0", status="SUCCEEDED", output="abc" + ), + ), + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + result = distributed_map_handler( + source=["a"], + processor=DistributedMapProcessor.report_item_results("proc"), + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("cs"), + config=DistributedMapConfig(collect_results=True, item_serdes=_UpperSerDes()), + ) + assert result.get_results() == ["ABC"] + + +@pytest.mark.parametrize("value", [0, False, "", [], {}]) +def test_falsy_output_round_trips(value): + """A falsy-but-present output survives wire round-trip and decode (not dropped).""" + wire = DistributedMapResultItemWire(item_id="0", status="SUCCEEDED", output=value) + assert wire.to_dict()["Output"] == value + + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="fo", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + status="SUCCEEDED", + completion_reason="ALL_COMPLETED", + success_count=1, + failure_count=0, + unprocessed_count=0, + results=(wire,), + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + result = distributed_map_handler( + source=["a"], + processor=DistributedMapProcessor.report_item_results("proc"), + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("fo"), + config=DistributedMapConfig(collect_results=True), + ) + assert result.get_results() == [value] + + +def test_processor_retry_duration_only(): + """A retry config with only a duration sends the duration and omits attempts.""" + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_batch_outcome( + "proc", + retry=ProcessorRetryConfig(max_retry_duration=Duration.from_minutes(5)), + ), + max_concurrency=1, + config=DistributedMapConfig(), + ) + processor = options["Processor"] + assert processor["MaxRetryDurationSeconds"] == 300 + assert "MaxRetryAttempts" not in processor