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 69f7bb164..745479a68 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 @@ -16,7 +16,10 @@ CompletionItemStatus, CompletionOutcome, CompletionStatus, + DistributedMapItemStatus, + DistributedMapStatus, ParallelBranch, + ProcessorResponseMode, complete_batch, continue_batch, ) @@ -27,6 +30,18 @@ durable_wait_for_callback, durable_with_child_context, ) +from aws_durable_execution_sdk_python.dmap.handlers import ( + distributed_map_batch_handler, + distributed_map_item_handler, + distributed_map_reader, + durable_distributed_map_batch_handler, + durable_distributed_map_item_handler, +) +from aws_durable_execution_sdk_python.dmap.models import ( + DistributedMapResult, + DistributedMapResultItem, + DistributedMapSummary, +) # Most common exceptions - users need to handle these exceptions from aws_durable_execution_sdk_python.exceptions import ( @@ -35,6 +50,7 @@ CallbackSubmitterError, CallbackTimeoutError, ChildContextError, + DistributedMapError, DurableExecutionsError, DurableOperationError, ExecutionError, @@ -69,6 +85,12 @@ "CompletionItemStatus", "CompletionOutcome", "CompletionStatus", + "DistributedMapError", + "DistributedMapItemStatus", + "DistributedMapResult", + "DistributedMapResultItem", + "DistributedMapStatus", + "DistributedMapSummary", "DurableContext", "DurableExecutionsError", "DurableOperationError", @@ -76,6 +98,7 @@ "InvocationError", "InvokeError", "ParallelBranch", + "ProcessorResponseMode", "PluginLoadError", "RetryableSerDesError", "SerDesError", @@ -87,6 +110,11 @@ "__version__", "complete_batch", "continue_batch", + "distributed_map_batch_handler", + "distributed_map_item_handler", + "distributed_map_reader", + "durable_distributed_map_batch_handler", + "durable_distributed_map_item_handler", "durable_execution", "durable_parallel_branch", "durable_step", 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 aa5a2b320..29bf1a321 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,17 +6,63 @@ 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 +class DistributedMapStatus(Enum): + SUCCEEDED = "SUCCEEDED" + FAILED = "FAILED" + STOPPED = "STOPPED" + TIMED_OUT = "TIMED_OUT" + + +class DistributedMapItemStatus(Enum): + """Terminal status of a single map run item.""" + + SUCCEEDED = "SUCCEEDED" + FAILED = "FAILED" + UNKNOWN_TO_SDK_VERSION = "UNKNOWN_TO_SDK_VERSION" + + +class DistributedMapSourceFormat(Enum): + JSON_LINES = "JSON_LINES" + JSON_ARRAY = "JSON_ARRAY" + CSV = "CSV" + + +class DistributedMapS3SourceType(Enum): + """What an S3 distributed map source reads.""" + + OBJECT = "OBJECT" + LISTED_OBJECT_METADATA = "LISTED_OBJECT_METADATA" + LISTED_OBJECTS = "LISTED_OBJECTS" + + +class ProcessorResponseMode(Enum): + """How a processor function reports item outcomes.""" + + BATCH = "BATCH" + ITEM_OUTCOMES = "ITEM_OUTCOMES" + + +class DistributedMapCsvDelimiter(Enum): + """Column delimiter for a CSV distributed map source.""" + + COMMA = "COMMA" + PIPE = "PIPE" + SEMICOLON = "SEMICOLON" + SPACE = "SPACE" + TAB = "TAB" + + P = TypeVar("P") # Payload type R = TypeVar("R") # Result type 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 @@ -490,6 +536,668 @@ class StepConfig: serdes: SerDes | None = None +# region map run configuration + + +@dataclass(frozen=True) +class S3Uri: + """A parsed ``s3://bucket/path`` URI.""" + + bucket: str + path: str | None = None + + @classmethod + def parse(cls, uri: str) -> S3Uri: + """Parse an ``s3://bucket/path`` URI, rejecting any other scheme.""" + if not uri.startswith("s3://"): + msg = f"S3 URI must start with s3://, got: {uri}" + raise ValidationError(msg) + bucket, _, path = uri.removeprefix("s3://").partition("/") + if not bucket: + msg = f"Invalid S3 URI: {uri}" + raise ValidationError(msg) + return cls(bucket=bucket, path=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.isascii() and value.isdigit()) # noqa: PLR2004 + ): + msg = f"expected_bucket_owner must be a 12-digit account id, got: {value}" + raise ValidationError(msg) + + +_MAX_CSV_COLUMNS = 1000 +_MAX_CSV_COLUMN_LENGTH = 256 +_MAX_ITEMS_TO_READ = 100_000_000 + + +def _validate_columns(name: str, columns: list[str] | None) -> None: + """Validate a CSV columns/headers list against its count, uniqueness and length limits.""" + if columns is None: + return + if len(columns) == 0: + msg = f"{name} must be non-empty" + raise ValidationError(msg) + if len(columns) > _MAX_CSV_COLUMNS: + msg = f"{name} must hold at most {_MAX_CSV_COLUMNS} names, got: {len(columns)}" + raise ValidationError(msg) + if len(set(columns)) != len(columns): + msg = f"{name} must not contain duplicates" + raise ValidationError(msg) + for column in columns: + if len(column) > _MAX_CSV_COLUMN_LENGTH: + msg = ( + f"each {name} name must be at most {_MAX_CSV_COLUMN_LENGTH} " + f"characters, got: {len(column)}" + ) + raise ValidationError(msg) + + +# NamespacedFunctionName max length; the service enforces the full grammar. +_MAX_FUNCTION_NAME_LENGTH = 170 + + +def _validate_function_name(name: str) -> None: + """Validate a Lambda function reference is present and within the length limit.""" + if not name: + msg = "function name must be non-empty" + raise ValidationError(msg) + if len(name) > _MAX_FUNCTION_NAME_LENGTH: + msg = f"function name must be at most {_MAX_FUNCTION_NAME_LENGTH} characters" + raise ValidationError(msg) + + +@dataclass(frozen=True) +class DistributedMapCompletionConfig: + """Failure-tolerance configuration for a map run.""" + + tolerated_failure_count: int | None = None + tolerated_failure_percentage: int | 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.tolerated_failure_count is None + and self.tolerated_failure_percentage is None + ): + # Omitting the config entirely means zero tolerated failures, so an + # empty one is rejected rather than resolved to the same thing. + msg = ( + "a completion config must set tolerated_failure_count or " + "tolerated_failure_percentage" + ) + 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 isinstance( + self.tolerated_failure_percentage, int + ): + msg = ( + "tolerated_failure_percentage must be an integer, got: " + f"{self.tolerated_failure_percentage}" + ) + 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: int, *, 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 DistributedMapProcessor: + """Processor configuration for a map run.""" + + UNLIMITED: ClassVar[str] = "unlimited" + + function_name: str + response_mode: ProcessorResponseMode = ProcessorResponseMode.BATCH + batch_size: int | None = None + max_retry_attempts: int | Literal["unlimited"] | None = None + max_retry_duration: Duration | 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) + attempts = self.max_retry_attempts + # A value that is neither an int nor UNLIMITED carries no retry limit to the + # wire, so the service would apply its own default instead of this value. + # Rejecting it here is what makes the omission impossible. + if ( + attempts is not None + and attempts != self.UNLIMITED + and (not isinstance(attempts, int) or attempts < 0) + ): + msg = ( + "max_retry_attempts must be a non-negative integer or " + "DistributedMapProcessor.UNLIMITED, " + f"got: {attempts!r}" + ) + 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) + if self.durable_execution_name_prefix is not None and not ( + 1 <= len(self.durable_execution_name_prefix) <= 36 # noqa: PLR2004 + ): + msg = ( + "durable_execution_name_prefix must be 1 to 36 characters, got: " + f"{len(self.durable_execution_name_prefix)}" + ) + raise ValidationError(msg) + + @classmethod + def report_batch_outcome( + cls, + name: str, + *, + batch_size: int | None = None, + max_retry_attempts: int | Literal["unlimited"] | None = None, + max_retry_duration: Duration | None = None, + durable_execution_name_prefix: str | None = None, + ) -> DistributedMapProcessor: + """Processor that reports one outcome for the whole batch. Items have no outputs.""" + return cls( + function_name=name, + response_mode=ProcessorResponseMode.BATCH, + batch_size=batch_size, + max_retry_attempts=max_retry_attempts, + max_retry_duration=max_retry_duration, + durable_execution_name_prefix=durable_execution_name_prefix, + ) + + @classmethod + def report_item_outcomes( + cls, + name: str, + *, + batch_size: int | None = None, + max_retry_attempts: int | Literal["unlimited"] | None = None, + max_retry_duration: Duration | None = None, + durable_execution_name_prefix: str | None = None, + ) -> DistributedMapProcessor: + """Processor that may report a failure or an output per item. Unlisted items succeed.""" + return cls( + function_name=name, + response_mode=ProcessorResponseMode.ITEM_OUTCOMES, + batch_size=batch_size, + max_retry_attempts=max_retry_attempts, + max_retry_duration=max_retry_duration, + durable_execution_name_prefix=durable_execution_name_prefix, + ) + + +def _parse_delimiter( + value: str | DistributedMapCsvDelimiter, +) -> DistributedMapCsvDelimiter: + """Normalize a delimiter passed as a string or enum member.""" + if isinstance(value, DistributedMapCsvDelimiter): + return value + try: + return DistributedMapCsvDelimiter(value) + except ValueError: + allowed = ", ".join(d.value for d in DistributedMapCsvDelimiter) + msg = f"delimiter must be one of ({allowed}), got: {value!r}" + raise ValidationError(msg) from None + + +@dataclass(frozen=True) +class S3Source: + """Resolved S3 source.""" + + source_type: DistributedMapS3SourceType + bucket: str + key: str | None = None + key_prefix: str | None = None + fmt: DistributedMapSourceFormat | None = None + delimiter: DistributedMapCsvDelimiter | None = None + headers: list[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) + if self.source_type is DistributedMapS3SourceType.OBJECT: + if not self.key: + msg = "an OBJECT source requires a key" + raise ValidationError(msg) + if self.fmt is None: + msg = "an OBJECT source requires a format" + raise ValidationError(msg) + else: + if self.key_prefix is None: + msg = f"a {self.source_type.value} source requires a key prefix" + raise ValidationError(msg) + if ( + self.source_type is DistributedMapS3SourceType.LISTED_OBJECTS + and self.fmt is None + ): + msg = "a LISTED_OBJECTS source requires a format" + raise ValidationError(msg) + if ( + self.source_type is DistributedMapS3SourceType.LISTED_OBJECT_METADATA + and self.fmt is not None + ): + msg = "a LISTED_OBJECT_METADATA source reads no contents, so it takes no format" + raise ValidationError(msg) + + +@dataclass(frozen=True) +class ReaderSource: + """Resolved reader-function source.""" + + 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.""" + + max_items: int | None = None + inline_items: list[Any] | None = None + inline_serdes: SerDes | None = None # None = DEFAULT_JSON_SERDES + s3: S3Source | None = None + reader: ReaderSource | None = None + + def __post_init__(self) -> None: + if self.max_items is not None and not ( + 1 <= self.max_items <= _MAX_ITEMS_TO_READ + ): + msg = ( + f"max_items must be between 1 and {_MAX_ITEMS_TO_READ}, " + f"got: {self.max_items}" + ) + raise ValidationError(msg) + populated = sum( + 1 + for value in (self.inline_items, self.s3, self.reader) + if value is not None + ) + if populated != 1: + msg = "exactly one of inline_items, s3 or reader must be set" + 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( + inline_items=list(items), + inline_serdes=serdes, + max_items=max_items, + ) + + @classmethod + def s3_json_lines( + cls, + uri: str, + *, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read a single object, treating each line as an item.""" + parsed_uri = S3Uri.parse(uri) + bucket, key = parsed_uri.bucket, parsed_uri.path + if key is None: + msg = "s3_json_lines requires an S3 object key" + raise ValidationError(msg) + return cls( + max_items=max_items, + s3=S3Source( + source_type=DistributedMapS3SourceType.OBJECT, + bucket=bucket, + key=key, + fmt=DistributedMapSourceFormat.JSON_LINES, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @classmethod + def s3_json_array( + cls, + uri: str, + *, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read a single object holding a JSON array, treating each element as an item.""" + parsed_uri = S3Uri.parse(uri) + bucket, key = parsed_uri.bucket, parsed_uri.path + if key is None: + msg = "s3_json_array requires an S3 object key" + raise ValidationError(msg) + return cls( + max_items=max_items, + s3=S3Source( + source_type=DistributedMapS3SourceType.OBJECT, + bucket=bucket, + key=key, + fmt=DistributedMapSourceFormat.JSON_ARRAY, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @classmethod + def s3_csv( + cls, + uri: str, + *, + headers: Sequence[str] | None = None, + delimiter: str | DistributedMapCsvDelimiter = DistributedMapCsvDelimiter.COMMA, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read a single object, treating each record as an item.""" + parsed_uri = S3Uri.parse(uri) + bucket, key = parsed_uri.bucket, parsed_uri.path + if key is None: + msg = "s3_csv requires an S3 object key" + raise ValidationError(msg) + return cls( + max_items=max_items, + s3=S3Source( + source_type=DistributedMapS3SourceType.OBJECT, + bucket=bucket, + key=key, + fmt=DistributedMapSourceFormat.CSV, + delimiter=_parse_delimiter(delimiter), + headers=list(headers) if headers is not None else None, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @classmethod + def s3_objects( + cls, + prefix_uri: str, + *, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read each object under a prefix as one item, without reading its contents.""" + parsed_uri = S3Uri.parse(prefix_uri) + bucket, prefix = parsed_uri.bucket, parsed_uri.path + return cls( + max_items=max_items, + s3=S3Source( + source_type=DistributedMapS3SourceType.LISTED_OBJECT_METADATA, + bucket=bucket, + key_prefix=prefix or "", + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @classmethod + def s3_flattened_json_lines( + cls, + prefix_uri: str, + *, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read every object under a prefix, treating each line as an item.""" + parsed_uri = S3Uri.parse(prefix_uri) + bucket, prefix = parsed_uri.bucket, parsed_uri.path + return cls( + max_items=max_items, + s3=S3Source( + source_type=DistributedMapS3SourceType.LISTED_OBJECTS, + bucket=bucket, + key_prefix=prefix or "", + fmt=DistributedMapSourceFormat.JSON_LINES, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @classmethod + def s3_flattened_json_array( + cls, + prefix_uri: str, + *, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read every object under a prefix, treating each array element as an item.""" + parsed_uri = S3Uri.parse(prefix_uri) + bucket, prefix = parsed_uri.bucket, parsed_uri.path + return cls( + max_items=max_items, + s3=S3Source( + source_type=DistributedMapS3SourceType.LISTED_OBJECTS, + bucket=bucket, + key_prefix=prefix or "", + fmt=DistributedMapSourceFormat.JSON_ARRAY, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @classmethod + def s3_flattened_csv( + cls, + prefix_uri: str, + *, + headers: Sequence[str] | None = None, + delimiter: str | DistributedMapCsvDelimiter = DistributedMapCsvDelimiter.COMMA, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read every object under a prefix, treating each record as an item.""" + parsed_uri = S3Uri.parse(prefix_uri) + bucket, prefix = parsed_uri.bucket, parsed_uri.path + return cls( + max_items=max_items, + s3=S3Source( + source_type=DistributedMapS3SourceType.LISTED_OBJECTS, + bucket=bucket, + key_prefix=prefix or "", + fmt=DistributedMapSourceFormat.CSV, + delimiter=_parse_delimiter(delimiter), + headers=list(headers) if headers is not None else None, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @classmethod + def reader_function( + cls, + 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 cls( + max_items=max_items, + reader=ReaderSource( + 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) + + @classmethod + def from_uri( + cls, + 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.""" + parsed_uri = S3Uri.parse(prefix_uri) + return cls( + bucket=parsed_uri.bucket, + prefix=parsed_uri.path or "", + include_input=include_input, + include_output=include_output, + expected_bucket_owner=expected_bucket_owner, + ) + + +@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) + + @classmethod + def from_uri( + cls, + 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.""" + parsed_uri = S3Uri.parse(prefix_uri) + return cls( + bucket=parsed_uri.bucket, + prefix=parsed_uri.path or "", + include_input=include_input, + include_error=include_error, + expected_bucket_owner=expected_bucket_owner, + ) + + +@dataclass(frozen=True) +class DistributedMapDestination: + """Destination routing for map run results.""" + + on_success: SuccessDestination | None = None + on_failure: FailureDestination | None = None + + +@dataclass(frozen=True) +class DistributedMapConfig: + """Configuration for map run operations.""" + + destination: DistributedMapDestination | None = None + completion_config: DistributedMapCompletionConfig | None = None + timeout: Duration | None = None + + 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) + + +@dataclass(frozen=True) +class DistributedMapResultConfig(DistributedMapConfig): + """Configuration that collects per-item results inline. + + Passing this instead of ``DistributedMapConfig`` makes ``ctx.distributed_map`` + return a ``DistributedMapResult``, which carries the per-item outcomes. + """ + + result_serdes: SerDes | None = None # None = DEFAULT_JSON_SERDES + + +# 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 16f5cced6..5f886832b 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 @@ -12,6 +12,7 @@ NoReturn, ParamSpec, TypeVar, + overload, ) from aws_durable_execution_sdk_python.config import ( @@ -20,6 +21,10 @@ Duration, InvokeConfig, MapConfig, + DistributedMapConfig, + DistributedMapProcessor, + DistributedMapResultConfig, + DistributedMapSource, ParallelBranch, ParallelConfig, StepConfig, @@ -28,6 +33,10 @@ from aws_durable_execution_sdk_python.concurrency.models import ( envelope_summary_generator, ) +from aws_durable_execution_sdk_python.dmap.models import ( + DistributedMapResult, + DistributedMapSummary, +) from aws_durable_execution_sdk_python.exceptions import ( CallbackError, CallbackExternalError, @@ -54,6 +63,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.dmap 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 @@ -683,6 +695,72 @@ def invoke( ) return executor.process() + @overload + def distributed_map( + self, + source: DistributedMapSource | Sequence[Any], + processor: DistributedMapProcessor, + max_concurrency: int, + name: str | None = None, + *, + config: DistributedMapResultConfig, + ) -> DistributedMapResult: ... + + @overload + def distributed_map( + self, + source: DistributedMapSource | Sequence[Any], + processor: DistributedMapProcessor, + max_concurrency: int, + name: str | None = None, + *, + config: DistributedMapConfig | None = None, + ) -> DistributedMapSummary: ... + + def distributed_map( + self, + source: DistributedMapSource | Sequence[Any], + processor: DistributedMapProcessor, + max_concurrency: int, + name: str | None = None, + *, + config: DistributedMapConfig | None = None, + ) -> DistributedMapSummary: + """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. Required, + with no service default. + name: Optional name for the operation + config: Optional run-level configuration + + Returns: + The map run's summary, or a DistributedMapResult when a + DistributedMapResultConfig is passed + """ + 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 < 1: + msg = f"max_concurrency must be greater than zero, got: {max_concurrency}" + raise ValidationError(msg) + if config is None: + config = DistributedMapConfig() + with self._operation_replay_aware( + OperationSubType.DISTRIBUTED_MAP, name + ) as operation_identifier: + executor: DistributedMapOperationExecutor = DistributedMapOperationExecutor( + source=source, + processor=processor, + max_concurrency=max_concurrency, + state=self.state, + operation_identifier=operation_identifier, + 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/dmap/__init__.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dmap/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dmap/handlers.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dmap/handlers.py new file mode 100644 index 000000000..6d14edab5 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dmap/handlers.py @@ -0,0 +1,484 @@ +"""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 functools +from collections.abc import MutableMapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Any, Callable + +from aws_durable_execution_sdk_python.config import ( + BatchItemStatus, + 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.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.serdes import ( + DEFAULT_JSON_SERDES, + SerDes, + 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 + + +@dataclass(frozen=True) +class ProcessorRecord: + """One item of the batch handed to a processor function.""" + + item_id: str + body: Any = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> ProcessorRecord: + return cls(item_id=data.get("itemId", ""), body=data.get("body")) + + def to_dict(self) -> MutableMapping[str, Any]: + return {"itemId": self.item_id, "body": self.body} + + +@dataclass(frozen=True) +class ProcessorEvent: + """The event a processor function is invoked with, one batch of items.""" + + records: list[ProcessorRecord] = field(default_factory=list) + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> ProcessorEvent: + records = data.get("records") + if not isinstance(records, list): + msg = "expected a distributed map processor envelope with a 'records' list" + raise ValidationError(msg) + return cls(records=[ProcessorRecord.from_dict(r) for r in records]) + + def to_dict(self) -> MutableMapping[str, Any]: + return {"records": [record.to_dict() for record in self.records]} + + +@dataclass(frozen=True) +class ItemResult: + """One item's output, reported by an item_outcomes processor.""" + + item_identifier: str + output: str | None = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> ItemResult: + return cls( + item_identifier=data.get("itemIdentifier", ""), output=data.get("output") + ) + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {"itemIdentifier": self.item_identifier} + if self.output is not None: + result["output"] = self.output + return result + + +@dataclass(frozen=True) +class ItemFailure: + """One item's failure, reported by an item_outcomes processor.""" + + item_identifier: str + error: ErrorObject | None = None + + @classmethod + def from_exception(cls, item_identifier: str, exc: Exception) -> ItemFailure: + return cls( + item_identifier=item_identifier, error=ErrorObject.from_exception(exc) + ) + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> ItemFailure: + error = data.get("error") + return cls( + item_identifier=data.get("itemIdentifier", ""), + error=ErrorObject( + message=error.get("errorMessage"), + type=error.get("errorType"), + data=error.get("errorData"), + stack_trace=error.get("stackTrace"), + ) + if error + else None, + ) + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {"itemIdentifier": self.item_identifier} + if self.error is None: + return result + # The processor response envelope is camelCase, so ErrorObject.to_dict + # cannot be reused here. Every member is optional and the service + # stores each one unchanged, so an absent member is omitted. + error: MutableMapping[str, Any] = {} + if self.error.type is not None: + error["errorType"] = self.error.type + if self.error.message is not None: + error["errorMessage"] = self.error.message + if self.error.data is not None: + error["errorData"] = self.error.data + if self.error.stack_trace is not None: + error["stackTrace"] = self.error.stack_trace + if error: + result["error"] = error + return result + + +@dataclass(frozen=True) +class ItemHandlerResponse: + """What an item handler returns: a failure or a result for each item.""" + + failures: list[ItemFailure] = field(default_factory=list) + results: list[ItemResult] = field(default_factory=list) + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> ItemHandlerResponse: + return cls( + failures=[ + ItemFailure.from_dict(f) for f in data.get("batchItemFailures", []) + ], + results=[ItemResult.from_dict(r) for r in data.get("batchItemResults", [])], + ) + + def to_dict(self) -> MutableMapping[str, Any]: + return { + "batchItemFailures": [f.to_dict() for f in self.failures], + "batchItemResults": [r.to_dict() for r in self.results], + } + + +@dataclass(frozen=True) +class ReaderEvent: + """The event a reader function is invoked with.""" + + max_items_per_page: int + state: str | None = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> ReaderEvent: + max_items_per_page = data.get("maxItemsPerPage") + if not isinstance(max_items_per_page, int): + msg = ( + "expected a distributed map reader envelope with an integer " + "'maxItemsPerPage'" + ) + raise ValidationError(msg) + return cls(max_items_per_page=max_items_per_page, state=data.get("state")) + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {"maxItemsPerPage": self.max_items_per_page} + if self.state is not None: + result["state"] = self.state + return result + + +@dataclass(frozen=True) +class ReaderResponse: + """What a reader function returns. A None ``next_state`` exhausts the source.""" + + items: list[str] = field(default_factory=list) + next_state: str | None = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> ReaderResponse: + return cls(items=list(data.get("items", [])), next_state=data.get("nextState")) + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {"items": self.items} + if self.next_state is not None: + result["nextState"] = self.next_state + return result + + +def _to_item(serdes: SerDes[Any], body: str, ctx: SerDesContext) -> Any: + """Recover a typed item from a record body.""" + return serdes.deserialize(body, ctx) + + +def _to_output(serdes: SerDes[Any], value: Any, ctx: SerDesContext) -> str | None: + """Serialize an item's output, or None when the handler returned nothing.""" + if value is None: + return None + return serdes.serialize(value, ctx) + + +def _validate_concurrency(concurrency: int) -> None: + if concurrency < 1: + msg = f"concurrency must be at least 1, got: {concurrency}" + raise ValidationError(msg) + + +def distributed_map_item_handler( + func: Callable[[Any], Any] | None = None, + *, + item_serdes: SerDes[Any] | None = None, + result_serdes: SerDes[Any] | None = None, + concurrency: int = 1, +) -> Callable[..., Any]: + """Wrap a Lambda to be used as an item_outcomes processor. + + ``func`` takes one item and returns its output or raises. The decorated name + becomes the Lambda handler and takes ``(event, context)``, so name it + ``handler``. An item whose handler returns ``None`` succeeds with no output. + + Pair this with ``DistributedMapProcessor.report_item_outcomes``. Paired with + ``report_batch_outcome`` instead, the service ignores the response this + returns, so every item in the batch is marked succeeded, including the items + whose handler raised. + + Items in a batch are processed one at a time. Raise ``concurrency`` to run + that many items at once, which requires ``func`` to be safe to call from + several threads. + """ + if func is None: + return functools.partial( + distributed_map_item_handler, + item_serdes=item_serdes, + result_serdes=result_serdes, + concurrency=concurrency, + ) + _validate_concurrency(concurrency) + serdes = item_serdes or DEFAULT_JSON_SERDES + out_serdes = result_serdes or DEFAULT_JSON_SERDES + + def handler(event: dict[str, Any], _context: Any = None) -> dict[str, Any]: + ctx = SerDesContext() + records = ProcessorEvent.from_dict(event).records + + def run(record: ProcessorRecord) -> str | None: + # Serializing here, not after the loop, keeps a serialization + # failure on the one item whose output could not be serialized. + return _to_output(out_serdes, func(_to_item(serdes, record.body, ctx)), ctx) + + outputs: list[str | None] = [None] * len(records) + errors: list[Exception | None] = [None] * len(records) + with ThreadPoolExecutor(max_workers=concurrency) 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[ItemResult] = [] + failures: list[ItemFailure] = [] + for i, record in enumerate(records): + err = errors[i] + if err is not None: + failures.append(ItemFailure.from_exception(record.item_id, err)) + else: + results.append( + ItemResult(item_identifier=record.item_id, output=outputs[i]) + ) + + response = ItemHandlerResponse(failures=failures, results=results) + return dict(response.to_dict()) + + return handler + + +def distributed_map_batch_handler( + func: Callable[[list[Any]], None] | None = None, + *, + item_serdes: SerDes[Any] | None = None, +) -> Callable[..., Any]: + """Wrap a Lambda to be used as a batch processor. + + ``func`` takes the whole batch of items. Returning succeeds every item; + raising fails every item. The decorated name becomes the Lambda handler and + takes ``(event, context)``, so name it ``handler``. + + A batch processor reports no per-item outcome, so it records no output. Use + ``distributed_map_item_handler`` to report an output per item. + """ + if func is None: + return functools.partial(distributed_map_batch_handler, item_serdes=item_serdes) + serdes = item_serdes or DEFAULT_JSON_SERDES + + def handler(event: dict[str, Any], _context: Any = None) -> None: + ctx = SerDesContext() + func( + [ + _to_item(serdes, record.body, ctx) + for record in ProcessorEvent.from_dict(event).records + ] + ) + + return handler + + +def distributed_map_reader( + func: Callable[[Any, int], ReaderPage] | None = None, + *, + item_serdes: SerDes[Any] | None = None, + state_serdes: SerDes[Any] | None = None, +) -> Callable[..., Any]: + """Wrap a Lambda to be used as a reader source. + + ``func`` takes the current state and the maximum items this page may + return, and returns a ReaderPage. A ``next_state`` of ``None`` signals the + source is exhausted. The decorated name becomes the Lambda handler and + takes ``(event, context)``, so name it ``handler``. + + ``item_serdes`` serializes each item of the page. It must be the same serdes + the processor is configured with, since the processor deserializes every + item body with its own ``item_serdes``. + """ + if func is None: + return functools.partial( + distributed_map_reader, + item_serdes=item_serdes, + state_serdes=state_serdes, + ) + item_sd = item_serdes or DEFAULT_JSON_SERDES + state_sd = state_serdes or DEFAULT_JSON_SERDES + + def handler(event: dict[str, Any], _context: Any = None) -> dict[str, Any]: + ctx = SerDesContext() + reader_event = ReaderEvent.from_dict(event) + state = ( + state_sd.deserialize(reader_event.state, ctx) + if reader_event.state is not None + else None + ) + + page = func(state, reader_event.max_items_per_page) + if len(page.items) > reader_event.max_items_per_page: + msg = ( + f"reader returned {len(page.items)} items, exceeding " + f"maxItemsPerPage {reader_event.max_items_per_page}" + ) + raise ValidationError(msg) + + items = [item_sd.serialize(item, ctx) for item in page.items] + + next_state: str | None = None + if page.next_state is not None: + next_state = state_sd.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 ValidationError(msg) + + return dict(ReaderResponse(items=items, next_state=next_state).to_dict()) + + return handler + + +def durable_distributed_map_item_handler( + func: Callable[..., Any] | None = None, + *, + item_serdes: SerDes[Any] | None = None, + result_serdes: SerDes[Any] | None = None, + concurrency: int = 1, +) -> Callable[..., Any]: + """Durable variant of the item handler. ``func`` receives (context, item). + + The Lambda must be deployed as a durable function. The decorated name becomes + the Lambda handler and takes ``(event, context)``, so name it ``handler``. + An item whose handler returns ``None`` succeeds with no output. + + Pair this with ``DistributedMapProcessor.report_item_outcomes``. Paired with + ``report_batch_outcome`` instead, the service ignores the response this + returns, so every item in the batch is marked succeeded, including the items + whose handler raised. + + Items in a batch are processed one at a time. Raise ``concurrency`` to run + that many items at once. + """ + if func is None: + return functools.partial( + durable_distributed_map_item_handler, + item_serdes=item_serdes, + result_serdes=result_serdes, + concurrency=concurrency, + ) + _validate_concurrency(concurrency) + serdes = item_serdes or DEFAULT_JSON_SERDES + out_serdes = result_serdes or DEFAULT_JSON_SERDES + + @durable_execution + def handler(event: dict[str, Any], context: Any) -> dict[str, Any]: + ctx = SerDesContext() + records = ProcessorEvent.from_dict(event).records + + def per_item(inner_ctx: Any, body: Any, _index: int, _inputs: Any) -> Any: + return func(inner_ctx, _to_item(serdes, body, ctx)) + + batch = context.map( + [record.body for record in records], + per_item, + config=MapConfig( + max_concurrency=concurrency, + completion_config=CompletionConfig.all_completed(), + ), + ) + + results: list[ItemResult] = [] + failures: list[ItemFailure] = [] + for bi in batch.all: + item_id = records[bi.index].item_id + if bi.status is BatchItemStatus.SUCCEEDED: + # Catching here, not around the loop, keeps a serialization + # failure on the one item whose output could not be serialized. + try: + output = _to_output(out_serdes, bi.result, ctx) + except Exception as exc: # noqa: BLE001 + failures.append(ItemFailure.from_exception(item_id, exc)) + else: + results.append(ItemResult(item_identifier=item_id, output=output)) + else: + failures.append(ItemFailure(item_identifier=item_id, error=bi.error)) + + response = ItemHandlerResponse(failures=failures, results=results) + return dict(response.to_dict()) + + return handler + + +def durable_distributed_map_batch_handler( + func: Callable[..., None] | None = None, + *, + item_serdes: SerDes[Any] | None = None, +) -> Callable[..., Any]: + """Durable variant of the batch handler. ``func`` receives (context, items). + + The Lambda must be deployed as a durable function. The decorated name becomes + the Lambda handler and takes ``(event, context)``, so name it ``handler``. + + A batch processor reports no per-item outcome, so it records no output. Use + ``durable_distributed_map_item_handler`` to report an output per item. + """ + if func is None: + return functools.partial( + durable_distributed_map_batch_handler, item_serdes=item_serdes + ) + serdes = item_serdes or DEFAULT_JSON_SERDES + + @durable_execution + def handler(event: dict[str, Any], context: Any) -> None: + ctx = SerDesContext() + func( + context, + [ + _to_item(serdes, record.body, ctx) + for record in ProcessorEvent.from_dict(event).records + ], + ) + + return handler diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dmap/models.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dmap/models.py new file mode 100644 index 000000000..b39919726 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dmap/models.py @@ -0,0 +1,250 @@ +"""Result models returned by ``ctx.distributed_map``.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Self, TypedDict + +from aws_durable_execution_sdk_python.exceptions import ( + DistributedMapError, + ExecutionError, +) +from aws_durable_execution_sdk_python.config import ( + DistributedMapItemStatus, + DistributedMapStatus, +) +from aws_durable_execution_sdk_python.lambda_service import ErrorObject, OperationStatus +from aws_durable_execution_sdk_python.types import DurableExecutionArn + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python.lambda_service import Operation + +_STATUS_FOR_OPERATION_STATUS = { + OperationStatus.SUCCEEDED: DistributedMapStatus.SUCCEEDED, + OperationStatus.FAILED: DistributedMapStatus.FAILED, + OperationStatus.STOPPED: DistributedMapStatus.STOPPED, + OperationStatus.TIMED_OUT: DistributedMapStatus.TIMED_OUT, +} + + +def _to_distributed_map_status(operation: Operation | None) -> DistributedMapStatus: + """Derive the map run status from the operation's terminal status.""" + op_status = operation.status if operation else None + resolved = ( + _STATUS_FOR_OPERATION_STATUS.get(op_status) if op_status is not None else None + ) + if resolved is None: + msg = ( + "Cannot derive distributed map status from operation status " + f"{op_status.value if op_status else 'UNKNOWN'}" + ) + raise ExecutionError(msg) + return resolved + + +class _SummaryFields(TypedDict): + """The summary fields both factories pass to the constructor.""" + + status: DistributedMapStatus + success_count: int + failure_count: int + unprocessed_count: int + distributed_map_run_arn: str | None + source_exhausted: bool | None + error: ErrorObject | None + total_count: int | None + + +def _summary_fields(operation: Operation | None) -> _SummaryFields: + """Read the summary fields a terminal operation must carry.""" + details = operation.distributed_map_details if operation else None + if details is None: + # The raw operation status, not _to_distributed_map_status, which raises + # on a non-terminal status and would replace this error with its own. + op_status = ( + operation.status.value if operation and operation.status else "UNKNOWN" + ) + msg = ( + f"DISTRIBUTED_MAP operation ended {op_status} but carried no " + f"DistributedMapDetails" + ) + raise ExecutionError(msg) + status = _to_distributed_map_status(operation) + # SourceExhausted is required at terminal once a run exists, and is absent + # only when the operation failed before one was created. + if details.distributed_map_run_arn is not None and details.source_exhausted is None: + msg = ( + f"DISTRIBUTED_MAP operation ended {status.value} but carried no " + f"SourceExhausted" + ) + raise ExecutionError(msg) + return { + "status": status, + "success_count": details.success_count, + "failure_count": details.failure_count, + "unprocessed_count": details.unprocessed_count, + "distributed_map_run_arn": details.distributed_map_run_arn, + "source_exhausted": details.source_exhausted, + "error": details.error, + "total_count": details.total_count, + } + + +_RUN_ARN_SEPARATOR = "/distributed-map-run/" + + +@dataclass(frozen=True) +class DistributedMapRunArn: + """Parsed components of a distributed map run ARN. + + Attributes: + durable_execution_arn: The parent durable execution's ARN. + run_id: The map run id, unique per account. + """ + + durable_execution_arn: str + run_id: str + + @classmethod + def from_arn(cls, arn: str) -> DistributedMapRunArn | None: + """Parse a map run ARN, returning ``None`` when it does not match the format.""" + parent, _, run_id = arn.partition(_RUN_ARN_SEPARATOR) + if not run_id or "/" in run_id or DurableExecutionArn.from_arn(parent) is None: + return None + return cls(durable_execution_arn=parent, run_id=run_id) + + +@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 rather than raising. Use :meth:`throw_if_error` to opt into raising. + """ + + status: DistributedMapStatus + success_count: int + failure_count: int + unprocessed_count: int + distributed_map_run_arn: str | None = None + source_exhausted: bool | None = None + error: ErrorObject | None = None + total_count: int | None = None + + @classmethod + def from_operation(cls, operation: Operation | None) -> Self: + """Rebuild the summary from a terminal DISTRIBUTED_MAP operation. + + Called on :class:`DistributedMapResult` this returns a result whose + ``all`` is empty, which is the correct value for a run with no items. + """ + return cls(**_summary_fields(operation)) + + @property + def distributed_map_id(self) -> str | None: + """Run id from :attr:`distributed_map_run_arn`, ``None`` if absent or unparseable.""" + if self.distributed_map_run_arn is None: + return None + parsed = DistributedMapRunArn.from_arn(self.distributed_map_run_arn) + return parsed.run_id if parsed else None + + @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: + error = self.error + msg = f"Map run ended {self.status.value}" + if error is not None and error.message: + msg = f"{msg}: {error.message}" + raise DistributedMapError( + message=msg, + error_type=error.type if error else None, + data=error.data if error else None, + stack_trace=error.stack_trace if error else None, + ) + 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 DistributedMapResultItem: + """Outcome of a single map run item.""" + + item_id: str + status: DistributedMapItemStatus + output: Any | None = None + error: ErrorObject | None = None + + +@dataclass(frozen=True) +class DistributedMapResult(DistributedMapSummary): + """Outcome of a map run with per-item results. + + Resolved by ``ctx.distributed_map`` when passed a + :class:`DistributedMapResultConfig`. Extends :class:`DistributedMapSummary` + with the retained per-item results. + """ + + all: list[DistributedMapResultItem] = field(default_factory=list) + + @classmethod + def from_operation_and_items( + cls, operation: Operation | None, items: list[DistributedMapResultItem] + ) -> Self: + """Rebuild the result from a terminal operation and its deserialized items.""" + return cls(**_summary_fields(operation), all=items) + + def succeeded(self) -> list[DistributedMapResultItem]: + """Return the items that succeeded.""" + return [ + item + for item in self.all + if item.status is DistributedMapItemStatus.SUCCEEDED + ] + + def failed(self) -> list[DistributedMapResultItem]: + """Return the items that permanently failed.""" + return [ + item for item in self.all if item.status is DistributedMapItemStatus.FAILED + ] + + def get_results(self) -> list[Any]: + """Return one output per succeeded item, None where none was recorded.""" + return [ + item.output + for item in self.all + if item.status is DistributedMapItemStatus.SUCCEEDED + ] + + def get_errors(self) -> list[ErrorObject | None]: + """Return one error per failed item, None where none was recorded.""" + return [ + item.error + for item in self.all + if item.status is DistributedMapItemStatus.FAILED + ] + + 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] + error = first.error + if error is not None: + raise DistributedMapError( + message=f"{error.type}: {error.message}", + error_type=error.type, + data=error.data, + stack_trace=error.stack_trace, + ) + msg = f"item {first.item_id} failed" + raise DistributedMapError(msg) + super().throw_if_error() 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 ea3f163f2..b2123d7af 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 @@ -330,6 +330,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 eb5ee78b8..4b793b3f2 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 @@ -7,15 +7,40 @@ from collections.abc import MutableMapping from dataclasses import dataclass, field from enum import Enum -from typing import TYPE_CHECKING, Any, NoReturn, Protocol, TypeAlias, cast +from typing import ( + TYPE_CHECKING, + Any, + NoReturn, + Protocol, + Self, + TypeAlias, + TypeVar, + cast, +) import boto3 from botocore.config import Config from aws_durable_execution_sdk_python.__about__ import __version__ +from aws_durable_execution_sdk_python.config import ( + DistributedMapCompletionConfig as ConfigCompletionConfig, +) +from aws_durable_execution_sdk_python.config import ( + DistributedMapCsvDelimiter, + DistributedMapDestination, + DistributedMapItemStatus, + DistributedMapProcessor, + DistributedMapS3SourceType, + DistributedMapSourceFormat, + FailureDestination, + ProcessorResponseMode, + S3Source, + SuccessDestination, +) from aws_durable_execution_sdk_python.exceptions import ( CheckpointError, DurableOperationError, + ExecutionError, GetExecutionStateError, SerDesError, ) @@ -74,6 +99,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 +112,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 @@ -104,6 +132,36 @@ class CallbackTimeoutType(Enum): HEARTBEAT = "Callback.Heartbeat" +class DistributedMapSourceType(Enum): + INLINE = "INLINE" + S3 = "S3" + READER_FUNCTION = "READER_FUNCTION" + + +class DistributedMapCsvHeaderLocation(Enum): + FIRST_ROW = "FIRST_ROW" + GIVEN = "GIVEN" + + +class DistributedMapFunctionResponseType(Enum): + REPORT_BATCH_ITEM_OUTCOMES = "REPORT_BATCH_ITEM_OUTCOMES" + + +class DistributedMapDestinationType(Enum): + S3 = "S3" + + +class DistributedMapDestinationInclude(Enum): + INPUT = "INPUT" + OUTPUT = "OUTPUT" + ERROR = "ERROR" + + +class DistributedMapResultCollectionMode(Enum): + NONE = "NONE" + INLINE = "INLINE" + + class OperationSubType(Enum): STEP = "Step" WAIT = "Wait" @@ -116,6 +174,7 @@ class OperationSubType(Enum): WAIT_FOR_CALLBACK = "WaitForCallback" WAIT_FOR_CONDITION = "WaitForCondition" CHAINED_INVOKE = "ChainedInvoke" + DISTRIBUTED_MAP = "DistributedMap" class InvocationStatus(Enum): @@ -383,6 +442,100 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> ChainedInvokeDetails: ) +_BackendEnumT = TypeVar("_BackendEnumT", bound=Enum) + + +def _parse_enum( + enum_cls: type[_BackendEnumT], + value: str, + field_name: str, + *, + unknown_fallback: _BackendEnumT | None = None, +) -> _BackendEnumT: + """Convert a backend enum string, returning unknown_fallback for an unrecognized value or raising ExecutionError when no fallback is given.""" + try: + return enum_cls(value) + except ValueError as e: + if unknown_fallback is not None: + return unknown_fallback + msg = f"Unknown distributed map {field_name} from the backend: {value!r}" + raise ExecutionError(msg) from e + + +@dataclass(frozen=True) +class DistributedMapResultItem: + """Represent a single map run item's outcome.""" + + item_id: str + status: DistributedMapItemStatus + output: str | None = None + error: ErrorObject | None = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapResultItem: + error_raw = data.get("Error") + status_raw = data.get("Status") + # Status is required, so its absence breaks the contract and must fail. + # An unrecognized value is a status added after this SDK version. + if not status_raw: + msg = ( + f"Distributed map result item {data.get('ItemId', '')!r} carried " + f"no item status" + ) + raise ExecutionError(msg) + return cls( + item_id=data.get("ItemId", ""), + status=_parse_enum( + DistributedMapItemStatus, + status_raw, + "item status", + unknown_fallback=DistributedMapItemStatus.UNKNOWN_TO_SDK_VERSION, + ), + 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.value, + } + 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: + distributed_map_run_arn: str | None = None + source_exhausted: bool | None = None + error: ErrorObject | None = None + total_count: int | None = None + success_count: int = 0 + failure_count: int = 0 + unprocessed_count: int = 0 + results: list[DistributedMapResultItem] | None = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapDetails: + results_raw = data.get("Results") + error_raw = data.get("Error") + return cls( + distributed_map_run_arn=data.get("DistributedMapRunArn"), + source_exhausted=data.get("SourceExhausted"), + error=ErrorObject.from_dict(error_raw) if error_raw else None, + total_count=data.get("TotalCount"), + success_count=data.get("SuccessCount", 0), + failure_count=data.get("FailureCount", 0), + unprocessed_count=data.get("UnprocessedCount", 0), + results=[DistributedMapResultItem.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 +631,688 @@ def to_dict(self) -> MutableMapping[str, Any]: return result +@dataclass(frozen=True) +class DistributedMapCsvFormatOptions: + """Represent CSV format options for an S3 source.""" + + header_location: DistributedMapCsvHeaderLocation + headers: list[str] | None = None + delimiter: DistributedMapCsvDelimiter | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = { + "HeaderLocation": self.header_location.value + } + if self.headers is not None: + result["Headers"] = self.headers + if self.delimiter is not None: + result["Delimiter"] = self.delimiter.value + return result + + @classmethod + def from_dict( + cls, data: MutableMapping[str, Any] + ) -> DistributedMapCsvFormatOptions: + headers = data.get("Headers") + delimiter = data.get("Delimiter") + return cls( + header_location=DistributedMapCsvHeaderLocation( + data.get("HeaderLocation", "FIRST_ROW") + ), + headers=list(headers) if headers is not None else None, + delimiter=DistributedMapCsvDelimiter(delimiter) + if delimiter is not None + else None, + ) + + @classmethod + def from_source(cls, s3: S3Source) -> Self | None: + """Build the CSV options for a CSV source, or None for any other format.""" + if s3.fmt is not DistributedMapSourceFormat.CSV: + return None + return cls( + header_location=( + DistributedMapCsvHeaderLocation.GIVEN + if s3.headers is not None + else DistributedMapCsvHeaderLocation.FIRST_ROW + ), + headers=s3.headers, + delimiter=s3.delimiter, + ) + + +@dataclass(frozen=True) +class DistributedMapS3ObjectConfig: + """Represent the object an OBJECT S3 source reads.""" + + key: str + fmt: DistributedMapSourceFormat + csv_format_options: DistributedMapCsvFormatOptions | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = { + "Key": self.key, + "Format": self.fmt.value, + } + if self.csv_format_options is not None: + result["CsvFormatOptions"] = self.csv_format_options.to_dict() + return result + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapS3ObjectConfig: + csv_raw = data.get("CsvFormatOptions") + return cls( + key=data.get("Key", ""), + fmt=DistributedMapSourceFormat(data["Format"]), + csv_format_options=DistributedMapCsvFormatOptions.from_dict(csv_raw) + if csv_raw is not None + else None, + ) + + +@dataclass(frozen=True) +class DistributedMapS3ListedObjectMetadataConfig: + """Represent the prefix a LISTED_OBJECT_METADATA S3 source lists.""" + + key_prefix: str + + def to_dict(self) -> MutableMapping[str, Any]: + return {"KeyPrefix": self.key_prefix} + + @classmethod + def from_dict( + cls, data: MutableMapping[str, Any] + ) -> DistributedMapS3ListedObjectMetadataConfig: + return cls(key_prefix=data.get("KeyPrefix", "")) + + +@dataclass(frozen=True) +class DistributedMapS3ListedObjectsConfig: + """Represent the prefix a LISTED_OBJECTS S3 source lists and reads.""" + + key_prefix: str + fmt: DistributedMapSourceFormat + csv_format_options: DistributedMapCsvFormatOptions | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = { + "KeyPrefix": self.key_prefix, + "Format": self.fmt.value, + } + if self.csv_format_options is not None: + result["CsvFormatOptions"] = self.csv_format_options.to_dict() + return result + + @classmethod + def from_dict( + cls, data: MutableMapping[str, Any] + ) -> DistributedMapS3ListedObjectsConfig: + csv_raw = data.get("CsvFormatOptions") + return cls( + key_prefix=data.get("KeyPrefix", ""), + fmt=DistributedMapSourceFormat(data["Format"]), + csv_format_options=DistributedMapCsvFormatOptions.from_dict(csv_raw) + if csv_raw is not None + else None, + ) + + +@dataclass(frozen=True) +class DistributedMapS3SourceConfig: + """Represent an S3 distributed map source config. + + ``source_type`` selects which of the three config blocks is present. The + other two are absent. + """ + + source_type: DistributedMapS3SourceType + bucket: str + expected_bucket_owner: str | None = None + object_config: DistributedMapS3ObjectConfig | None = None + listed_object_metadata_config: DistributedMapS3ListedObjectMetadataConfig | None = ( + None + ) + listed_objects_config: DistributedMapS3ListedObjectsConfig | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = { + "Type": self.source_type.value, + "Bucket": self.bucket, + } + if self.expected_bucket_owner is not None: + result["ExpectedBucketOwner"] = self.expected_bucket_owner + if self.object_config is not None: + result["ObjectConfig"] = self.object_config.to_dict() + if self.listed_object_metadata_config is not None: + result["ListedObjectMetadataConfig"] = ( + self.listed_object_metadata_config.to_dict() + ) + if self.listed_objects_config is not None: + result["ListedObjectsConfig"] = self.listed_objects_config.to_dict() + return result + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapS3SourceConfig: + object_raw = data.get("ObjectConfig") + metadata_raw = data.get("ListedObjectMetadataConfig") + listed_raw = data.get("ListedObjectsConfig") + return cls( + source_type=DistributedMapS3SourceType(data["Type"]), + bucket=data.get("Bucket", ""), + expected_bucket_owner=data.get("ExpectedBucketOwner"), + object_config=DistributedMapS3ObjectConfig.from_dict(object_raw) + if object_raw is not None + else None, + listed_object_metadata_config=( + DistributedMapS3ListedObjectMetadataConfig.from_dict(metadata_raw) + if metadata_raw is not None + else None + ), + listed_objects_config=DistributedMapS3ListedObjectsConfig.from_dict( + listed_raw + ) + if listed_raw is not None + else None, + ) + + @classmethod + def from_source(cls, s3: S3Source) -> Self: + """Translate the S3 source into its matching nested config block.""" + csv_options = DistributedMapCsvFormatOptions.from_source(s3) + if s3.source_type is DistributedMapS3SourceType.OBJECT: + # S3Source.__post_init__ requires a key and a format for OBJECT. + return cls( + source_type=s3.source_type, + bucket=s3.bucket, + expected_bucket_owner=s3.expected_bucket_owner, + object_config=DistributedMapS3ObjectConfig( + key=s3.key or "", + fmt=s3.fmt or DistributedMapSourceFormat.JSON_LINES, + csv_format_options=csv_options, + ), + ) + if s3.source_type is DistributedMapS3SourceType.LISTED_OBJECT_METADATA: + return cls( + source_type=s3.source_type, + bucket=s3.bucket, + expected_bucket_owner=s3.expected_bucket_owner, + listed_object_metadata_config=( + DistributedMapS3ListedObjectMetadataConfig( + key_prefix=s3.key_prefix or "" + ) + ), + ) + return cls( + source_type=s3.source_type, + bucket=s3.bucket, + expected_bucket_owner=s3.expected_bucket_owner, + listed_objects_config=DistributedMapS3ListedObjectsConfig( + key_prefix=s3.key_prefix or "", + fmt=s3.fmt or DistributedMapSourceFormat.JSON_LINES, + csv_format_options=csv_options, + ), + ) + + +@dataclass(frozen=True) +class DistributedMapReaderFunctionSourceConfig: + """Represent a reader-function distributed map source config.""" + + function_name: str + initial_state: str | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {"FunctionName": self.function_name} + if self.initial_state is not None: + result["InitialState"] = self.initial_state + return result + + @classmethod + def from_dict( + cls, data: MutableMapping[str, Any] + ) -> DistributedMapReaderFunctionSourceConfig: + return cls( + function_name=data.get("FunctionName", ""), + initial_state=data.get("InitialState"), + ) + + +@dataclass(frozen=True) +class DistributedMapInlineSourceConfig: + """Represent the items an inline map run source reads from.""" + + items: list[str] = field(default_factory=list) + + def to_dict(self) -> MutableMapping[str, Any]: + return {"Items": self.items} + + @classmethod + def from_dict( + cls, data: MutableMapping[str, Any] + ) -> DistributedMapInlineSourceConfig: + return cls(items=list(data.get("Items", []))) + + +@dataclass(frozen=True) +class DistributedMapSourceConfig: + """Represent a map run source.""" + + source_type: DistributedMapSourceType + max_items: int | None = None + inline_source_config: DistributedMapInlineSourceConfig | None = None + s3_config: DistributedMapS3SourceConfig | None = None + reader_config: DistributedMapReaderFunctionSourceConfig | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {"Type": self.source_type.value} + if self.source_type is DistributedMapSourceType.INLINE: + inline = self.inline_source_config or DistributedMapInlineSourceConfig() + result["InlineSourceConfig"] = inline.to_dict() + elif ( + self.source_type is DistributedMapSourceType.S3 + and self.s3_config is not None + ): + result["S3SourceConfig"] = self.s3_config.to_dict() + elif ( + self.source_type is DistributedMapSourceType.READER_FUNCTION + and self.reader_config is not None + ): + result["ReaderFunctionSourceConfig"] = self.reader_config.to_dict() + if self.max_items is not None: + result["MaxItemsToRead"] = self.max_items + return result + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapSourceConfig: + source_type = DistributedMapSourceType(data.get("Type", "INLINE")) + inline_cfg = data.get("InlineSourceConfig") or {} + s3_raw = data.get("S3SourceConfig") + reader_raw = data.get("ReaderFunctionSourceConfig") + return cls( + source_type=source_type, + max_items=data.get("MaxItemsToRead"), + inline_source_config=DistributedMapInlineSourceConfig.from_dict(inline_cfg) + if source_type is DistributedMapSourceType.INLINE + else None, + s3_config=DistributedMapS3SourceConfig.from_dict(s3_raw) + if s3_raw is not None + else None, + reader_config=DistributedMapReaderFunctionSourceConfig.from_dict(reader_raw) + if reader_raw is not None + else None, + ) + + @classmethod + def create_inline( + cls, items: list[str], *, max_items: int | None = None + ) -> DistributedMapSourceConfig: + """Build an inline source from items the caller already serialized.""" + return cls( + source_type=DistributedMapSourceType.INLINE, + max_items=max_items, + inline_source_config=DistributedMapInlineSourceConfig(items=items), + ) + + @classmethod + def create_s3( + cls, s3: S3Source, *, max_items: int | None = None + ) -> DistributedMapSourceConfig: + """Build an S3 source from its resolved configuration.""" + return cls( + source_type=DistributedMapSourceType.S3, + max_items=max_items, + s3_config=DistributedMapS3SourceConfig.from_source(s3), + ) + + @classmethod + def create_reader( + cls, + function_name: str, + initial_state: str | None, + *, + max_items: int | None = None, + ) -> DistributedMapSourceConfig: + """Build a reader-function source from a state the caller already serialized.""" + return cls( + source_type=DistributedMapSourceType.READER_FUNCTION, + max_items=max_items, + reader_config=DistributedMapReaderFunctionSourceConfig( + function_name=function_name, initial_state=initial_state + ), + ) + + +_UNLIMITED_RETRY_WIRE = -1 + +_RESPONSE_TYPE_FOR_MODE = { + ProcessorResponseMode.ITEM_OUTCOMES: DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_OUTCOMES, +} + + +@dataclass(frozen=True) +class DistributedMapProcessorConfig: + """Represent a map run processor.""" + + function_name: str + function_response_types: list[DistributedMapFunctionResponseType] | 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"] = [ + t.value for t in 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]) -> DistributedMapProcessorConfig: + response_types = data.get("FunctionResponseTypes") + return cls( + function_name=data.get("FunctionName", ""), + function_response_types=[ + DistributedMapFunctionResponseType(t) for t in 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"), + ) + + @classmethod + def from_processor(cls, processor: DistributedMapProcessor) -> Self: + """Translate the processor, mapping the response mode and unlimited retries.""" + response_type = _RESPONSE_TYPE_FOR_MODE.get(processor.response_mode) + max_retry_attempts: int | None = None + attempts = processor.max_retry_attempts + if attempts == DistributedMapProcessor.UNLIMITED: + max_retry_attempts = _UNLIMITED_RETRY_WIRE + elif isinstance(attempts, int): + max_retry_attempts = attempts + return cls( + function_name=processor.function_name, + function_response_types=[response_type] if response_type else None, + batch_size=processor.batch_size, + max_retry_attempts=max_retry_attempts, + max_retry_duration_seconds=processor.max_retry_duration.to_seconds() + if processor.max_retry_duration is not None + else None, + durable_execution_name_prefix=processor.durable_execution_name_prefix, + ) + + +@dataclass(frozen=True) +class DistributedMapCompletionConfig: + """Represent a map run completion (failure-tolerance) config.""" + + tolerated_failure_count: int | None = None + tolerated_failure_percentage: int | 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] + ) -> DistributedMapCompletionConfig: + return cls( + tolerated_failure_count=data.get("ToleratedFailureCount"), + tolerated_failure_percentage=data.get("ToleratedFailurePercentage"), + minimum_sample_size=data.get("MinimumSampleSize"), + ) + + @classmethod + def from_config(cls, completion_config: ConfigCompletionConfig) -> Self: + """Translate the completion config.""" + return cls( + tolerated_failure_count=completion_config.tolerated_failure_count, + tolerated_failure_percentage=completion_config.tolerated_failure_percentage, + minimum_sample_size=completion_config.minimum_sample_size, + ) + + +@dataclass(frozen=True) +class DistributedMapS3DestinationConfig: + """Represent an S3 distributed map destination config.""" + + bucket: str + key_prefix: str + expected_bucket_owner: str | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = { + "Bucket": self.bucket, + "KeyPrefix": self.key_prefix, + } + if self.expected_bucket_owner is not None: + result["ExpectedBucketOwner"] = self.expected_bucket_owner + return result + + @classmethod + def from_dict( + cls, data: MutableMapping[str, Any] + ) -> DistributedMapS3DestinationConfig: + return cls( + bucket=data.get("Bucket", ""), + key_prefix=data.get("KeyPrefix", ""), + expected_bucket_owner=data.get("ExpectedBucketOwner"), + ) + + +@dataclass(frozen=True) +class _DistributedMapDestinationEntry: + """Hold the members the OnSuccess and OnFailure shapes have in common.""" + + type: DistributedMapDestinationType + include: list[DistributedMapDestinationInclude] + s3_destination_config: DistributedMapS3DestinationConfig + + def to_dict(self) -> MutableMapping[str, Any]: + return { + "Type": self.type.value, + "Include": [i.value for i in self.include], + "S3DestinationConfig": self.s3_destination_config.to_dict(), + } + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> Self: + s3_raw = data.get("S3DestinationConfig") or {} + return cls( + type=DistributedMapDestinationType(data.get("Type", "S3")), + include=[ + DistributedMapDestinationInclude(i) for i in data.get("Include", []) + ], + s3_destination_config=DistributedMapS3DestinationConfig.from_dict(s3_raw), + ) + + @classmethod + def _create( + cls, + destination: SuccessDestination | FailureDestination, + include: list[DistributedMapDestinationInclude], + ) -> Self: + """Build an entry from a destination and the include set its subclass computed.""" + return cls( + type=DistributedMapDestinationType.S3, + include=include, + s3_destination_config=DistributedMapS3DestinationConfig( + bucket=destination.bucket, + key_prefix=destination.prefix, + expected_bucket_owner=destination.expected_bucket_owner, + ), + ) + + +@dataclass(frozen=True) +class DistributedMapOnSuccessConfig(_DistributedMapDestinationEntry): + """Represent where a map run writes the outcome of its succeeded items.""" + + @classmethod + def from_destination(cls, destination: SuccessDestination) -> Self: + """Translate the success destination.""" + include: list[DistributedMapDestinationInclude] = [] + if destination.include_input: + include.append(DistributedMapDestinationInclude.INPUT) + if destination.include_output: + include.append(DistributedMapDestinationInclude.OUTPUT) + return cls._create(destination, include) + + +@dataclass(frozen=True) +class DistributedMapOnFailureConfig(_DistributedMapDestinationEntry): + """Represent where a map run writes the outcome of its failed items.""" + + @classmethod + def from_destination(cls, destination: FailureDestination) -> Self: + """Translate the failure destination.""" + include: list[DistributedMapDestinationInclude] = [] + if destination.include_input: + include.append(DistributedMapDestinationInclude.INPUT) + if destination.include_error: + include.append(DistributedMapDestinationInclude.ERROR) + return cls._create(destination, include) + + +@dataclass(frozen=True) +class DistributedMapDestinationConfig: + """Represent a map run destination config.""" + + on_success: DistributedMapOnSuccessConfig | None = None + on_failure: DistributedMapOnFailureConfig | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {} + if self.on_success is not None: + result["OnSuccess"] = self.on_success.to_dict() + if self.on_failure is not None: + result["OnFailure"] = self.on_failure.to_dict() + return result + + @classmethod + def from_dict( + cls, data: MutableMapping[str, Any] + ) -> DistributedMapDestinationConfig: + success_raw = data.get("OnSuccess") + failure_raw = data.get("OnFailure") + return cls( + on_success=DistributedMapOnSuccessConfig.from_dict(success_raw) + if success_raw is not None + else None, + on_failure=DistributedMapOnFailureConfig.from_dict(failure_raw) + if failure_raw is not None + else None, + ) + + @classmethod + def from_destination(cls, destination: DistributedMapDestination) -> Self | None: + """Translate the destination config, or None when neither side is set.""" + if destination.on_success is None and destination.on_failure is None: + return None + return cls( + on_success=DistributedMapOnSuccessConfig.from_destination( + destination.on_success + ) + if destination.on_success is not None + else None, + on_failure=DistributedMapOnFailureConfig.from_destination( + destination.on_failure + ) + if destination.on_failure is not None + else None, + ) + + +@dataclass(frozen=True) +class DistributedMapResultCollectionConfig: + """Represent the map run result-collection setting.""" + + mode: DistributedMapResultCollectionMode + + def to_dict(self) -> MutableMapping[str, Any]: + return {"Mode": self.mode.value} + + @classmethod + def from_dict( + cls, data: MutableMapping[str, Any] + ) -> DistributedMapResultCollectionConfig: + return cls(mode=DistributedMapResultCollectionMode(data.get("Mode", "NONE"))) + + +@dataclass(frozen=True) +class DistributedMapOptions: + """Configuration options for starting a map run.""" + + max_concurrency: int + source: DistributedMapSourceConfig + processor: DistributedMapProcessorConfig + destination: DistributedMapDestinationConfig | None = None + completion_config: DistributedMapCompletionConfig | None = None + result_collection: DistributedMapResultCollectionConfig | 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=DistributedMapSourceConfig.from_dict(source_raw), + processor=DistributedMapProcessorConfig.from_dict(processor_raw), + destination=DistributedMapDestinationConfig.from_dict(destination_raw) + if destination_raw is not None + else None, + completion_config=DistributedMapCompletionConfig.from_dict(completion_raw) + if completion_raw is not None + else None, + result_collection=DistributedMapResultCollectionConfig.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 +1345,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 +1374,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 +1404,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 +1424,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 +1608,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 +1751,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 +1796,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 +1817,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 +1882,33 @@ 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] = { + "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.source_exhausted is not None: + distributed_map_details_dict["SourceExhausted"] = ( + self.distributed_map_details.source_exhausted + ) + if self.distributed_map_details.error is not None: + distributed_map_details_dict["Error"] = ( + self.distributed_map_details.error.to_dict() + ) + 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/dmap.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dmap.py new file mode 100644 index 000000000..097cdd628 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dmap.py @@ -0,0 +1,350 @@ +"""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.config import ( + DistributedMapResultConfig, + DistributedMapSource, +) +from aws_durable_execution_sdk_python.dmap.models import ( + DistributedMapResult, + DistributedMapResultItem, + DistributedMapSummary, +) +from aws_durable_execution_sdk_python.exceptions import ( + INVALID_PARAMETER_VALUE_EXCEPTION, + CheckpointError, + ExecutionError, + ValidationError, +) +from aws_durable_execution_sdk_python.lambda_service import ( + DistributedMapCompletionConfig, + DistributedMapDestinationConfig, + DistributedMapOptions, + DistributedMapProcessorConfig, + DistributedMapResultCollectionConfig, + DistributedMapResultCollectionMode, + DistributedMapSourceConfig, + 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, + DistributedMapProcessor, + ReaderSource, + ) + from aws_durable_execution_sdk_python.identifier import OperationIdentifier + from aws_durable_execution_sdk_python.lambda_service import ( + DistributedMapDetails, + Operation, + ) + from aws_durable_execution_sdk_python.serdes import SerDes + from aws_durable_execution_sdk_python.state import ( + CheckpointedResult, + ExecutionState, + ) + +logger = logging.getLogger(__name__) + +# 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 _rejected_config_error(error: CheckpointError) -> Exception: + """Return the error to raise for a rejected START checkpoint. + + The service rejects an invalid configuration with + ``InvalidParameterValueException``, which the caller must see as a + ``ValidationError`` whether the SDK or the service caught it. A retryable + rejection is a stale checkpoint token rather than a bad configuration, so it + is left as it is. + """ + code = error.error.get("Code") if error.error else None + if code == INVALID_PARAMETER_VALUE_EXCEPTION and not error.is_retryable(): + return ValidationError(str(error)) + return error + + +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 _serialize(self, serdes: SerDes | None, value: Any) -> str: + """Serialize a value with this operation's id and execution ARN.""" + return serialize( + serdes=serdes or DEFAULT_JSON_SERDES, + value=value, + operation_id=self.operation_identifier.operation_id, + durable_execution_arn=self.state.durable_execution_arn, + ) + + def _serialize_inline_items( + self, items: Sequence[Any], serdes: SerDes | None + ) -> list[str]: + """Serialize each inline item, enforcing the 1 MB cap on the whole list.""" + serialized_items = [self._serialize(serdes, item) for item in items] + total = len(json.dumps(serialized_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 serialized_items + + def _reader_initial_state(self, reader: ReaderSource) -> str | None: + """Serialize the reader's initial state, enforcing the 32 KB cap.""" + if reader.initial_state is None: + return None + state = self._serialize(reader.state_serdes, reader.initial_state) + 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) + return state + + def _build_source_config(self) -> DistributedMapSourceConfig: + """Translate the source (typed or plain-list shorthand) into its source config.""" + source = self.source + if not isinstance(source, DistributedMapSource): + # A plain list is treated as an inline source with the default serializer. + return DistributedMapSourceConfig.create_inline( + self._serialize_inline_items(source, None) + ) + if source.inline_items is not None: + return DistributedMapSourceConfig.create_inline( + self._serialize_inline_items(source.inline_items, source.inline_serdes), + max_items=source.max_items, + ) + if source.s3 is not None: + return DistributedMapSourceConfig.create_s3( + source.s3, max_items=source.max_items + ) + if source.reader is not None: + return DistributedMapSourceConfig.create_reader( + source.reader.function_name, + self._reader_initial_state(source.reader), + max_items=source.max_items, + ) + msg = "Distributed map source has no configured items" + raise ExecutionError(msg) + + def _build_options(self) -> DistributedMapOptions: + """Assemble the DistributedMapOptions payload from the operands and config.""" + config = self.config + return DistributedMapOptions( + max_concurrency=self.max_concurrency, + source=self._build_source_config(), + processor=DistributedMapProcessorConfig.from_processor(self.processor), + destination=( + DistributedMapDestinationConfig.from_destination(config.destination) + if config.destination is not None + else None + ), + completion_config=( + DistributedMapCompletionConfig.from_config(config.completion_config) + if config.completion_config is not None + else None + ), + result_collection=( + DistributedMapResultCollectionConfig( + mode=DistributedMapResultCollectionMode.INLINE + ) + if isinstance(config, DistributedMapResultConfig) + else None + ), + timeout_seconds=config.timeout.to_seconds() + if config.timeout is not None + else None, + ) + + def _resolve_summary(self, operation: Operation | None) -> DistributedMapSummary: + """Reconstruct the resolved summary/result from the terminal operation.""" + if not isinstance(self.config, DistributedMapResultConfig): + return DistributedMapSummary.from_operation(operation) + details = operation.distributed_map_details if operation else None + items = ( + self._deserialize_items(details, self.config.result_serdes) + if details is not None + else [] + ) + return DistributedMapResult.from_operation_and_items(operation, items) + + def _deserialize_items( + self, details: DistributedMapDetails, result_serdes: SerDes | None + ) -> list[DistributedMapResultItem]: + """Deserialize the service result items into customer result items.""" + items: list[DistributedMapResultItem] = [] + for entry in details.results or (): + output: Any | None = None + if entry.output is not None: + output = deserialize( + serdes=result_serdes or DEFAULT_JSON_SERDES, + data=entry.output, + operation_id=self.operation_identifier.operation_id, + durable_execution_arn=self.state.durable_execution_arn, + ) + items.append( + DistributedMapResultItem( + item_id=entry.item_id, + status=entry.status, + output=output, + error=entry.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._get_checkpoint_result() + + # Terminal success - build the summary/result from the operation + if checkpointed_result.is_succeeded(): + operation = checkpointed_result.operation + summary = self._resolve_summary(operation) + return CheckResult.create_completed(summary) + + # Operation-level terminal failure. Every terminal state resolves with the + # summary, and throw_if_error opts into raising. + if ( + checkpointed_result.is_failed() + or checkpointed_result.is_timed_out() + or checkpointed_result.is_stopped() + ): + operation = checkpointed_result.operation + if operation is not None and operation.distributed_map_details is not None: + return CheckResult.create_completed(self._resolve_summary(operation)) + status_value = ( + checkpointed_result.status.value + if checkpointed_result.status + else "UNKNOWN" + ) + msg = ( + f"DISTRIBUTED_MAP operation ended {status_value} but carried no " + f"DistributedMapDetails" + ) + raise ExecutionError(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=self._build_options(), + ) + ) + # Checkpoint map run START with blocking (is_sync=True). + # Must ensure the map run is recorded before suspending execution. + try: + self.state.create_checkpoint( + operation_update=start_operation, is_sync=True + ) + except CheckpointError as e: + raise _rejected_config_error(e) from e + + 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 e9549d308..67d83f2f1 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 f5ce7214a..e762beb87 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 @@ -223,6 +223,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/config_test.py b/packages/aws-durable-execution-sdk-python/tests/config_test.py index 5f230974a..a5183478d 100644 --- a/packages/aws-durable-execution-sdk-python/tests/config_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/config_test.py @@ -22,6 +22,15 @@ WaitForConditionConfig, WaitForConditionDecision, ) +from aws_durable_execution_sdk_python.config import ( + DistributedMapCompletionConfig, + DistributedMapConfig, + SuccessDestination, + FailureDestination, + DistributedMapProcessor, + DistributedMapSource, + ReaderSource, +) def test_completion_config_defaults(): @@ -279,3 +288,230 @@ def test_completion_decision_rejects_outcome_when_not_complete(): # endregion Config validation + + +# ========================================================================== +# Distributed map config validation +# ========================================================================== + + +def test_retry_duration_out_of_range_rejected(): + with pytest.raises(ValidationError, match="between 1 minute and 6 hours"): + DistributedMapProcessor.report_batch_outcome( + "p", max_retry_duration=Duration.from_seconds(30) + ) + with pytest.raises(ValidationError, match="between 1 minute and 6 hours"): + DistributedMapProcessor.report_batch_outcome( + "p", 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_delimiter_invalid_string_rejected(): + with pytest.raises(ValidationError, match="delimiter must be one of"): + DistributedMapSource.s3_csv("s3://b/data.csv", delimiter="BAR") + + +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_empty_function_name_rejected(): + with pytest.raises(ValidationError, match="non-empty"): + DistributedMapProcessor.report_batch_outcome("") + + +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_function_name_over_max_length_rejected(): + with pytest.raises(ValidationError, match="at most 170 characters"): + DistributedMapProcessor.report_batch_outcome("f" * 171) + + +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_config_without_a_threshold_rejected(): + """Omitting the config means zero tolerated failures, so an empty one is an error.""" + with pytest.raises(ValidationError, match="must set tolerated_failure_count"): + DistributedMapCompletionConfig() + + +def test_completion_sample_size_requires_percentage(): + with pytest.raises(ValidationError, match="minimum_sample_size"): + DistributedMapCompletionConfig( + tolerated_failure_count=1, 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_fractional_percentage_rejected(): + """The service's percentage is an integer, so a fraction is rejected at the call site.""" + with pytest.raises(ValidationError, match="must be an integer"): + DistributedMapCompletionConfig(tolerated_failure_percentage=12.5) # type: ignore[arg-type] + + +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_retry_duration_below_minimum_rejected(): + with pytest.raises(ValidationError, match="1 minute and 6 hours"): + DistributedMapProcessor.report_batch_outcome( + "p", max_retry_duration=Duration.from_seconds(30) + ) + + +def test_max_items_below_one_rejected(): + with pytest.raises(ValidationError, match="between 1 and 100000000"): + DistributedMapSource.s3_json_lines("s3://b/k.jsonl", max_items=0) + + +def test_max_items_above_the_service_maximum_rejected(): + with pytest.raises(ValidationError, match="between 1 and 100000000"): + DistributedMapSource.s3_json_lines("s3://b/k.jsonl", max_items=100_000_001) + + +def test_csv_headers_above_the_count_limit_rejected(): + with pytest.raises(ValidationError, match="at most 1000 names"): + DistributedMapSource.s3_csv( + "s3://b/f.csv", headers=[f"c{i}" for i in range(1001)] + ) + + +def test_csv_header_name_above_the_length_limit_rejected(): + with pytest.raises(ValidationError, match="at most 256 characters"): + DistributedMapSource.s3_csv("s3://b/f.csv", headers=["a" * 257]) + + +def test_csv_headers_at_the_limits_accepted(): + source = DistributedMapSource.s3_csv( + "s3://b/f.csv", headers=[f"c{i}" for i in range(999)] + ["a" * 256] + ) + assert source.s3 is not None + assert source.s3.headers is not None + assert len(source.s3.headers) == 1000 + + +def test_csv_requires_object_key(): + with pytest.raises(ValidationError, match="csv requires an S3 object key"): + DistributedMapSource.s3_csv("s3://bucket") + + +def test_source_without_any_kind_rejected(): + with pytest.raises(ValidationError, match="exactly one of"): + DistributedMapSource() + + +def test_source_with_two_kinds_rejected(): + with pytest.raises(ValidationError, match="exactly one of"): + DistributedMapSource( + inline_items=("a",), + reader=ReaderSource(function_name="r"), + ) + + +def test_success_destination_all_false_rejected(): + with pytest.raises(ValidationError, match="success destination must include"): + SuccessDestination.from_uri( + "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"): + FailureDestination.from_uri( + "s3://out/bad", include_input=False, include_error=False + ) + + +def test_invalid_s3_uri_rejected(): + with pytest.raises(ValidationError, match="Invalid S3 URI"): + DistributedMapSource.s3_json_lines("s3:///key.jsonl") + + +def test_non_s3_scheme_uri_rejected(): + with pytest.raises(ValidationError, match="must start with s3://"): + DistributedMapSource.s3_json_lines("http://foo/bar") + + +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"): + DistributedMapProcessor.report_batch_outcome("p", max_retry_attempts=-5) + + +@pytest.mark.parametrize("attempts", ["5", "forever", "UNLIMITED", 2.5]) +def test_retry_attempts_that_is_not_an_int_or_unlimited_rejected(attempts): + with pytest.raises(ValidationError, match="non-negative integer"): + DistributedMapProcessor.report_batch_outcome("p", max_retry_attempts=attempts) + + +def test_unlimited_retry_attempts_accepted(): + processor = DistributedMapProcessor.report_batch_outcome( + "p", max_retry_attempts=DistributedMapProcessor.UNLIMITED + ) + assert processor.max_retry_attempts == DistributedMapProcessor.UNLIMITED + + +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_durable_execution_name_prefix_too_long_rejected(): + with pytest.raises(ValidationError, match="1 to 36 characters"): + DistributedMapProcessor.report_batch_outcome( + "p", durable_execution_name_prefix="x" * 37 + ) 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 1b75325b7..cb4128914 100644 --- a/packages/aws-durable-execution-sdk-python/tests/context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/context_test.py @@ -17,6 +17,8 @@ Duration, InvokeConfig, MapConfig, + DistributedMapConfig, + DistributedMapProcessor, ParallelBranch, ParallelConfig, StepConfig, @@ -1093,6 +1095,214 @@ 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_max_concurrency_below_one(max_concurrency: int): + """The SDK validates only that max_concurrency is greater than zero.""" + 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, + ) + + +@patch("aws_durable_execution_sdk_python.context.DistributedMapOperationExecutor") +def test_distributed_map_basic(mock_executor_class): + """distributed_map builds a DISTRIBUTED_MAP executor and returns its process() result.""" + mock_executor = MagicMock() + mock_executor.process.return_value = "map_summary" + mock_executor_class.return_value = mock_executor + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = ( + "arn:aws:durable:us-east-1:123456789012:execution/test" + ) + processor = DistributedMapProcessor.report_batch_outcome("test_processor") + + context = create_test_context(state=mock_state) + expected_operation_id = next(operation_id_sequence()) + + result = context.distributed_map(["a", "b"], processor, max_concurrency=4) + + assert result == "map_summary" + mock_executor_class.assert_called_once_with( + state=mock_state, + operation_identifier=OperationIdentifier( + expected_operation_id, OperationSubType.DISTRIBUTED_MAP, None, None + ), + source=["a", "b"], + processor=processor, + max_concurrency=4, + config=ANY, + ) + mock_executor.process.assert_called_once() + + +@patch("aws_durable_execution_sdk_python.context.DistributedMapOperationExecutor") +def test_distributed_map_with_name_and_config(mock_executor_class): + """distributed_map forwards name into the identifier and passes the given config through.""" + mock_executor = MagicMock() + mock_executor.process.return_value = "configured_summary" + mock_executor_class.return_value = mock_executor + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = ( + "arn:aws:durable:us-east-1:123456789012:execution/test" + ) + processor = DistributedMapProcessor.report_batch_outcome("test_processor") + config = DistributedMapConfig() + + context = create_test_context(state=mock_state) + [context._create_step_id() for _ in range(5)] # Set counter to 5 # noqa: SLF001 + + result = context.distributed_map( + ["a"], processor, max_concurrency=2, name="named_map", config=config + ) + + seq = operation_id_sequence() + [next(seq) for _ in range(5)] + expected_id = next(seq) + + assert result == "configured_summary" + mock_executor_class.assert_called_once_with( + state=mock_state, + operation_identifier=OperationIdentifier( + expected_id, OperationSubType.DISTRIBUTED_MAP, None, "named_map" + ), + source=["a"], + processor=processor, + max_concurrency=2, + config=config, + ) + mock_executor.process.assert_called_once() + + +@patch("aws_durable_execution_sdk_python.context.DistributedMapOperationExecutor") +def test_distributed_map_with_parent_id(mock_executor_class): + """distributed_map propagates the parent_id into the operation identifier.""" + mock_executor = MagicMock() + mock_executor.process.return_value = "parent_summary" + mock_executor_class.return_value = mock_executor + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = ( + "arn:aws:durable:us-east-1:123456789012:execution/test" + ) + processor = DistributedMapProcessor.report_batch_outcome("test_processor") + + context = create_test_context(state=mock_state, parent_id="parent123") + [context._create_step_id() for _ in range(2)] # Set counter to 2 # noqa: SLF001 + + context.distributed_map(["a"], processor, max_concurrency=1) + + seq = operation_id_sequence("parent123") + [next(seq) for _ in range(2)] + expected_id = next(seq) + + mock_executor_class.assert_called_once_with( + state=mock_state, + operation_identifier=OperationIdentifier( + expected_id, OperationSubType.DISTRIBUTED_MAP, "parent123", None + ), + source=["a"], + processor=processor, + max_concurrency=1, + config=ANY, + ) + mock_executor.process.assert_called_once() + + +@patch("aws_durable_execution_sdk_python.context.DistributedMapOperationExecutor") +def test_distributed_map_increments_counter(mock_executor_class): + """distributed_map increments the step counter once per call.""" + mock_executor = MagicMock() + mock_executor.process.return_value = "result" + mock_executor_class.return_value = mock_executor + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = ( + "arn:aws:durable:us-east-1:123456789012:execution/test" + ) + processor = DistributedMapProcessor.report_batch_outcome("test_processor") + + context = create_test_context(state=mock_state) + [context._create_step_id() for _ in range(10)] # Set counter to 10 # noqa: SLF001 + + context.distributed_map(["a"], processor, max_concurrency=1) + context.distributed_map(["b"], processor, max_concurrency=1) + + seq = operation_id_sequence() + [next(seq) for _ in range(10)] + expected_id1 = next(seq) + expected_id2 = next(seq) + + assert context._step_counter.get_current() == 12 # noqa: SLF001 + assert mock_executor_class.call_args_list[0][1][ + "operation_identifier" + ] == OperationIdentifier(expected_id1, OperationSubType.DISTRIBUTED_MAP, None, None) + assert mock_executor_class.call_args_list[1][1][ + "operation_identifier" + ] == OperationIdentifier(expected_id2, OperationSubType.DISTRIBUTED_MAP, None, None) + + +@patch("aws_durable_execution_sdk_python.context.DistributedMapOperationExecutor") +def test_distributed_map_defaults_config_when_none(mock_executor_class): + """distributed_map builds a default DistributedMapConfig when none is given.""" + mock_executor = MagicMock() + mock_executor.process.return_value = "summary" + mock_executor_class.return_value = mock_executor + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = ( + "arn:aws:durable:us-east-1:123456789012:execution/test" + ) + processor = DistributedMapProcessor.report_batch_outcome("test_processor") + + context = create_test_context(state=mock_state) + context.distributed_map(["a"], processor, max_concurrency=1) + + passed_config = mock_executor_class.call_args[1]["config"] + assert isinstance(passed_config, DistributedMapConfig) + + +@patch("aws_durable_execution_sdk_python.context.DistributedMapOperationExecutor") +def test_distributed_map_returns_process_result(mock_executor_class): + """distributed_map returns whatever executor.process() returns (summary or result).""" + mock_executor = MagicMock() + sentinel = object() + mock_executor.process.return_value = sentinel + mock_executor_class.return_value = mock_executor + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = ( + "arn:aws:durable:us-east-1:123456789012:execution/test" + ) + processor = DistributedMapProcessor.report_batch_outcome("test_processor") + + context = create_test_context(state=mock_state) + result = context.distributed_map(["a"], processor, max_concurrency=1) + + assert result is sentinel + + +# 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/dmap/handlers_test.py b/packages/aws-durable-execution-sdk-python/tests/dmap/handlers_test.py new file mode 100644 index 000000000..02c892cd7 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/dmap/handlers_test.py @@ -0,0 +1,323 @@ +"""Unit tests for the non-durable distributed map authoring wrappers.""" + +from __future__ import annotations + +import json +import threading +import time +from typing import Any + +import pytest + +from aws_durable_execution_sdk_python.dmap.handlers import ( + ItemFailure, + ReaderPage, + distributed_map_batch_handler, + distributed_map_item_handler, + durable_distributed_map_item_handler, + distributed_map_reader, +) +from aws_durable_execution_sdk_python.exceptions import SerDesError, ValidationError +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.serdes import SerDes, SerDesContext + + +def test_item_handler_rejects_concurrency_below_one(): + with pytest.raises(ValidationError, match="concurrency must be at least 1"): + distributed_map_item_handler(lambda x: x, concurrency=0) + with pytest.raises(ValidationError, match="concurrency must be at least 1"): + distributed_map_item_handler(lambda x: x, concurrency=-1) + + +def test_durable_item_handler_rejects_concurrency_below_one(): + with pytest.raises(ValidationError, match="concurrency must be at least 1"): + durable_distributed_map_item_handler(lambda _ctx, item: item, concurrency=0) + with pytest.raises(ValidationError, match="concurrency must be at least 1"): + durable_distributed_map_item_handler(lambda _ctx, item: item, concurrency=-1) + + +def test_item_handler_runs_items_one_at_a_time_by_default(): + lock = threading.Lock() + in_flight = 0 + peak = 0 + + def process(x): + nonlocal in_flight, peak + with lock: + in_flight += 1 + peak = max(peak, in_flight) + # Hold the worker so that a wider pool would overlap this item with the + # next one, which is what makes the peak below meaningful. + time.sleep(0.02) + with lock: + in_flight -= 1 + return x + + handler = distributed_map_item_handler(process) + handler({"records": [{"itemId": str(i), "body": "1"} for i in range(4)]}) + assert peak == 1 + + +def test_item_handler_honors_requested_concurrency(): + # Each item blocks until all four have arrived, so the batch only finishes + # if the pool really runs four items at once. At a smaller pool size the + # first item waits out the timeout and every item fails. + barrier = threading.Barrier(4, timeout=5) + + def process(x): + barrier.wait() + return x + + handler = distributed_map_item_handler(process, concurrency=4) + resp = handler({"records": [{"itemId": str(i), "body": str(i)} for i in range(4)]}) + assert resp["batchItemFailures"] == [] + # The four items are released together, so completion order is arbitrary + # while the reported results stay in item order. + assert resp["batchItemResults"] == [ + {"itemIdentifier": "0", "output": "0"}, + {"itemIdentifier": "1", "output": "1"}, + {"itemIdentifier": "2", "output": "2"}, + {"itemIdentifier": "3", "output": "3"}, + ] + + +def test_item_handler_reports_results_in_order(): + handler = 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 = 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_reports_every_member_of_a_rich_error(): + """An error carrying data and a stack trace reaches the service with all four members.""" + + def process(_x: Any) -> Any: + raise SerDesError( + message="bad payload", + error_type="ignored, SerDesError pins its own type", + data="payload", + stack_trace=["frame one", "frame two"], + ) + + handler = distributed_map_item_handler(process) + resp = handler({"records": [{"itemId": "0", "body": "1"}]}) + + error = resp["batchItemFailures"][0]["error"] + assert error["errorMessage"] == "bad payload" + assert error["errorData"] == "payload" + assert error["stackTrace"] == ["frame one", "frame two"] + + +def test_item_handler_omits_error_members_the_exception_lacks(): + """A plain exception carries no data or trace, so neither key is written.""" + + def process(_x: Any) -> Any: + msg = "boom" + raise ValueError(msg) + + handler = distributed_map_item_handler(process) + resp = handler({"records": [{"itemId": "0", "body": "1"}]}) + + assert resp["batchItemFailures"][0]["error"] == { + "errorType": "ValueError", + "errorMessage": "boom", + } + + +def test_item_failure_wire_round_trip_keeps_every_member(): + failure = ItemFailure( + item_identifier="7", + error=ErrorObject( + message="boom", type="E", data="payload", stack_trace=["frame"] + ), + ) + parsed = ItemFailure.from_dict(failure.to_dict()) + assert parsed == failure + + +def test_item_handler_omits_the_output_of_a_handler_returning_none(): + handler = distributed_map_item_handler(lambda _x: None) + resp = handler({"records": [{"itemId": "0", "body": "1"}]}) + assert resp == { + "batchItemFailures": [], + "batchItemResults": [{"itemIdentifier": "0"}], + } + + +def test_item_handler_fails_only_the_item_whose_output_cannot_be_serialized(): + """A result serdes failure is that item's failure, not the batch's.""" + + class _RejectsBad(SerDes[Any]): + def serialize(self, value: Any, _serdes_context: SerDesContext) -> str: + if value == "bad": + msg = "cannot serialize" + raise TypeError(msg) + return json.dumps(value) + + def deserialize(self, data: str, _serdes_context: SerDesContext) -> Any: + return json.loads(data) + + handler = distributed_map_item_handler(lambda x: x, result_serdes=_RejectsBad()) + 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": "TypeError", "errorMessage": "cannot serialize"}, + } + ] + + +def test_batch_handler_success_and_propagates_error(): + seen: list = [] + + def process(items): + seen.extend(items) + return "ignored" + + handler = distributed_map_batch_handler(process) + # A batch processor reports no output, so the return value is not forwarded. + assert ( + handler( + {"records": [{"itemId": "0", "body": "1"}, {"itemId": "1", "body": "2"}]} + ) + is None + ) + assert seen == [1, 2] + + def boom(_items): + msg = "batch failed" + raise RuntimeError(msg) + + failing = 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, max_items_per_page): + if state is None: + return ReaderPage(items=[1, 2][:max_items_per_page], next_state={"page": 1}) + return ReaderPage(items=[3]) + + handler = distributed_map_reader(read) + first = handler({"state": None, "maxItemsPerPage": 10}) + assert first["items"] == ["1", "2"] + assert first["nextState"] == '{"page": 1}' + + second = handler({"state": '{"page": 1}', "maxItemsPerPage": 10}) + assert second["items"] == ["3"] + assert "nextState" not in second + + +def test_reader_serializes_items_with_the_item_serdes(): + """An item serdes encodes each item, so the processor's can decode it.""" + + class _UpperSerDes(SerDes[Any]): + def serialize(self, value: Any, _serdes_context: SerDesContext) -> str: + return json.dumps(value.upper()) + + def deserialize(self, data: str, _serdes_context: SerDesContext) -> Any: + return json.loads(data).lower() + + serdes = _UpperSerDes() + reader = distributed_map_reader( + lambda _s, _max: ReaderPage(items=["a", "b"]), item_serdes=serdes + ) + page = reader({"state": None, "maxItemsPerPage": 10}) + assert page["items"] == ['"A"', '"B"'] + + seen = [] + processor = distributed_map_batch_handler(seen.extend, item_serdes=serdes) + processor( + { + "records": [ + {"itemId": str(i), "body": body} for i, body in enumerate(page["items"]) + ] + } + ) + assert seen == ["a", "b"] + + +def test_reader_item_serdes_is_separate_from_the_state_serdes(): + """Items use the item serdes and next_state uses the state serdes.""" + + class _PrefixSerDes(SerDes[Any]): + def __init__(self, prefix: str) -> None: + self.prefix = prefix + + def serialize(self, value: Any, _serdes_context: SerDesContext) -> str: + return self.prefix + json.dumps(value) + + def deserialize(self, data: str, _serdes_context: SerDesContext) -> Any: + return json.loads(data.removeprefix(self.prefix)) + + handler = distributed_map_reader( + lambda state, _max: ReaderPage(items=[state], next_state="s"), + item_serdes=_PrefixSerDes("item:"), + state_serdes=_PrefixSerDes("state:"), + ) + page = handler({"state": 'state:"in"', "maxItemsPerPage": 10}) + assert page["items"] == ['item:"in"'] + assert page["nextState"] == 'state:"s"' + + +def test_reader_receives_the_page_bound(): + """The page bound reaches func, so a reader can page to it.""" + handler = distributed_map_reader( + lambda _s, max_items_per_page: ReaderPage(items=list(range(max_items_per_page))) + ) + assert handler({"state": None, "maxItemsPerPage": 3})["items"] == ["0", "1", "2"] + + +def test_reader_rejects_page_over_max_items_per_page(): + handler = distributed_map_reader(lambda _s, _max: ReaderPage(items=[1, 2, 3])) + with pytest.raises(ValidationError, match="exceeding maxItemsPerPage"): + handler({"state": None, "maxItemsPerPage": 2}) + + +def test_reader_rejects_oversized_next_state(): + handler = distributed_map_reader( + lambda _s, _max: ReaderPage(items=[1], next_state="x" * 40_000) + ) + with pytest.raises(ValidationError, match="32 KB limit"): + handler({"state": None, "maxItemsPerPage": 10}) + + +def test_item_handler_rejects_non_processor_envelope(): + handler = distributed_map_item_handler(lambda x: x) + with pytest.raises(ValidationError, match="processor envelope"): + handler({}) + + +def test_reader_rejects_non_reader_envelope(): + handler = distributed_map_reader(lambda _state, _max: ReaderPage(items=[])) + with pytest.raises(ValidationError, match="reader envelope"): + handler({"state": None}) diff --git a/packages/aws-durable-execution-sdk-python/tests/dmap/models_test.py b/packages/aws-durable-execution-sdk-python/tests/dmap/models_test.py new file mode 100644 index 000000000..edbca77e3 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/dmap/models_test.py @@ -0,0 +1,423 @@ +"""Unit tests for the distributed map result types.""" + +from __future__ import annotations + +import pytest + +from aws_durable_execution_sdk_python.config import ( + DistributedMapItemStatus, + DistributedMapStatus, +) +from aws_durable_execution_sdk_python.dmap.models import ( + DistributedMapResult, + DistributedMapResultItem, + DistributedMapRunArn, + DistributedMapSummary, + _to_distributed_map_status, +) +from aws_durable_execution_sdk_python.exceptions import ( + DistributedMapError, + ExecutionError, +) +from aws_durable_execution_sdk_python.lambda_service import ( + DistributedMapDetails, + ErrorObject, + Operation, + OperationStatus, + OperationType, +) + + +def _terminal_operation( + status: OperationStatus, details: DistributedMapDetails +) -> Operation: + return Operation( + operation_id="dmap-id", + operation_type=OperationType.DISTRIBUTED_MAP, + status=status, + distributed_map_details=details, + ) + + +def _result(status, failure_count, items): + return DistributedMapResult( + status=status, + success_count=len(items) - failure_count, + failure_count=failure_count, + unprocessed_count=0, + all=items, + ) + + +def test_summary_throw_if_error(): + ok = DistributedMapSummary( + status=DistributedMapStatus.SUCCEEDED, + success_count=2, + failure_count=0, + unprocessed_count=0, + ) + ok.throw_if_error() # no raise + + failed = DistributedMapSummary( + status=DistributedMapStatus.FAILED, + success_count=0, + failure_count=1, + unprocessed_count=0, + ) + with pytest.raises(DistributedMapError): + failed.throw_if_error() + + succeeded_with_failures = DistributedMapSummary( + status=DistributedMapStatus.SUCCEEDED, + 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=DistributedMapItemStatus.SUCCEEDED, output=1 + ), + DistributedMapResultItem( + item_id="1", + status=DistributedMapItemStatus.FAILED, + error=ErrorObject(message="boom", type="E", data=None, stack_trace=None), + ), + ] + result = DistributedMapResult( + status=DistributedMapStatus.SUCCEEDED, + 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].message == "boom" + + +def test_map_run_result_accessors_keep_items_with_nothing_recorded(): + items = [ + DistributedMapResultItem( + item_id="0", status=DistributedMapItemStatus.SUCCEEDED, output=None + ), + DistributedMapResultItem( + item_id="1", status=DistributedMapItemStatus.SUCCEEDED, output=1 + ), + DistributedMapResultItem( + item_id="2", status=DistributedMapItemStatus.FAILED, error=None + ), + DistributedMapResultItem( + item_id="3", + status=DistributedMapItemStatus.FAILED, + error=ErrorObject(message="boom", type="E", data=None, stack_trace=None), + ), + ] + result = DistributedMapResult( + status=DistributedMapStatus.SUCCEEDED, + success_count=2, + failure_count=2, + unprocessed_count=0, + all=items, + ) + # One entry per item of that status, so both lists stay aligned with + # succeeded() and failed(). + assert result.get_results() == [None, 1] + assert len(result.get_results()) == len(result.succeeded()) + assert [e.message if e else None for e in result.get_errors()] == [ + None, + "boom", + ] + assert len(result.get_errors()) == len(result.failed()) + + +def test_map_run_result_accessors_omit_an_unknown_item_status(): + """An item of an unrecognized status stays in all and is in no accessor.""" + items = [ + DistributedMapResultItem( + item_id="0", status=DistributedMapItemStatus.SUCCEEDED, output=1 + ), + DistributedMapResultItem( + item_id="1", + status=DistributedMapItemStatus.UNKNOWN_TO_SDK_VERSION, + output=2, + ), + DistributedMapResultItem( + item_id="2", + status=DistributedMapItemStatus.FAILED, + error=ErrorObject(message="boom", type="E", data=None, stack_trace=None), + ), + ] + result = DistributedMapResult( + status=DistributedMapStatus.SUCCEEDED, + success_count=1, + failure_count=1, + unprocessed_count=0, + all=items, + ) + assert [i.item_id for i in result.all] == ["0", "1", "2"] + assert [i.item_id for i in result.succeeded()] == ["0"] + assert [i.item_id for i in result.failed()] == ["2"] + assert result.get_results() == [1] + assert [e.message if e else None for e in result.get_errors()] == ["boom"] + + +def test_summary_has_failure_false_without_failures(): + summary = DistributedMapSummary( + status=DistributedMapStatus.SUCCEEDED, + success_count=0, + failure_count=0, + unprocessed_count=0, + ) + assert summary.has_failure is False + + +_RUN_ARN = ( + "arn:aws:lambda:us-east-1:123456789012:function:fn:$LATEST" + "/durable-execution/exec1/invoke1/distributed-map-run/abc" +) + + +def _summary_with_arn(arn): + return DistributedMapSummary( + status=DistributedMapStatus.SUCCEEDED, + success_count=0, + failure_count=0, + unprocessed_count=0, + distributed_map_run_arn=arn, + source_exhausted=True, + ) + + +def test_run_arn_from_arn_splits_parent_and_run_id(): + parsed = DistributedMapRunArn.from_arn(_RUN_ARN) + assert parsed is not None + assert parsed.run_id == "abc" + assert parsed.durable_execution_arn == ( + "arn:aws:lambda:us-east-1:123456789012:function:fn:$LATEST" + "/durable-execution/exec1/invoke1" + ) + + +@pytest.mark.parametrize( + "arn", + [ + "arn:aws:lambda:us-east-1:123456789012:function:fn:$LATEST/durable-execution/exec1/invoke1", + "arn:aws:states:us-east-1:123456789012:mapRun:sm/abc", + _RUN_ARN + "/extra", + "not-an-arn/distributed-map-run/abc", + "abc", + "", + ], +) +def test_run_arn_from_arn_returns_none_for_other_shapes(arn): + assert DistributedMapRunArn.from_arn(arn) is None + + +def test_summary_distributed_map_id_parses_run_arn(): + assert _summary_with_arn(_RUN_ARN).distributed_map_id == "abc" + + +def test_summary_distributed_map_id_none_without_arn(): + assert _summary_with_arn(None).distributed_map_id is None + + +def test_summary_distributed_map_id_none_for_unparseable_arn(): + assert _summary_with_arn("not-an-arn").distributed_map_id is None + + +def test_summary_carries_source_exhausted_and_run_error(): + """The terminal details' SourceExhausted and Error reach the summary.""" + details = DistributedMapDetails( + distributed_map_run_arn=_RUN_ARN, + source_exhausted=False, + error=ErrorObject( + message="tolerance exceeded", + type="DistributedMapFailure", + data="d", + stack_trace=["line"], + ), + success_count=1, + failure_count=2, + unprocessed_count=3, + ) + summary = DistributedMapSummary.from_operation( + _terminal_operation(OperationStatus.FAILED, details) + ) + assert summary.source_exhausted is False + assert summary.error is not None + assert summary.error.message == "tolerance exceeded" + assert summary.error.data == "d" + assert summary.error.stack_trace == ["line"] + + +def test_summary_throw_if_error_carries_the_run_error(): + """A thrown run-level error carries every member of the summary's error object.""" + details = DistributedMapDetails( + distributed_map_run_arn=_RUN_ARN, + source_exhausted=True, + error=ErrorObject( + message="boom", type="SourceFailure", data="d", stack_trace=["line"] + ), + failure_count=1, + ) + summary = DistributedMapSummary.from_operation( + _terminal_operation(OperationStatus.FAILED, details) + ) + with pytest.raises(DistributedMapError, match="Map run ended FAILED: boom") as exc: + summary.throw_if_error() + assert exc.value.error_type == "SourceFailure" + assert exc.value.data == "d" + assert exc.value.stack_trace == ["line"] + + +def test_summary_missing_details_names_the_operation_status(): + """The missing-details error reports the status the operation actually ended in.""" + op = _terminal_operation(OperationStatus.FAILED, None) + with pytest.raises( + ExecutionError, match="ended FAILED but carried no DistributedMapDetails" + ): + DistributedMapSummary.from_operation(op) + + +def test_summary_missing_source_exhausted_raises_once_a_run_exists(): + """SourceExhausted is required at terminal once a run was created.""" + details = DistributedMapDetails( + distributed_map_run_arn=_RUN_ARN, success_count=1, failure_count=0 + ) + op = _terminal_operation(OperationStatus.SUCCEEDED, details) + with pytest.raises(ExecutionError, match="carried no SourceExhausted"): + DistributedMapSummary.from_operation(op) + + +def test_summary_allows_missing_source_exhausted_before_run_creation(): + """An operation that failed before run creation carries neither ARN nor SourceExhausted.""" + details = DistributedMapDetails( + error=ErrorObject( + message="admission denied", type="E", data=None, stack_trace=None + ) + ) + summary = DistributedMapSummary.from_operation( + _terminal_operation(OperationStatus.FAILED, details) + ) + assert summary.distributed_map_run_arn is None + assert summary.source_exhausted is None + assert summary.success_count == 0 + assert summary.failure_count == 0 + assert summary.unprocessed_count == 0 + + +def test_result_throw_if_error_raises_first_item_error(): + result = _result( + DistributedMapStatus.SUCCEEDED, + 1, + [ + DistributedMapResultItem( + item_id="0", status=DistributedMapItemStatus.SUCCEEDED, output=1 + ), + DistributedMapResultItem( + item_id="1", + status=DistributedMapItemStatus.FAILED, + error=ErrorObject( + message="boom", type="E", data=None, stack_trace=None + ), + ), + ], + ) + 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, + 1, + [ + DistributedMapResultItem( + item_id="7", status=DistributedMapItemStatus.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, + 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, + 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, + success_count=1, + failure_count=0, + unprocessed_count=0, + all=[ + DistributedMapResultItem( + item_id="0", status=DistributedMapItemStatus.SUCCEEDED, output=1 + ) + ], + ) + result.throw_if_error() + + +def test_result_from_the_inherited_factory_has_no_items(): + """The inherited summary factory builds a result whose all is empty.""" + details = DistributedMapDetails( + distributed_map_run_arn=_RUN_ARN, + source_exhausted=True, + success_count=0, + failure_count=0, + unprocessed_count=0, + ) + result = DistributedMapResult.from_operation( + _terminal_operation(OperationStatus.SUCCEEDED, details) + ) + assert isinstance(result, DistributedMapResult) + assert result.all == [] + assert result.status is DistributedMapStatus.SUCCEEDED + + +def test_status_derived_from_operation_when_details_omit_status(): + """Status comes from the operation, not from details.""" + details = DistributedMapDetails( + success_count=3, + total_count=3, + ) + op = _terminal_operation(OperationStatus.SUCCEEDED, details) + assert _to_distributed_map_status(op) is DistributedMapStatus.SUCCEEDED + + +def test_to_distributed_map_status_rejects_non_terminal(): + """A non-terminal operation status cannot map to a DistributedMapStatus.""" + op = _terminal_operation( + OperationStatus.STARTED, + DistributedMapDetails(success_count=0), + ) + with pytest.raises(ExecutionError, match="Cannot derive distributed map status"): + _to_distributed_map_status(op) diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/dmap_helpers_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/dmap_helpers_int_test.py new file mode 100644 index 000000000..114c8a29f --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/dmap_helpers_int_test.py @@ -0,0 +1,261 @@ +"""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 +import threading +import time +from typing import Any +from unittest.mock import Mock, patch + +from aws_durable_execution_sdk_python.dmap.handlers import ( + durable_distributed_map_batch_handler, + durable_distributed_map_item_handler, +) +from aws_durable_execution_sdk_python.exceptions import DurableOperationError +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, +) +from aws_durable_execution_sdk_python.serdes import SerDes, SerDesContext + +_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 = durable_distributed_map_item_handler(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 = durable_distributed_map_item_handler(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_item_handler_forwards_every_error_member(): + """The branch error's data and stack trace reach the service, not just type and message.""" + + def process(_ctx, _item): + raise DurableOperationError( + message="boom", + error_type="Inner", + data="payload", + stack_trace=["frame one", "frame two"], + ) + + handler = durable_distributed_map_item_handler(process) + result = _run(handler, [{"itemId": "0", "body": "1"}]) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + data = json.loads(result["Result"]) + + error = data["batchItemFailures"][0]["error"] + # A branch error is checkpointed and rebuilt, which normalizes errorType to + # the raised class's qualified name. + assert error["errorType"] == ( + f"{DurableOperationError.__module__}.{DurableOperationError.__qualname__}" + ) + assert error["errorMessage"] == "boom" + assert error["errorData"] == "payload" + assert error["stackTrace"] == ["frame one", "frame two"] + + +def test_durable_batch_handler_records_no_output(): + """A batch processor reports no output, so the return value is not forwarded.""" + seen: list = [] + + def process(_ctx, items): + seen.extend(items) + return {"count": len(items)} + + handler = durable_distributed_map_batch_handler(process) + result = _run(handler, [{"itemId": "0", "body": "1"}, {"itemId": "1", "body": "2"}]) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + assert json.loads(result["Result"]) is None + assert seen == [1, 2] + + +def test_durable_item_handler_fails_only_the_item_whose_output_cannot_be_serialized(): + """A result serdes failure is that item's failure, not the batch's.""" + + class _RejectsBad(SerDes[Any]): + def serialize(self, value: Any, _serdes_context: SerDesContext) -> str: + if value == "bad": + msg = "cannot serialize" + raise TypeError(msg) + return json.dumps(value) + + def deserialize(self, data: str, _serdes_context: SerDesContext) -> Any: + return json.loads(data) + + handler = durable_distributed_map_item_handler( + lambda _ctx, item: item, result_serdes=_RejectsBad() + ) + 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"] == [ + { + "itemIdentifier": "1", + "error": {"errorType": "TypeError", "errorMessage": "cannot serialize"}, + } + ] + + +def test_durable_item_handler_omits_the_output_of_a_handler_returning_none(): + handler = durable_distributed_map_item_handler(lambda _ctx, _item: None) + result = _run(handler, [{"itemId": "0", "body": "1"}]) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + data = json.loads(result["Result"]) + assert data == { + "batchItemFailures": [], + "batchItemResults": [{"itemIdentifier": "0"}], + } + + +def test_durable_item_handler_runs_items_one_at_a_time_by_default(): + lock = threading.Lock() + in_flight = 0 + peak = 0 + + def process(_ctx, item): + nonlocal in_flight, peak + with lock: + in_flight += 1 + peak = max(peak, in_flight) + # Hold the branch so that a wider bound would overlap this item with the + # next one, which is what makes the peak below meaningful. + time.sleep(0.02) + with lock: + in_flight -= 1 + return item + + handler = durable_distributed_map_item_handler(process) + result = _run(handler, [{"itemId": str(i), "body": "1"} for i in range(4)]) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + assert peak == 1 + + +def test_durable_item_handler_honors_requested_concurrency(): + # Each item blocks until all four have arrived, so the batch only finishes + # if the map really runs four items at once. At a smaller bound the first + # item waits out the timeout and every item fails. + barrier = threading.Barrier(4, timeout=5) + + def process(_ctx, item): + barrier.wait() + return item + + handler = durable_distributed_map_item_handler(process, concurrency=4) + result = _run(handler, [{"itemId": str(i), "body": str(i)} for i in range(4)]) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + data = json.loads(result["Result"]) + assert data["batchItemFailures"] == [] + assert data["batchItemResults"] == [ + {"itemIdentifier": "0", "output": "0"}, + {"itemIdentifier": "1", "output": "1"}, + {"itemIdentifier": "2", "output": "2"}, + {"itemIdentifier": "3", "output": "3"}, + ] diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/dmap_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/dmap_int_test.py new file mode 100644 index 000000000..3c2d9f59e --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/dmap_int_test.py @@ -0,0 +1,251 @@ +"""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.dmap.models import DistributedMapResult +from aws_durable_execution_sdk_python.config import ( + DistributedMapResultConfig, + 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", + "SubType": "DistributedMap", + "Status": "SUCCEEDED", + "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, + sub_type=update.sub_type, + name=update.name, + 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_run_arn": summary.distributed_map_run_arn, + } + + # 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( + { + "SourceExhausted": True, + "SuccessCount": 2, + "FailureCount": 0, + "UnprocessedCount": 0, + "TotalCount": 2, + "DistributedMapRunArn": "arn:aws:lambda:us-east-1:123456789012:function:fn:$LATEST/durable-execution/exec1/invoke1/distributed-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_run_arn": "arn:aws:lambda:us-east-1:123456789012:function:fn:$LATEST/durable-execution/exec1/invoke1/distributed-map-run/abc", + } + + +def test_map_run_result_config_returns_items(): + @durable_execution + def handler(event, context: DurableContext) -> dict[str, Any]: + result = context.distributed_map( + ["a", "b"], + DistributedMapProcessor.report_item_outcomes("proc"), + max_concurrency=2, + config=DistributedMapResultConfig(), + ) + assert isinstance(result, DistributedMapResult) + return { + "results": result.get_results(), + "errors": [ + e.message if e is not None else None for e in result.get_errors() + ], + } + + _map_run_id, replay_event = _replay_event( + { + "SourceExhausted": True, + "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( + { + "SourceExhausted": False, + "Error": { + "ErrorType": "DistributedMapFailure", + "ErrorMessage": "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/exceptions_test.py b/packages/aws-durable-execution-sdk-python/tests/exceptions_test.py index 540f09fa5..9f1a906e7 100644 --- a/packages/aws-durable-execution-sdk-python/tests/exceptions_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/exceptions_test.py @@ -16,6 +16,7 @@ ChildContextError, CheckpointError, CheckpointErrorCategory, + DistributedMapError, DurableApiErrorCategory, DurableExecutionsError, DurableOperationError, @@ -257,6 +258,7 @@ def test_durable_operation_error_with_none_message(): [ StepError, InvokeError, + DistributedMapError, ChildContextError, WaitForConditionError, CallbackError, diff --git a/packages/aws-durable-execution-sdk-python/tests/lambda_service_test.py b/packages/aws-durable-execution-sdk-python/tests/lambda_service_test.py index 516930e20..c582165fc 100644 --- a/packages/aws-durable-execution-sdk-python/tests/lambda_service_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/lambda_service_test.py @@ -7,9 +7,11 @@ import pytest from aws_durable_execution_sdk_python.__about__ import __version__ +from aws_durable_execution_sdk_python.config import DistributedMapItemStatus from aws_durable_execution_sdk_python.exceptions import ( CheckpointError, DurableOperationError, + ExecutionError, GetExecutionStateError, SerDesError, StepError, @@ -42,6 +44,10 @@ WaitOptions, _is_in_var_dir, ) +from aws_durable_execution_sdk_python.lambda_service import ( + DistributedMapDetails, + DistributedMapResultItem, +) # ============================================================================= @@ -1271,6 +1277,42 @@ def test_operation_from_dict_no_options(): assert operation.operation_id == "test-id" +def test_operation_from_dict_in_flight_distributed_map_details_parsed_without_status(): + """In-flight DistributedMapDetails (no SourceExhausted) parses without raising.""" + data = { + "Id": "dmap-id", + "Type": "DISTRIBUTED_MAP", + "Status": "STARTED", + "DistributedMapDetails": { + "TotalCount": 3, + "SuccessCount": 1, + }, + } + operation = Operation.from_dict(data) + assert operation.distributed_map_details is not None + assert operation.distributed_map_details.source_exhausted is None + assert operation.distributed_map_details.error is None + assert operation.distributed_map_details.total_count == 3 + assert operation.distributed_map_details.success_count == 1 + + +def test_operation_from_dict_terminal_distributed_map_details_parsed(): + """Terminal DistributedMapDetails parses into distributed_map_details.""" + data = { + "Id": "dmap-id", + "Type": "DISTRIBUTED_MAP", + "Status": "SUCCEEDED", + "DistributedMapDetails": { + "SourceExhausted": True, + "TotalCount": 3, + "SuccessCount": 3, + }, + } + operation = Operation.from_dict(data) + assert operation.distributed_map_details is not None + assert operation.distributed_map_details.source_exhausted is True + + def test_operation_from_dict_individual_options(): """Test Operation.from_dict with each option type individually.""" # Test with just ContextOptions @@ -3029,3 +3071,112 @@ def test_operation_to_dict_omits_absent_context_error_and_replay_children(): context = op.to_dict()["ContextDetails"] assert context == {"Result": '"hello"'} + + +# ========================================================================== +# Distributed map wire round-trips +# ========================================================================== + + +def test_map_run_details_from_dict_parses_results(): + data = { + "Status": "SUCCEEDED", + "SourceExhausted": True, + "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.total_count == 2 + assert details.source_exhausted is True + 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_details_missing_source_exhausted_parses_with_none(): + # In-flight runs omit SourceExhausted, so parsing must succeed with None. + details = DistributedMapDetails.from_dict({"SuccessCount": 2}) + assert details.source_exhausted is None + assert details.success_count == 2 + + +def test_details_parses_the_run_level_error(): + """The run-level Error parses into the shared ErrorObject with every member.""" + details = DistributedMapDetails.from_dict( + { + "SourceExhausted": False, + "Error": { + "ErrorType": "SourceFailure", + "ErrorMessage": "boom", + "ErrorData": "d", + "StackTrace": ["line"], + }, + } + ) + assert details.error is not None + assert details.error.type == "SourceFailure" + assert details.error.message == "boom" + assert details.error.data == "d" + assert details.error.stack_trace == ["line"] + + +def test_result_item_wire_round_trip(): + item = DistributedMapResultItem.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 = DistributedMapResultItem.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_result_item_unknown_status_maps_to_unknown_to_sdk_version(): + """A status added after this SDK version parses, and the other items are kept.""" + details = DistributedMapDetails.from_dict( + { + "SourceExhausted": True, + "Results": [ + {"ItemId": "0", "Status": "SUCCEEDED", "Output": "1"}, + {"ItemId": "1", "Status": "PARTIALLY_SUCCEEDED"}, + { + "ItemId": "2", + "Status": "FAILED", + "Error": {"ErrorType": "E", "ErrorMessage": "boom"}, + }, + ], + } + ) + + assert details.results is not None + assert [item.status for item in details.results] == [ + DistributedMapItemStatus.SUCCEEDED, + DistributedMapItemStatus.UNKNOWN_TO_SDK_VERSION, + DistributedMapItemStatus.FAILED, + ] + + +def test_result_item_missing_status_raises(): + """Status is required, so its absence breaks the contract.""" + with pytest.raises(ExecutionError, match="carried no item status"): + DistributedMapResultItem.from_dict({"ItemId": "0", "Output": "1"}) diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dmap_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dmap_test.py new file mode 100644 index 000000000..4d1b7f319 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dmap_test.py @@ -0,0 +1,1713 @@ +"""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.dmap.models import ( + DistributedMapResult, + DistributedMapSummary, +) +from aws_durable_execution_sdk_python.config import ( + Duration, + DistributedMapItemStatus, + DistributedMapStatus, + DistributedMapCompletionConfig, + DistributedMapConfig, + DistributedMapResultConfig, + DistributedMapCsvDelimiter, + DistributedMapDestination, + DistributedMapSource, + SuccessDestination, + FailureDestination, + DistributedMapProcessor, +) +from aws_durable_execution_sdk_python.exceptions import ( + ExecutionError, + CheckpointError, + DistributedMapError, + DurableApiErrorCategory, + NonDeterministicExecutionError, + SuspendExecution, + ValidationError, +) +from aws_durable_execution_sdk_python.identifier import OperationIdentifier +from aws_durable_execution_sdk_python.lambda_service import ( + ErrorObject, + DistributedMapDetails, + DistributedMapFunctionResponseType, + DistributedMapOptions, + DistributedMapResultCollectionMode, + DistributedMapResultItem as DistributedMapResultItemApi, + DistributedMapSourceType, + Operation, + OperationAction, + OperationStatus, + OperationSubType, + OperationType, +) +from aws_durable_execution_sdk_python.operation.dmap import ( + DistributedMapOperationExecutor, +) +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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + success_count=5, + failure_count=0, + unprocessed_count=0, + total_count=5, + distributed_map_run_arn="arn:aws:lambda:us-east-1:123456789012:function:fn:$LATEST/durable-execution/exec1/invoke1/distributed-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.source_exhausted is True + assert result.error is None + assert result.success_count == 5 + assert result.failure_count == 0 + assert result.total_count == 5 + assert result.distributed_map_run_arn.endswith("/distributed-map-run/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 run failed, so distributed_map must still return the summary carrying + the counts and let the caller opt into raising. + """ + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + operation = Operation( + operation_id="mr2", + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.FAILED, + distributed_map_details=DistributedMapDetails( + source_exhausted=False, + error=ErrorObject( + message="failure tolerance exceeded", + type="DistributedMapFailure", + data=None, + stack_trace=None, + ), + success_count=3, + failure_count=2, + unprocessed_count=1, + distributed_map_run_arn="arn:aws:lambda:us-east-1:123456789012:function:fn:$LATEST/durable-execution/exec1/invoke1/distributed-map-run/abc", + ), + ) + 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.source_exhausted is False + assert result.error is not None + assert result.error.message == "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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + 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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + 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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + 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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + 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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + 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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + 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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + 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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + 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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + ), + ) + 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 + + +def _rejecting_state(error: CheckpointError) -> Mock: + """A state whose START checkpoint is rejected with the given error.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_not_found() + ) + mock_state.create_checkpoint.side_effect = error + return mock_state + + +def _start(mock_state: Mock) -> DistributedMapSummary: + return distributed_map_handler( + source=["a"], + processor="p", + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("mrv"), + ) + + +def test_rejected_configuration_raises_validation_error(): + """A configuration the service rejects surfaces as ValidationError, as the spec requires.""" + error = CheckpointError( + "Source.MaxItemsToRead is out of range", + error_category=DurableApiErrorCategory.EXECUTION, + error={"Code": "InvalidParameterValueException", "Message": "bad config"}, + ) + with pytest.raises(ValidationError, match="MaxItemsToRead"): + _start(_rejecting_state(error)) + + +def test_stale_checkpoint_token_stays_a_checkpoint_error(): + """A retryable InvalidParameterValueException is a stale token, not a bad configuration.""" + error = CheckpointError( + "Invalid checkpoint token", + error_category=DurableApiErrorCategory.INVOCATION, + error={ + "Code": "InvalidParameterValueException", + "Message": "Invalid checkpoint token", + }, + ) + with pytest.raises(CheckpointError): + _start(_rejecting_state(error)) + + +def test_other_checkpoint_failures_stay_checkpoint_errors(): + """Only an invalid configuration is remapped; every other rejection passes through.""" + error = CheckpointError( + "throttled", + error_category=DurableApiErrorCategory.INVOCATION, + error={"Code": "TooManyRequestsException", "Message": "slow down"}, + ) + with pytest.raises(CheckpointError): + _start(_rejecting_state(error)) + + +@patch( + "aws_durable_execution_sdk_python.operation.dmap.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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + 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() + + +# Serialization and result-collection tests +# ============================================================================ + + +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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.STARTED, + ) + ) + return [not_found, started] + + +def test_processor_item_outcomes_sets_response_types(): + """report_item_outcomes serializes FunctionResponseTypes=REPORT_BATCH_ITEM_OUTCOMES.""" + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_item_outcomes("proc", batch_size=25), + max_concurrency=4, + config=DistributedMapConfig(), + ) + processor = options["Processor"] + assert processor["FunctionName"] == "proc" + assert processor["FunctionResponseTypes"] == ["REPORT_BATCH_ITEM_OUTCOMES"] + assert processor["BatchSize"] == 25 + + +def test_processor_batch_outcome_sets_no_response_types(): + """report_batch_outcome omits FunctionResponseTypes, selecting default batch mode.""" + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=4, + config=DistributedMapConfig(), + ) + assert "FunctionResponseTypes" not in options["Processor"] + + +def test_processor_unlimited_retries_maps_to_negative_one(): + """DistributedMapProcessor.UNLIMITED serializes to MaxRetryAttempts=-1.""" + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_item_outcomes( + "proc", + max_retry_attempts=DistributedMapProcessor.UNLIMITED, + max_retry_duration=Duration.from_hours(1), + ), + max_concurrency=1, + config=DistributedMapConfig(), + ) + processor = options["Processor"] + assert processor["FunctionResponseTypes"] == ["REPORT_BATCH_ITEM_OUTCOMES"] + 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", 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 + s3 = source["S3SourceConfig"] + assert s3["Type"] == "OBJECT" + assert s3["Bucket"] == "bucket" + assert s3["ObjectConfig"]["Key"] == "data.jsonl" + assert s3["ObjectConfig"]["Format"] == "JSON_LINES" + assert "ListedObjectMetadataConfig" not in s3 + assert "ListedObjectsConfig" not in s3 + + +def test_full_config_serializes_all_blocks(): + """Completion, destination, timeout, and result-collection blocks all serialize.""" + config = DistributedMapResultConfig( + completion_config=DistributedMapCompletionConfig.failure_percentage( + 5, minimum_sample_size=200 + ), + destination=DistributedMapDestination( + on_success=SuccessDestination.from_uri("s3://out/ok"), + on_failure=FailureDestination.from_uri("s3://out/bad"), + ), + timeout=Duration.from_minutes(30), + ) + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_item_outcomes("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_result_config_returns_map_run_result_with_items(): + """With a DistributedMapResultConfig, 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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + success_count=1, + failure_count=1, + unprocessed_count=0, + total_count=2, + results=( + DistributedMapResultItemApi( + item_id="0", status=DistributedMapItemStatus.SUCCEEDED, output="42" + ), + DistributedMapResultItemApi( + item_id="1", status=DistributedMapItemStatus.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_outcomes("proc"), + max_concurrency=2, + state=mock_state, + operation_identifier=_identifier("mrr"), + config=DistributedMapResultConfig(), + ) + + 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].type == "ItemError" + assert errors[0].message == "boom" + + +def test_item_error_keeps_every_error_member(): + """An item's error reaches the caller as the full ErrorObject, data and trace included.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="mre", + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + success_count=0, + failure_count=1, + unprocessed_count=0, + total_count=1, + results=( + DistributedMapResultItemApi( + item_id="0", + status=DistributedMapItemStatus.FAILED, + error=ErrorObject( + message="boom", + type="ItemError", + data="payload", + stack_trace=["frame one", "frame two"], + ), + ), + ), + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + result = distributed_map_handler( + source=["a"], + processor=DistributedMapProcessor.report_item_outcomes("proc"), + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("mre"), + config=DistributedMapResultConfig(), + ) + + error = result.get_errors()[0] + assert error is not None + assert error.data == "payload" + assert error.stack_trace == ["frame one", "frame two"] + + # The raised item failure carries the same members. + with pytest.raises(DistributedMapError) as exc: + result.throw_if_error() + assert exc.value.error_type == "ItemError" + assert exc.value.data == "payload" + assert exc.value.stack_trace == ["frame one", "frame two"] + + +def test_item_error_members_absent_from_the_backend_stay_none(): + """An omitted error member stays None rather than becoming an empty string.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="mren", + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + success_count=0, + failure_count=1, + unprocessed_count=0, + total_count=1, + results=( + DistributedMapResultItemApi( + item_id="0", + status=DistributedMapItemStatus.FAILED, + error=ErrorObject( + message=None, type=None, data=None, stack_trace=None + ), + ), + ), + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + result = distributed_map_handler( + source=["a"], + processor=DistributedMapProcessor.report_item_outcomes("proc"), + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("mren"), + config=DistributedMapResultConfig(), + ) + + error = result.get_errors()[0] + assert error is not None + assert error.type is None + assert error.message is None + + +def test_plain_config_returns_plain_summary(): + """With a plain DistributedMapConfig, a DistributedMapSummary is returned.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="mrs", + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + 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(): + """Explicit headers map to GIVEN, omitting them maps to FIRST_ROW.""" + # 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"]["ObjectConfig"]["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(), + ) + object_cfg = first_row["Source"]["S3SourceConfig"]["ObjectConfig"] + assert object_cfg["CsvFormatOptions"]["HeaderLocation"] == "FIRST_ROW" + assert "Headers" not in object_cfg["CsvFormatOptions"] + + +# Call-site validation and serdes tests +# ============================================================================ + + +def test_csv_delimiter_accepts_enum(): + opts = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.s3_csv( + "s3://b/data.csv", delimiter=DistributedMapCsvDelimiter.PIPE + ), + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + assert ( + opts["Source"]["S3SourceConfig"]["ObjectConfig"]["CsvFormatOptions"][ + "Delimiter" + ] + == "PIPE" + ) + + +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_and_capped(): + # typed initial_state is serialized into the opaque state string + options = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.reader_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_function( + "reader", initial_state="x" * 40_000 + ), + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + + +def test_inline_custom_serdes_applied_to_items(): + """A custom inline serdes transforms each item's serialized 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"'] + + +# Result helpers, from_dict, destinations, and source variants +# ============================================================================ + + +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 is DistributedMapSourceType.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=DistributedMapDestination( + on_success=SuccessDestination.from_uri("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=DistributedMapDestination( + on_failure=FailureDestination.from_uri("s3://out/bad") + ) + ), + ) + dest = opts["Destination"] + assert "OnFailure" in dest + assert "OnSuccess" not in dest + + +def test_s3_objects_source_config(): + 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["Type"] == "LISTED_OBJECT_METADATA" + assert s3["ListedObjectMetadataConfig"]["KeyPrefix"] == "prefix/" + assert "ObjectConfig" not in s3 + assert "ListedObjectsConfig" not in s3 + + +def test_s3_flattened_json_lines_source_config(): + 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["Type"] == "LISTED_OBJECTS" + assert s3["ListedObjectsConfig"]["KeyPrefix"] == "prefix/" + assert s3["ListedObjectsConfig"]["Format"] == "JSON_LINES" + + +# Validations, config branches, and 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 translation --- + + +def test_completion_config_percentage_and_sample_size_sent(): + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig( + completion_config=DistributedMapCompletionConfig.failure_percentage( + 25, minimum_sample_size=50 + ) + ), + ) + assert options["CompletionConfig"] == { + "ToleratedFailurePercentage": 25, + "MinimumSampleSize": 50, + } + + +# --- Source translation variants --- + + +def test_objects_whole_bucket_prefix_config(): + 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["Type"] == "LISTED_OBJECT_METADATA" + assert s3["ListedObjectMetadataConfig"]["KeyPrefix"] == "" + assert "ObjectConfig" not in s3 + + +def test_flattened_csv_source_config(): + 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(), + ) + listed = options["Source"]["S3SourceConfig"]["ListedObjectsConfig"] + assert listed["Format"] == "CSV" + assert listed["CsvFormatOptions"]["HeaderLocation"] == "GIVEN" + assert listed["CsvFormatOptions"]["Headers"] == ["a", "b"] + + +def test_reader_source_without_initial_state_omits_state(): + options = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.reader_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_non_json_serdes_supported(): + """Items are opaque strings, so a serdes need not produce JSON.""" + from aws_durable_execution_sdk_python.serdes import SerDes + + class _PlainTextSerDes(SerDes): + def serialize(self, value, _serdes_context): # noqa: ANN001, ANN201 + return f"item-{value}" + + def deserialize(self, data, _serdes_context): # noqa: ANN001, ANN201, ARG002 + return data + + options = _start_options( + _new_op_state_calls(), + source=DistributedMapSource.inline([1, 2], serdes=_PlainTextSerDes()), + processor=DistributedMapProcessor.report_batch_outcome("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + assert options["Source"]["InlineSourceConfig"]["Items"] == ["item-1", "item-2"] + + +# --- Destination translation permutations --- + + +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=DistributedMapDestination( + on_success=SuccessDestination.from_uri( + "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=DistributedMapDestination( + on_failure=FailureDestination.from_uri( + "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" + + +# --- Operation round-trips through the checkpoint --- + + +def test_operation_to_dict_round_trip_preserves_results(): + op = Operation( + operation_id="opx", + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + success_count=1, + failure_count=1, + unprocessed_count=0, + total_count=2, + distributed_map_run_arn="arn:aws:lambda:us-east-1:123456789012:function:fn:$LATEST/durable-execution/exec1/invoke1/distributed-map-run/z", + error=ErrorObject( + message="tolerated", + type="DistributedMapFailure", + data=None, + stack_trace=None, + ), + results=( + DistributedMapResultItemApi( + item_id="0", status=DistributedMapItemStatus.SUCCEEDED, output=5 + ), + DistributedMapResultItemApi( + item_id="1", + status=DistributedMapItemStatus.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("/distributed-map-run/z") + assert block["SourceExhausted"] is True + assert block["Error"]["ErrorType"] == "DistributedMapFailure" + assert block["TotalCount"] == 2 + + parsed = Operation.from_dict(op.to_dict()) + assert parsed.distributed_map_details is not None + assert parsed.distributed_map_details.source_exhausted is True + assert parsed.distributed_map_details.error.message == "tolerated" + 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_outcomes( + "proc", + batch_size=5, + max_retry_attempts=DistributedMapProcessor.UNLIMITED, + max_retry_duration=Duration.from_minutes(10), + durable_execution_name_prefix="pfx", + ), + max_concurrency=3, + config=DistributedMapResultConfig( + destination=DistributedMapDestination( + on_success=SuccessDestination.from_uri("s3://o/ok"), + on_failure=FailureDestination.from_uri("s3://o/bad"), + ), + completion_config=DistributedMapCompletionConfig.failure_count(2), + timeout=Duration.from_minutes(5), + ), + ) + parsed = DistributedMapOptions.from_dict(options_dict) + assert parsed.max_concurrency == 3 + assert parsed.source.source_type is DistributedMapSourceType.S3 + assert parsed.processor.function_response_types == [ + DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_OUTCOMES + ] + 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 is DistributedMapResultCollectionMode.INLINE + assert parsed.timeout_seconds == 300 + assert parsed.to_dict()["MaxConcurrency"] == 3 + + +# Remaining validation and config branches +# ============================================================================ + + +def test_csv_first_row_no_headers(): + 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"]["ObjectConfig"]["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(): + 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=DistributedMapDestination()), + ) + 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=DistributedMapDestination( + on_success=SuccessDestination.from_uri( + "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=DistributedMapDestination( + on_failure=FailureDestination.from_uri( + "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_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 is DistributedMapSourceType.READER_FUNCTION + assert parsed.source.reader_config is not None + assert parsed.source.reader_config.function_name == "reader" + + +def test_operation_to_dict_minimal_details_omits_optionals(): + op = Operation( + operation_id="opm", + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + success_count=1, + failure_count=0, + unprocessed_count=0, + ), + ) + block = op.to_dict()["DistributedMapDetails"] + 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_without_details_raises(status): + """A terminal failure carrying no details raises, since counts cannot be built.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="mrf", + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=status, + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + with pytest.raises(ExecutionError): + distributed_map_handler( + source=["a"], + processor="p", + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("mrf"), + ) + + +@pytest.mark.parametrize( + ("status", "expected"), + [ + (OperationStatus.FAILED, DistributedMapStatus.FAILED), + (OperationStatus.TIMED_OUT, DistributedMapStatus.TIMED_OUT), + (OperationStatus.STOPPED, DistributedMapStatus.STOPPED), + ], +) +def test_terminal_failure_with_details_resolves_with_summary(status, expected): + """A terminal non-success run resolves with its summary rather than raising.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="mrf", + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=status, + distributed_map_details=DistributedMapDetails( + source_exhausted=False, + success_count=1, + failure_count=2, + unprocessed_count=0, + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + summary = distributed_map_handler( + source=["a"], + processor="p", + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("mrf"), + ) + assert summary.status is expected + assert summary.failure_count == 2 + # The caller opts into raising rather than having it forced on them. + with pytest.raises(DistributedMapError): + summary.throw_if_error() + + +def test_terminal_without_source_exhausted_raises(): + """A terminal operation whose run exists but omits SourceExhausted is surfaced, not papered over.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="mrf", + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.STOPPED, + distributed_map_details=DistributedMapDetails( + distributed_map_run_arn="arn:aws:lambda:us-east-1:123456789012:function:fn:$LATEST/durable-execution/exec1/invoke1/distributed-map-run/abc", + success_count=0, + failure_count=0, + unprocessed_count=3, + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + with pytest.raises(ExecutionError, match="carried no SourceExhausted"): + distributed_map_handler( + source=["a"], + processor="p", + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("mrf"), + ) + + +# Serdes decode, falsy output, and duration-only retry +# ============================================================================ + + +def test_custom_result_serdes_applied_on_decode(): + """A custom result_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, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + success_count=1, + failure_count=0, + unprocessed_count=0, + results=( + DistributedMapResultItemApi( + item_id="0", + status=DistributedMapItemStatus.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_outcomes("proc"), + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("cs"), + config=DistributedMapResultConfig(result_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 the round-trip and decode (not dropped).""" + entry = DistributedMapResultItemApi( + item_id="0", + status=DistributedMapItemStatus.SUCCEEDED, + output=json.dumps(value), + ) + assert entry.to_dict()["Output"] == json.dumps(value) + + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="fo", + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + success_count=1, + failure_count=0, + unprocessed_count=0, + results=(entry,), + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + result = distributed_map_handler( + source=["a"], + processor=DistributedMapProcessor.report_item_outcomes("proc"), + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("fo"), + config=DistributedMapResultConfig(), + ) + 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", + max_retry_duration=Duration.from_minutes(5), + ), + max_concurrency=1, + config=DistributedMapConfig(), + ) + processor = options["Processor"] + assert processor["MaxRetryDurationSeconds"] == 300 + assert "MaxRetryAttempts" not in processor + + +def test_resolved_summary_uses_operation_status_not_details(): + """_resolve_summary carries the derived status when details omit it.""" + operation = Operation( + operation_id="dmap-id", + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + success_count=2, + failure_count=0, + unprocessed_count=0, + total_count=2, + ), + ) + executor = DistributedMapOperationExecutor( + source=["a", "b"], + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=1, + state=Mock(spec=ExecutionState), + operation_identifier=_identifier("dmap-id"), + config=DistributedMapConfig(), + ) + + summary = executor._resolve_summary(operation) + + assert summary.status is DistributedMapStatus.SUCCEEDED + assert summary.source_exhausted is True + assert summary.success_count == 2 + + +def test_checkpoint_from_a_different_map_is_rejected(): + """A checkpoint whose identity differs from the current call is not consumed.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="dmap-id", + operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="a_different_map", + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + source_exhausted=True, + success_count=999, + failure_count=0, + unprocessed_count=0, + total_count=999, + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + with pytest.raises(NonDeterministicExecutionError, match="name checkpoint="): + distributed_map_handler( + source=["a", "b"], + processor=DistributedMapProcessor.report_batch_outcome("proc"), + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("dmap-id"), + )