From 1ae73e8bf2dc59b81c18818e8e75b75c30f70cc7 Mon Sep 17 00:00:00 2001 From: nvasiu Date: Tue, 22 Sep 2026 20:39:27 +0000 Subject: [PATCH 1/3] feat(dmap): add distributed map operation Add ctx.distributed_map with inline, S3, and reader sources, the config, processor, completion, and destination types, the result types, and function-authoring helpers for item and batch handlers. Serialization lives in lambda_service.py with one dataclass per API shape. Translation happens only in the executor, so config.py has no runtime import from lambda_service. Tests are split by layer. Items, Output, and the record body are opaque strings, matching the API model, so any serdes works rather than JSON only. Every terminal state resolves with the summary, and throw_if_error() opts into raising. Missing details or a missing completion reason raise ExecutionError. The processor factories are batch, item_failures, and item_results. Source and destination factories are InlineSource, S3Source, ReaderSource, and S3Destination. Retry is max_retry_attempts and max_retry_duration on the processor. Passing DistributedMapResultConfig selects the DistributedMapResult return type statically. Only the result types and their enums are exported from the package root. --- .../__init__.py | 30 + .../config.py | 669 +++++++- .../context.py | 82 + .../dmap/__init__.py | 0 .../dmap/handlers.py | 453 ++++++ .../dmap/models.py | 127 ++ .../exceptions.py | 4 + .../lambda_service.py | 631 +++++++- .../operation/dmap.py | 582 +++++++ .../plugin.py | 1 + .../aws_durable_execution_sdk_python/state.py | 3 + .../tests/config_test.py | 171 ++ .../tests/context_test.py | 210 +++ .../tests/dmap/handlers_test.py | 253 +++ .../tests/dmap/models_test.py | 173 ++ .../tests/e2e/dmap_helpers_int_test.py | 146 ++ .../tests/e2e/dmap_int_test.py | 246 +++ .../tests/exceptions_test.py | 2 + .../tests/lambda_service_test.py | 96 ++ .../tests/operation/dmap_test.py | 1388 +++++++++++++++++ 20 files changed, 5264 insertions(+), 3 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dmap/__init__.py create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dmap/handlers.py create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dmap/models.py create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dmap.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/dmap/handlers_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/dmap/models_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/dmap_helpers_int_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/dmap_int_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dmap_test.py diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py index 69f7bb164..bd0199539 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,6 +16,9 @@ CompletionItemStatus, CompletionOutcome, CompletionStatus, + DistributedMapCompletionReason, + DistributedMapItemStatus, + DistributedMapStatus, ParallelBranch, complete_batch, continue_batch, @@ -27,6 +30,19 @@ 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 ( + DistributedMapItemError, + DistributedMapResult, + DistributedMapResultItem, + DistributedMapSummary, +) # Most common exceptions - users need to handle these exceptions from aws_durable_execution_sdk_python.exceptions import ( @@ -35,6 +51,7 @@ CallbackSubmitterError, CallbackTimeoutError, ChildContextError, + DistributedMapError, DurableExecutionsError, DurableOperationError, ExecutionError, @@ -69,6 +86,14 @@ "CompletionItemStatus", "CompletionOutcome", "CompletionStatus", + "DistributedMapCompletionReason", + "DistributedMapError", + "DistributedMapItemError", + "DistributedMapItemStatus", + "DistributedMapResult", + "DistributedMapResultItem", + "DistributedMapStatus", + "DistributedMapSummary", "DurableContext", "DurableExecutionsError", "DurableOperationError", @@ -87,6 +112,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..111af6309 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,71 @@ 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" + + +class DistributedMapCompletionReason(Enum): + ALL_COMPLETED = "ALL_COMPLETED" + ITEM_LIMIT_REACHED = "ITEM_LIMIT_REACHED" + STOPPED = "STOPPED" + TIMED_OUT = "TIMED_OUT" + FAILURE_TOLERANCE_EXCEEDED = "FAILURE_TOLERANCE_EXCEEDED" + SOURCE_FAILED = "SOURCE_FAILED" + DESTINATION_FAILED = "DESTINATION_FAILED" + INLINE_RESULT_LIMIT_EXCEEDED = "INLINE_RESULT_LIMIT_EXCEEDED" + INVALID_CONFIGURATION = "INVALID_CONFIGURATION" + QUOTA_EXCEEDED = "QUOTA_EXCEEDED" + KMS_ACCESS_DENIED = "KMS_ACCESS_DENIED" + INTERNAL_ERROR = "INTERNAL_ERROR" + UNKNOWN_TO_SDK_VERSION = "UNKNOWN_TO_SDK_VERSION" + + +class DistributedMapSourceFormat(Enum): + JSON_LINES = "JSON_LINES" + JSON_ARRAY = "JSON_ARRAY" + CSV = "CSV" + + +class ProcessorResponseMode(Enum): + """How a processor function reports per-item outcomes.""" + + BATCH = "BATCH" + ITEM_FAILURES = "ITEM_FAILURES" + ITEM_RESULTS = "ITEM_RESULTS" + + +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 +544,617 @@ 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) + + +def _validate_columns(name: str, columns: tuple[str, ...] | None) -> None: + """Validate a CSV columns/headers tuple is non-empty with no duplicates.""" + if columns is None: + return + if len(columns) == 0: + msg = f"{name} must be non-empty" + raise ValidationError(msg) + if len(set(columns)) != len(columns): + msg = f"{name} must not contain duplicates" + raise ValidationError(msg) + + +# 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: float | None = None + minimum_sample_size: int | None = None + + def __post_init__(self) -> None: + if ( + self.tolerated_failure_count is not None + and self.tolerated_failure_percentage is not None + ): + msg = ( + "tolerated_failure_count and tolerated_failure_percentage " + "are mutually exclusive" + ) + raise ValidationError(msg) + if ( + self.minimum_sample_size is not None + and self.tolerated_failure_percentage is None + ): + msg = "minimum_sample_size is only valid with tolerated_failure_percentage" + raise ValidationError(msg) + if ( + self.tolerated_failure_count is not None + and self.tolerated_failure_count < 0 + ): + msg = ( + "tolerated_failure_count must be non-negative, got: " + f"{self.tolerated_failure_count}" + ) + raise ValidationError(msg) + if self.tolerated_failure_percentage is not None and not ( + 0 <= self.tolerated_failure_percentage <= 100 # noqa: PLR2004 + ): + msg = ( + "tolerated_failure_percentage must be between 0 and 100, got: " + f"{self.tolerated_failure_percentage}" + ) + raise ValidationError(msg) + if self.minimum_sample_size is not None and self.minimum_sample_size < 1: + msg = ( + "minimum_sample_size must be at least 1, got: " + f"{self.minimum_sample_size}" + ) + raise ValidationError(msg) + + @staticmethod + def failure_count(count: int) -> DistributedMapCompletionConfig: + """Abort once this many items have permanently failed.""" + return DistributedMapCompletionConfig(tolerated_failure_count=count) + + @staticmethod + def failure_percentage( + percentage: float, *, minimum_sample_size: int | None = None + ) -> DistributedMapCompletionConfig: + """Abort once the failure rate exceeds this percentage.""" + return DistributedMapCompletionConfig( + tolerated_failure_percentage=percentage, + minimum_sample_size=minimum_sample_size, + ) + + +@dataclass(frozen=True) +class 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) + if isinstance(self.max_retry_attempts, int) and self.max_retry_attempts < 0: + msg = ( + "max_retry_attempts must be non-negative or " + "DistributedMapProcessor.UNLIMITED, " + f"got: {self.max_retry_attempts}" + ) + raise ValidationError(msg) + if self.max_retry_duration is not None and not ( + 60 <= self.max_retry_duration.to_seconds() <= 21600 # noqa: PLR2004 + ): + msg = ( + "max_retry_duration must be between 1 minute and 6 hours, got: " + f"{self.max_retry_duration.to_seconds()}s" + ) + raise ValidationError(msg) + 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 batch( + 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 a single pass/fail outcome for the whole batch, with no per-item results.""" + 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 item_failures( + 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 the ids of failed items, with all others marked succeeded.""" + return cls( + function_name=name, + response_mode=ProcessorResponseMode.ITEM_FAILURES, + 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 item_results( + 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 the results (output or error) for every item.""" + return cls( + function_name=name, + response_mode=ProcessorResponseMode.ITEM_RESULTS, + 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 S3SourceConfig: + """Resolved S3 source configuration.""" + + bucket: str + key: str | None = None + prefix: str | None = None + fmt: DistributedMapSourceFormat | None = None + delimiter: DistributedMapCsvDelimiter | None = None + headers: tuple[str, ...] | None = None + expected_bucket_owner: str | None = None + + def __post_init__(self) -> None: + _validate_bucket_owner(self.expected_bucket_owner) + _validate_columns("headers", self.headers) + + +@dataclass(frozen=True) +class ReaderSourceConfig: + """Resolved reader-function source configuration.""" + + function_name: str + initial_state: Any = None + state_serdes: SerDes | None = None # None = DEFAULT_JSON_SERDES + + def __post_init__(self) -> None: + _validate_function_name(self.function_name) + + +@dataclass(frozen=True) +class DistributedMapSource: + """Source configuration for a map run.""" + + max_items: int | None = None + inline_items: tuple[Any, ...] | None = None + inline_serdes: SerDes | None = None # None = DEFAULT_JSON_SERDES + s3: S3SourceConfig | None = None + reader: ReaderSourceConfig | None = None + + def __post_init__(self) -> None: + if self.max_items is not None and self.max_items < 1: + msg = f"max_items must be at least 1, got: {self.max_items}" + raise ValidationError(msg) + 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) + + +class InlineSource: + """Inline source factory.""" + + @staticmethod + def of( + 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 DistributedMapSource( + inline_items=tuple(items), + inline_serdes=serdes, + max_items=max_items, + ) + + +class S3Source: + """S3 source factories.""" + + @staticmethod + def json_lines( + uri: str, + *, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read a single object, treating each line as an item.""" + parsed_uri = S3Uri.parse(uri) + bucket, key = parsed_uri.bucket, parsed_uri.path + if key is None: + msg = "json_lines requires an S3 object key" + raise ValidationError(msg) + return DistributedMapSource( + max_items=max_items, + s3=S3SourceConfig( + bucket=bucket, + key=key, + fmt=DistributedMapSourceFormat.JSON_LINES, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @staticmethod + def json_array( + 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 = "json_array requires an S3 object key" + raise ValidationError(msg) + return DistributedMapSource( + max_items=max_items, + s3=S3SourceConfig( + bucket=bucket, + key=key, + fmt=DistributedMapSourceFormat.JSON_ARRAY, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @staticmethod + def csv( + 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 = "csv requires an S3 object key" + raise ValidationError(msg) + return DistributedMapSource( + max_items=max_items, + s3=S3SourceConfig( + bucket=bucket, + key=key, + fmt=DistributedMapSourceFormat.CSV, + delimiter=_parse_delimiter(delimiter), + headers=tuple(headers) if headers is not None else None, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @staticmethod + def objects( + prefix_uri: str, + *, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read each object under a prefix as one item.""" + parsed_uri = S3Uri.parse(prefix_uri) + bucket, prefix = parsed_uri.bucket, parsed_uri.path + return DistributedMapSource( + max_items=max_items, + s3=S3SourceConfig( + bucket=bucket, + prefix=prefix or "", + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @staticmethod + def flattened_json_lines( + prefix_uri: str, + *, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read a prefix, flattening each object's lines into items.""" + parsed_uri = S3Uri.parse(prefix_uri) + bucket, prefix = parsed_uri.bucket, parsed_uri.path + return DistributedMapSource( + max_items=max_items, + s3=S3SourceConfig( + bucket=bucket, + prefix=prefix or "", + fmt=DistributedMapSourceFormat.JSON_LINES, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @staticmethod + def flattened_json_array( + prefix_uri: str, + *, + expected_bucket_owner: str | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Read a prefix, flattening each object's JSON array elements into items.""" + parsed_uri = S3Uri.parse(prefix_uri) + bucket, prefix = parsed_uri.bucket, parsed_uri.path + return DistributedMapSource( + max_items=max_items, + s3=S3SourceConfig( + bucket=bucket, + prefix=prefix or "", + fmt=DistributedMapSourceFormat.JSON_ARRAY, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + @staticmethod + def flattened_csv( + 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 a prefix, flattening each object's records into items.""" + parsed_uri = S3Uri.parse(prefix_uri) + bucket, prefix = parsed_uri.bucket, parsed_uri.path + return DistributedMapSource( + max_items=max_items, + s3=S3SourceConfig( + bucket=bucket, + prefix=prefix or "", + fmt=DistributedMapSourceFormat.CSV, + delimiter=_parse_delimiter(delimiter), + headers=tuple(headers) if headers is not None else None, + expected_bucket_owner=expected_bucket_owner, + ), + ) + + +class ReaderSource: + """Reader-function source factories.""" + + @staticmethod + def from_function( + name: str, + *, + initial_state: Any = None, + state_serdes: SerDes | None = None, + max_items: int | None = None, + ) -> DistributedMapSource: + """Page items from a customer-supplied reader Lambda function.""" + return DistributedMapSource( + max_items=max_items, + reader=ReaderSourceConfig( + function_name=name, + initial_state=initial_state, + state_serdes=state_serdes, + ), + ) + + +@dataclass(frozen=True) +class SuccessDestination: + """S3 destination for succeeded items.""" + + bucket: str + prefix: str + include_input: bool = False + include_output: bool = True + expected_bucket_owner: str | None = None + + def __post_init__(self) -> None: + _validate_bucket_owner(self.expected_bucket_owner) + if not (self.include_input or self.include_output): + msg = "success destination must include input or output" + raise ValidationError(msg) + + +@dataclass(frozen=True) +class FailureDestination: + """S3 destination for permanently-failed items.""" + + bucket: str + prefix: str + include_input: bool = True + include_error: bool = True + expected_bucket_owner: str | None = None + + def __post_init__(self) -> None: + _validate_bucket_owner(self.expected_bucket_owner) + if not (self.include_input or self.include_error): + msg = "failure destination must include input or error" + raise ValidationError(msg) + + +@dataclass(frozen=True) +class DistributedMapDestinationConfig: + """Destination routing for map run results.""" + + on_success: SuccessDestination | None = None + on_failure: FailureDestination | None = None + + +class S3Destination: + """S3 destination factories.""" + + @staticmethod + def successes( + prefix_uri: str, + *, + include_input: bool = False, + include_output: bool = True, + expected_bucket_owner: str | None = None, + ) -> SuccessDestination: + """Route succeeded item records to an S3 prefix.""" + parsed_uri = S3Uri.parse(prefix_uri) + bucket, prefix = parsed_uri.bucket, parsed_uri.path + return SuccessDestination( + bucket=bucket, + prefix=prefix or "", + include_input=include_input, + include_output=include_output, + expected_bucket_owner=expected_bucket_owner, + ) + + @staticmethod + def failures( + prefix_uri: str, + *, + include_input: bool = True, + include_error: bool = True, + expected_bucket_owner: str | None = None, + ) -> FailureDestination: + """Route permanently-failed item records to an S3 prefix.""" + parsed_uri = S3Uri.parse(prefix_uri) + bucket, prefix = parsed_uri.bucket, parsed_uri.path + return FailureDestination( + bucket=bucket, + prefix=prefix or "", + include_input=include_input, + include_error=include_error, + expected_bucket_owner=expected_bucket_owner, + ) + + +@dataclass(frozen=True) +class DistributedMapConfig: + """Configuration for map run operations.""" + + destination: DistributedMapDestinationConfig | None = None + completion_config: DistributedMapCompletionConfig | None = None + timeout: Duration | None = None + + 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..e7de3162b 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 @@ -102,6 +114,8 @@ PASS_THROUGH_SERDES: SerDes[Any] = PassThroughSerDes() +_MAX_CONCURRENCY_LIMIT = 10000 + @dataclass(frozen=True) class ExecutionContext: @@ -683,6 +697,74 @@ 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 + 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 not 1 <= max_concurrency <= _MAX_CONCURRENCY_LIMIT: + msg = ( + f"max_concurrency must be between 1 and {_MAX_CONCURRENCY_LIMIT}, " + f"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..952744f71 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dmap/handlers.py @@ -0,0 +1,453 @@ +"""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, Literal + +from aws_durable_execution_sdk_python.concurrency.models import BatchItemStatus +from aws_durable_execution_sdk_python.config import CompletionConfig, MapConfig +from aws_durable_execution_sdk_python.exceptions import ValidationError +from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python.serdes import ( + DEFAULT_JSON_SERDES, + SerDes, + SerDesContext, +) + +_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: tuple[ProcessorRecord, ...] = () + + @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=tuple(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_results processor.""" + + item_identifier: str + output: Any = 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]: + return {"itemIdentifier": self.item_identifier, "output": self.output} + + +@dataclass(frozen=True) +class ItemFailure: + """One item's failure, reported by an item_failures or item_results processor.""" + + item_identifier: str + error_type: str = "" + error_message: str = "" + + @classmethod + def from_exception(cls, item_identifier: str, exc: BaseException) -> ItemFailure: + return cls( + item_identifier=item_identifier, + error_type=type(exc).__name__, + error_message=str(exc), + ) + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> ItemFailure: + error = data.get("error") or {} + return cls( + item_identifier=data.get("itemIdentifier", ""), + error_type=error.get("errorType", ""), + error_message=error.get("errorMessage", ""), + ) + + def to_dict(self) -> MutableMapping[str, Any]: + return { + "itemIdentifier": self.item_identifier, + "error": { + "errorType": self.error_type, + "errorMessage": self.error_message, + }, + } + + +@dataclass(frozen=True) +class ItemHandlerResponse: + """What an item handler returns. A None ``results`` reports failures only.""" + + failures: tuple[ItemFailure, ...] = () + results: tuple[ItemResult, ...] | None = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> ItemHandlerResponse: + results = data.get("batchItemResults") + return cls( + failures=tuple( + ItemFailure.from_dict(f) for f in data.get("batchItemFailures", []) + ), + results=tuple(ItemResult.from_dict(r) for r in results) + if results is not None + else None, + ) + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = { + "batchItemFailures": [f.to_dict() for f in self.failures] + } + if self.results is not None: + result["batchItemResults"] = [r.to_dict() for r in self.results] + return result + + +@dataclass(frozen=True) +class ReaderEvent: + """The event a reader function is invoked with.""" + + max_items: int + state: str | None = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> ReaderEvent: + max_items = data.get("maxItems") + if not isinstance(max_items, int): + msg = ( + "expected a distributed map reader envelope with an integer 'maxItems'" + ) + raise ValidationError(msg) + return cls(max_items=max_items, state=data.get("state")) + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {"maxItems": self.max_items} + 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: tuple[Any, ...] = () + next_state: str | None = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> ReaderResponse: + return cls(items=tuple(data.get("items", ())), next_state=data.get("nextState")) + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {"items": list(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: + """Serialize an item's output for the Output field.""" + return serdes.serialize(value, ctx) + + +def _validate_report(report: str) -> None: + if report not in ("results", "failures"): + msg = f"report must be 'results' or 'failures', got: {report!r}" + raise ValidationError(msg) + + +def _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, + report: Literal["results", "failures"] = "results", +) -> Callable[..., Any]: + """Wrap a Lambda to be used as an item_results processor. + + Pass ``report="failures"`` for an item_failures 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``. + + 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, + report=report, + ) + _validate_report(report) + _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) -> Any: + return func(_to_item(serdes, record.body, ctx)) + + outputs: list[Any] = [None] * len(records) + errors: list[BaseException | 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)) + elif report == "results": + results.append( + ItemResult( + item_identifier=record.item_id, + output=_to_output(out_serdes, outputs[i], ctx), + ) + ) + + response = ItemHandlerResponse( + failures=tuple(failures), + results=tuple(results) if report == "results" else None, + ) + return dict(response.to_dict()) + + return handler + + +def distributed_map_batch_handler( + func: Callable[[list[Any]], Any] | 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``. + """ + 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) -> Any: + ctx = SerDesContext() + return func( + [ + _to_item(serdes, record.body, ctx) + for record in ProcessorEvent.from_dict(event).records + ] + ) + + return handler + + +def distributed_map_reader( + func: Callable[[Any], ReaderPage] | 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 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``. + """ + if func is None: + return functools.partial(distributed_map_reader, state_serdes=state_serdes) + serdes = 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 = ( + serdes.deserialize(reader_event.state, ctx) + if reader_event.state is not None + else None + ) + + page = func(state) + if len(page.items) > reader_event.max_items: + msg = ( + f"reader returned {len(page.items)} items, exceeding maxItems " + f"{reader_event.max_items}" + ) + raise ValidationError(msg) + + next_state: str | None = None + if page.next_state is not None: + next_state = serdes.serialize(page.next_state, ctx) + if len(next_state.encode("utf-8")) > _READER_STATE_LIMIT: + msg = f"reader next_state exceeds the {_READER_STATE_LIMIT // 1024} KB limit" + raise ValidationError(msg) + + return dict( + ReaderResponse(items=tuple(page.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, + report: Literal["results", "failures"] = "results", +) -> 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``. + """ + if func is None: + return functools.partial( + durable_distributed_map_item_handler, + item_serdes=item_serdes, + result_serdes=result_serdes, + report=report, + ) + _validate_report(report) + 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(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: + if report == "results": + results.append( + ItemResult( + item_identifier=item_id, + output=_to_output(out_serdes, bi.result, ctx), + ) + ) + else: + err = bi.error + failures.append( + ItemFailure( + item_identifier=item_id, + error_type=(err.type or "") if err else "", + error_message=(err.message or "") if err else "", + ) + ) + + response = ItemHandlerResponse( + failures=tuple(failures), + results=tuple(results) if report == "results" else None, + ) + return dict(response.to_dict()) + + return handler + + +def durable_distributed_map_batch_handler( + func: Callable[..., Any] | 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``. + """ + 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) -> Any: + ctx = SerDesContext() + return 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..7d9ad80c6 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dmap/models.py @@ -0,0 +1,127 @@ +"""Result models returned by ``ctx.distributed_map``.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from aws_durable_execution_sdk_python.exceptions import ( + DistributedMapError, +) +from aws_durable_execution_sdk_python.config import ( + DistributedMapCompletionReason, + DistributedMapItemStatus, + DistributedMapStatus, +) + + +@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 + completion_reason: DistributedMapCompletionReason + success_count: int + failure_count: int + unprocessed_count: int + distributed_map_run_arn: str | None = None + completion_details: str | None = None + total_count: int | None = None + + @property + def has_failure(self) -> bool: + """``True`` when any item permanently failed.""" + return self.failure_count > 0 + + def throw_if_error(self) -> None: + """Raise :class:`DistributedMapError` on any non-success outcome.""" + if self.status is not DistributedMapStatus.SUCCEEDED: + detail = f", {self.completion_details}" if self.completion_details else "" + msg = ( + f"Map run ended {self.status.value} " + f"(reason: {self.completion_reason.value}{detail})" + ) + raise DistributedMapError(msg) + if self.failure_count > 0: + msg = ( + f"Map run succeeded but {self.failure_count} item(s) permanently failed" + ) + raise DistributedMapError(msg) + + +@dataclass(frozen=True) +class DistributedMapItemError: + """Error for a single failed map run item.""" + + error_type: str + error_message: str + + +@dataclass(frozen=True) +class DistributedMapResultItem: + """Outcome of a single map run item.""" + + item_id: str + status: DistributedMapItemStatus + output: Any | None = None + error: DistributedMapItemError | None = None + + +@dataclass(frozen=True) +class DistributedMapResult(DistributedMapSummary): + """Outcome of a map run with per-item results. + + Resolved by ``ctx.distributed_map`` when passed a + :class:`DistributedMapResultConfig`. Extends :class:`DistributedMapSummary` + with the retained per-item results. + """ + + all: list[DistributedMapResultItem] = field(default_factory=list) + + def succeeded(self) -> list[DistributedMapResultItem]: + """Return the items that succeeded.""" + return [ + item + for item in self.all + if item.status 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 the outputs of the succeeded items.""" + return [ + item.output + for item in self.all + if item.status is DistributedMapItemStatus.SUCCEEDED + and item.output is not None + ] + + def get_errors(self) -> list[DistributedMapItemError]: + """Return the errors of the failed items.""" + return [ + item.error + for item in self.all + if item.status is DistributedMapItemStatus.FAILED and item.error is not None + ] + + def throw_if_error(self) -> None: + """Raise the first failed item's error, otherwise defer to the summary rule.""" + if self.status is DistributedMapStatus.SUCCEEDED and self.failure_count > 0: + failed = self.failed() + if failed: + first = failed[0] + if first.error is not None: + msg = f"{first.error.error_type}: {first.error.error_message}" + raise DistributedMapError(msg) + msg = f"item {first.item_id} failed" + raise DistributedMapError(msg) + super().throw_if_error() 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..a51568d16 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,31 @@ 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 ( + DistributedMapCompletionReason, + DistributedMapCsvDelimiter, + DistributedMapItemStatus, + DistributedMapSourceFormat, +) from aws_durable_execution_sdk_python.exceptions import ( CheckpointError, DurableOperationError, + ExecutionError, GetExecutionStateError, SerDesError, ) @@ -74,6 +90,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 +103,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 +123,42 @@ class CallbackTimeoutType(Enum): HEARTBEAT = "Callback.Heartbeat" +class DistributedMapSourceType(Enum): + INLINE = "INLINE" + S3 = "S3" + READER_FUNCTION = "READER_FUNCTION" + + +class DistributedMapS3SourceTransform(Enum): + NONE = "NONE" + LOAD_AND_FLATTEN = "LOAD_AND_FLATTEN" + + +class DistributedMapCsvHeaderLocation(Enum): + FIRST_ROW = "FIRST_ROW" + GIVEN = "GIVEN" + + +class DistributedMapFunctionResponseType(Enum): + REPORT_BATCH_ITEM_FAILURES = "REPORT_BATCH_ITEM_FAILURES" + REPORT_BATCH_ITEM_RESULTS = "REPORT_BATCH_ITEM_RESULTS" + + +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 +171,7 @@ class OperationSubType(Enum): WAIT_FOR_CALLBACK = "WaitForCallback" WAIT_FOR_CONDITION = "WaitForCondition" CHAINED_INVOKE = "ChainedInvoke" + DISTRIBUTED_MAP = "DistributedMap" class InvocationStatus(Enum): @@ -383,6 +439,101 @@ 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") + return cls( + item_id=data.get("ItemId", ""), + status=_parse_enum( + DistributedMapItemStatus, data.get("Status", ""), "item status" + ), + output=data.get("Output"), + error=ErrorObject.from_dict(error_raw) if error_raw else None, + ) + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = { + "ItemId": self.item_id, + "Status": self.status.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: + completion_reason: DistributedMapCompletionReason | None = None + distributed_map_run_arn: str | None = None + completion_details: str | None = None + total_count: int | None = None + success_count: int = 0 + failure_count: int = 0 + unprocessed_count: int = 0 + results: tuple[DistributedMapResultItem, ...] | None = None + + @classmethod + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapDetails: + results_raw = data.get("Results") + reason_raw = data.get("CompletionReason") + # CompletionReason is absent while a run is in flight, so parse it only + # when present rather than requiring it. + return cls( + completion_reason=( + _parse_enum( + DistributedMapCompletionReason, + reason_raw, + "completion reason", + unknown_fallback=DistributedMapCompletionReason.UNKNOWN_TO_SDK_VERSION, + ) + if reason_raw is not None + else None + ), + distributed_map_run_arn=data.get("DistributedMapRunArn"), + completion_details=data.get("CompletionDetails"), + total_count=data.get("TotalCount"), + success_count=data.get("SuccessCount", 0), + failure_count=data.get("FailureCount", 0), + unprocessed_count=data.get("UnprocessedCount", 0), + results=tuple( + 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 +629,419 @@ 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: tuple[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"] = list(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=tuple(headers) if headers is not None else None, + delimiter=DistributedMapCsvDelimiter(delimiter) + if delimiter is not None + else None, + ) + + +@dataclass(frozen=True) +class DistributedMapS3SourceConfig: + """Represent an S3 distributed map source config.""" + + bucket: str + key: str | None = None + key_prefix: str | None = None + transform: DistributedMapS3SourceTransform | None = None + expected_bucket_owner: str | None = None + fmt: DistributedMapSourceFormat | None = None + csv_format_options: DistributedMapCsvFormatOptions | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {"Bucket": self.bucket} + if self.key is not None: + result["Key"] = self.key + if self.key_prefix is not None: + result["KeyPrefix"] = self.key_prefix + if self.transform is not None: + result["Transform"] = self.transform.value + if self.expected_bucket_owner is not None: + result["ExpectedBucketOwner"] = self.expected_bucket_owner + if self.fmt is not None: + result["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]) -> DistributedMapS3SourceConfig: + transform = data.get("Transform") + fmt = data.get("Format") + csv_raw = data.get("CsvFormatOptions") + return cls( + bucket=data.get("Bucket", ""), + key=data.get("Key"), + key_prefix=data.get("KeyPrefix"), + transform=DistributedMapS3SourceTransform(transform) + if transform is not None + else None, + expected_bucket_owner=data.get("ExpectedBucketOwner"), + fmt=DistributedMapSourceFormat(fmt) if fmt is not None else None, + csv_format_options=DistributedMapCsvFormatOptions.from_dict(csv_raw) + if csv_raw is not None + else None, + ) + + +@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: tuple[str, ...] = () + + def to_dict(self) -> MutableMapping[str, Any]: + return {"Items": list(self.items)} + + @classmethod + def from_dict( + cls, data: MutableMapping[str, Any] + ) -> DistributedMapInlineSourceConfig: + return cls(items=tuple(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, + ) + + +@dataclass(frozen=True) +class DistributedMapProcessorConfig: + """Represent a map run processor.""" + + function_name: str + function_response_types: tuple[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=tuple( + 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"), + ) + + +@dataclass(frozen=True) +class DistributedMapCompletionConfig: + """Represent a map run completion (failure-tolerance) config.""" + + tolerated_failure_count: int | None = None + tolerated_failure_percentage: float | None = None + minimum_sample_size: int | None = None + + def to_dict(self) -> MutableMapping[str, Any]: + result: MutableMapping[str, Any] = {} + if self.tolerated_failure_count is not None: + result["ToleratedFailureCount"] = self.tolerated_failure_count + if self.tolerated_failure_percentage is not None: + result["ToleratedFailurePercentage"] = self.tolerated_failure_percentage + if self.minimum_sample_size is not None: + result["MinimumSampleSize"] = self.minimum_sample_size + return result + + @classmethod + def from_dict( + cls, data: MutableMapping[str, Any] + ) -> DistributedMapCompletionConfig: + return cls( + tolerated_failure_count=data.get("ToleratedFailureCount"), + tolerated_failure_percentage=data.get("ToleratedFailurePercentage"), + minimum_sample_size=data.get("MinimumSampleSize"), + ) + + +@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: tuple[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=tuple( + DistributedMapDestinationInclude(i) for i in data.get("Include", []) + ), + s3_destination_config=DistributedMapS3DestinationConfig.from_dict(s3_raw), + ) + + +@dataclass(frozen=True) +class DistributedMapOnSuccessConfig(_DistributedMapDestinationEntry): + """Represent where a map run writes the outcome of its succeeded items.""" + + +@dataclass(frozen=True) +class DistributedMapOnFailureConfig(_DistributedMapDestinationEntry): + """Represent where a map run writes the outcome of its failed items.""" + + +@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, + ) + + +@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 +1074,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 +1103,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 +1133,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 +1153,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 +1337,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 +1480,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 +1525,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 +1546,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 +1611,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.completion_reason is not None: + distributed_map_details_dict["CompletionReason"] = ( + self.distributed_map_details.completion_reason.value + ) + if self.distributed_map_details.distributed_map_run_arn: + distributed_map_details_dict["DistributedMapRunArn"] = ( + self.distributed_map_details.distributed_map_run_arn + ) + if self.distributed_map_details.completion_details: + distributed_map_details_dict["CompletionDetails"] = ( + self.distributed_map_details.completion_details + ) + if self.distributed_map_details.total_count is not None: + distributed_map_details_dict["TotalCount"] = ( + self.distributed_map_details.total_count + ) + if self.distributed_map_details.results is not None: + distributed_map_details_dict["Results"] = [ + item.to_dict() for item in self.distributed_map_details.results + ] + result["DistributedMapDetails"] = distributed_map_details_dict return result def to_json_dict(self) -> MutableMapping[str, Any]: diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dmap.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dmap.py new file mode 100644 index 000000000..33ed11e15 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dmap.py @@ -0,0 +1,582 @@ +"""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 ( + DistributedMapProcessor, + DistributedMapResultConfig, + DistributedMapSource, + DistributedMapSourceFormat, + DistributedMapStatus, + ProcessorResponseMode, +) +from aws_durable_execution_sdk_python.dmap.models import ( + DistributedMapItemError, + DistributedMapResult, + DistributedMapResultItem, + DistributedMapSummary, +) +from aws_durable_execution_sdk_python.exceptions import ( + ExecutionError, + ValidationError, +) +from aws_durable_execution_sdk_python.lambda_service import ( + DistributedMapCompletionConfig, + DistributedMapCsvFormatOptions, + DistributedMapCsvHeaderLocation, + DistributedMapDestinationConfig, + DistributedMapDestinationInclude, + DistributedMapDestinationType, + DistributedMapFunctionResponseType, + DistributedMapInlineSourceConfig, + DistributedMapOnFailureConfig, + DistributedMapOnSuccessConfig, + DistributedMapOptions, + DistributedMapProcessorConfig, + DistributedMapReaderFunctionSourceConfig, + DistributedMapResultCollectionConfig, + DistributedMapResultCollectionMode, + DistributedMapS3DestinationConfig, + DistributedMapS3SourceConfig, + DistributedMapS3SourceTransform, + DistributedMapSourceConfig, + DistributedMapSourceType, + OperationStatus, + 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, + ) + 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.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 +_UNLIMITED_RETRY_WIRE = -1 + +_RESPONSE_TYPE_FOR_MODE = { + ProcessorResponseMode.ITEM_FAILURES: DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_FAILURES, + ProcessorResponseMode.ITEM_RESULTS: DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_RESULTS, +} + + +def _build_inline_items( + items: tuple[Any, ...], + serdes: Any, + operation_id: str, + durable_execution_arn: str, +) -> tuple[str, ...]: + """Serialize each inline item, enforcing the 1 MB cap on the whole list.""" + serialized_items: list[str] = [] + for item in items: + serialized_items.append( + serialize( + serdes=serdes, + value=item, + operation_id=operation_id, + durable_execution_arn=durable_execution_arn, + ) + ) + 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 tuple(serialized_items) + + +def _build_completion_config( + completion_config: Any, +) -> DistributedMapCompletionConfig | None: + """Translate the completion config, or None when it carries no threshold.""" + if ( + completion_config.tolerated_failure_count is None + and completion_config.tolerated_failure_percentage is None + and completion_config.minimum_sample_size is None + ): + return None + return DistributedMapCompletionConfig( + tolerated_failure_count=completion_config.tolerated_failure_count, + tolerated_failure_percentage=completion_config.tolerated_failure_percentage, + minimum_sample_size=completion_config.minimum_sample_size, + ) + + +def _build_processor_config( + processor: DistributedMapProcessor, +) -> DistributedMapProcessorConfig: + """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 + max_retry_duration_seconds: 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 + if processor.max_retry_duration is not None: + max_retry_duration_seconds = processor.max_retry_duration.to_seconds() + return DistributedMapProcessorConfig( + 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=max_retry_duration_seconds, + durable_execution_name_prefix=processor.durable_execution_name_prefix, + ) + + +def _transform_for(s3: Any) -> DistributedMapS3SourceTransform | None: + """A prefix source flattens when a format is set, and lists keys when it is not.""" + if s3.prefix is None: + return None + if s3.fmt is None: + return DistributedMapS3SourceTransform.NONE + return DistributedMapS3SourceTransform.LOAD_AND_FLATTEN + + +def _build_s3_source_config(s3: Any) -> DistributedMapS3SourceConfig: + """Translate the S3 source, deriving the transform and CSV header location.""" + csv_format_options: DistributedMapCsvFormatOptions | None = None + if s3.fmt is DistributedMapSourceFormat.CSV: + csv_format_options = DistributedMapCsvFormatOptions( + header_location=( + DistributedMapCsvHeaderLocation.GIVEN + if s3.headers is not None + else DistributedMapCsvHeaderLocation.FIRST_ROW + ), + headers=s3.headers, + delimiter=s3.delimiter, + ) + return DistributedMapS3SourceConfig( + bucket=s3.bucket, + key=s3.key, + key_prefix=s3.prefix, + transform=_transform_for(s3), + expected_bucket_owner=s3.expected_bucket_owner, + fmt=s3.fmt, + csv_format_options=csv_format_options, + ) + + +def _destination_entry_fields( + destination: Any, *, second_include: DistributedMapDestinationInclude +) -> dict[str, Any]: + """Build the members the OnSuccess and OnFailure shapes share.""" + include: list[DistributedMapDestinationInclude] = [] + if destination.include_input: + include.append(DistributedMapDestinationInclude.INPUT) + if second_include is DistributedMapDestinationInclude.OUTPUT: + if destination.include_output: + include.append(second_include) + elif destination.include_error: + include.append(second_include) + return { + "type": DistributedMapDestinationType.S3, + "include": tuple(include), + "s3_destination_config": DistributedMapS3DestinationConfig( + bucket=destination.bucket, + key_prefix=destination.prefix, + expected_bucket_owner=destination.expected_bucket_owner, + ), + } + + +def _build_destination_config( + destination: Any, +) -> DistributedMapDestinationConfig | None: + """Translate the destination config, or None when neither side is set.""" + on_success = ( + DistributedMapOnSuccessConfig( + **_destination_entry_fields( + destination.on_success, + second_include=DistributedMapDestinationInclude.OUTPUT, + ) + ) + if destination.on_success is not None + else None + ) + on_failure = ( + DistributedMapOnFailureConfig( + **_destination_entry_fields( + destination.on_failure, + second_include=DistributedMapDestinationInclude.ERROR, + ) + ) + if destination.on_failure is not None + else None + ) + if on_success is None and on_failure is None: + return None + return DistributedMapDestinationConfig(on_success=on_success, on_failure=on_failure) + + +def _build_source_config( + source: DistributedMapSource | Sequence[Any], + operation_id: str, + durable_execution_arn: str, +) -> DistributedMapSourceConfig: + """Translate a source (typed or plain-list shorthand) into its source config.""" + if not isinstance(source, DistributedMapSource): + # A plain list is treated as an inline source with the default serializer. + serialized_items = _build_inline_items( + tuple(source), DEFAULT_JSON_SERDES, operation_id, durable_execution_arn + ) + return DistributedMapSourceConfig( + source_type=DistributedMapSourceType.INLINE, + inline_source_config=DistributedMapInlineSourceConfig( + items=serialized_items + ), + ) + + if source.inline_items is not None: + serialized_items = _build_inline_items( + source.inline_items, + source.inline_serdes or DEFAULT_JSON_SERDES, + operation_id, + durable_execution_arn, + ) + return DistributedMapSourceConfig( + source_type=DistributedMapSourceType.INLINE, + inline_source_config=DistributedMapInlineSourceConfig( + items=serialized_items + ), + max_items=source.max_items, + ) + if source.s3 is not None: + return DistributedMapSourceConfig( + source_type=DistributedMapSourceType.S3, + max_items=source.max_items, + s3_config=_build_s3_source_config(source.s3), + ) + if source.reader is not None: + reader_config = DistributedMapReaderFunctionSourceConfig( + function_name=source.reader.function_name + ) + if source.reader.initial_state is not None: + state = serialize( + serdes=source.reader.state_serdes or DEFAULT_JSON_SERDES, + value=source.reader.initial_state, + operation_id=operation_id, + durable_execution_arn=durable_execution_arn, + ) + if len(state.encode("utf-8")) > _READER_STATE_LIMIT: + msg = ( + f"reader initial_state exceeds the " + f"{_READER_STATE_LIMIT // 1024} KB limit" + ) + raise ValidationError(msg) + reader_config = DistributedMapReaderFunctionSourceConfig( + function_name=source.reader.function_name, initial_state=state + ) + return DistributedMapSourceConfig( + source_type=DistributedMapSourceType.READER_FUNCTION, + max_items=source.max_items, + reader_config=reader_config, + ) + msg = "Distributed map source has no configured items" + raise ExecutionError(msg) + + +def _build_distributed_map_options( + source: DistributedMapSource | Sequence[Any], + processor: DistributedMapProcessor, + max_concurrency: int, + config: DistributedMapConfig, + operation_id: str, + durable_execution_arn: str, +) -> DistributedMapOptions: + """Assemble the DistributedMapOptions payload from the operands and config.""" + result_collection = ( + DistributedMapResultCollectionConfig( + mode=DistributedMapResultCollectionMode.INLINE + ) + if isinstance(config, DistributedMapResultConfig) + else None + ) + return DistributedMapOptions( + max_concurrency=max_concurrency, + source=_build_source_config(source, operation_id, durable_execution_arn), + processor=_build_processor_config(processor), + destination=( + _build_destination_config(config.destination) + if config.destination is not None + else None + ), + completion_config=( + _build_completion_config(config.completion_config) + if config.completion_config is not None + else None + ), + result_collection=result_collection, + timeout_seconds=config.timeout.to_seconds() + if config.timeout is not None + else None, + ) + + +def _distributed_map_status_from_operation( + operation: Operation | None, +) -> DistributedMapStatus: + """Derive the map run status from the operation's terminal status.""" + op_status = operation.status if operation else None + resolved = ( + { + OperationStatus.SUCCEEDED: DistributedMapStatus.SUCCEEDED, + OperationStatus.FAILED: DistributedMapStatus.FAILED, + OperationStatus.STOPPED: DistributedMapStatus.STOPPED, + OperationStatus.TIMED_OUT: DistributedMapStatus.TIMED_OUT, + }.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 DistributedMapOperationExecutor(OperationExecutor[DistributedMapSummary]): + """Executor for map run operations. + + Creates the START checkpoint if none exists, then suspends until the + backend completes the run and re-invokes the parent. On resume, a + ``DistributedMapSummary`` (or ``DistributedMapResult`` when result collection is enabled) + is built from the checkpointed ``DistributedMapDetails``. + """ + + def __init__( + self, + source: DistributedMapSource | Sequence[Any], + processor: DistributedMapProcessor, + max_concurrency: int, + state: ExecutionState, + operation_identifier: OperationIdentifier, + config: DistributedMapConfig, + ): + """Initialize the map run operation executor. + + Args: + source: The items to process (typed source or plain-list shorthand) + processor: The processor configuration + max_concurrency: Maximum concurrent processor invocations + state: The execution state + operation_identifier: The operation identifier + config: Configuration for the map run operation + """ + self.source = source + self.processor = processor + self.max_concurrency = max_concurrency + self.state = state + self.operation_identifier = operation_identifier + self.config = config + + def _resolve_summary(self, operation: Operation | None) -> DistributedMapSummary: + """Reconstruct the resolved summary/result from the terminal operation.""" + details = operation.distributed_map_details if operation else None + if details is None: + msg = "DISTRIBUTED_MAP operation succeeded but carried no DistributedMapDetails" + raise ExecutionError(msg) + status = _distributed_map_status_from_operation(operation) + completion_reason = details.completion_reason + if completion_reason is None: + msg = ( + f"DISTRIBUTED_MAP operation ended {status.value} but carried no " + f"CompletionReason" + ) + raise ExecutionError(msg) + if not isinstance(self.config, DistributedMapResultConfig): + return DistributedMapSummary( + status=status, + completion_reason=completion_reason, + success_count=details.success_count, + failure_count=details.failure_count, + unprocessed_count=details.unprocessed_count, + distributed_map_run_arn=details.distributed_map_run_arn, + completion_details=details.completion_details, + total_count=details.total_count, + ) + return DistributedMapResult( + status=status, + completion_reason=completion_reason, + success_count=details.success_count, + failure_count=details.failure_count, + unprocessed_count=details.unprocessed_count, + distributed_map_run_arn=details.distributed_map_run_arn, + completion_details=details.completion_details, + total_count=details.total_count, + all=self._deserialize_items(details, self.config.result_serdes), + ) + + def _deserialize_items( + self, details: DistributedMapDetails, result_serdes: Any + ) -> 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, + ) + error = ( + DistributedMapItemError( + error_type=entry.error.type or "", + error_message=entry.error.message or "", + ) + if entry.error is not None + else None + ) + items.append( + DistributedMapResultItem( + item_id=entry.item_id, + status=entry.status, + output=output, + error=error, + ) + ) + return items + + def check_result_status(self) -> CheckResult[DistributedMapSummary]: + """Check operation status and create the START checkpoint if needed. + + Called twice by process() when creating synchronous checkpoints: once before + and once after, to detect if the operation completed immediately. + + Returns: + CheckResult indicating the next action to take + + Raises: + SuspendExecution: For STARTED operations waiting for completion + """ + checkpointed_result: CheckpointedResult = self.state.get_checkpoint_result( + self.operation_identifier.operation_id + ) + + # Terminal success - build the summary/result from 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=_build_distributed_map_options( + source=self.source, + processor=self.processor, + max_concurrency=self.max_concurrency, + config=self.config, + operation_id=self.operation_identifier.operation_id, + durable_execution_arn=self.state.durable_execution_arn, + ), + ) + ) + # Checkpoint map run START with blocking (is_sync=True). + # Must ensure the map run is recorded before suspending execution. + self.state.create_checkpoint(operation_update=start_operation, is_sync=True) + + logger.debug( + "🚀 Map run %s started, will check for immediate completion", + self.operation_identifier.name + or self.operation_identifier.operation_id, + ) + + # Signal to process() that checkpoint was created - to recheck status + # for immediate completion before proceeding. + return CheckResult.create_started() + + # Ready to suspend (checkpoint exists but not in a terminal or started state) + return CheckResult.create_is_ready_to_execute(checkpointed_result) + + def execute( + self, _checkpointed_result: CheckpointedResult + ) -> DistributedMapSummary: + """Execute map run operation by suspending to wait for async completion. + + The map run operation doesn't execute synchronously - it suspends and + the backend runs the map run asynchronously. + + Args: + checkpointed_result: The checkpoint data (unused, but required by interface) + + Returns: + Never returns - always suspends + + Raises: + Always suspends via suspend_with_optional_resume_delay + ExecutionError: If suspend doesn't raise (should never happen) + """ + msg: str = f"Map run {self.operation_identifier.operation_id} started, suspending for completion" + suspend_with_optional_resume_delay(msg) + # This line should never be reached since suspend_with_optional_resume_delay always raises + error_msg: str = "suspend_with_optional_resume_delay should have raised an exception, but did not." + raise ExecutionError(error_msg) from None diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py index 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..9b70a9b01 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, + S3Destination, + S3Source, + DistributedMapProcessor, + DistributedMapSource, + ReaderSourceConfig, +) def test_completion_config_defaults(): @@ -279,3 +288,165 @@ 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.batch("p", max_retry_duration=Duration.from_seconds(30)) + with pytest.raises(ValidationError, match="between 1 minute and 6 hours"): + DistributedMapProcessor.batch("p", max_retry_duration=Duration.from_hours(7)) + + +def test_expected_bucket_owner_must_be_12_digits(): + with pytest.raises(ValidationError, match="12-digit"): + S3Source.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"): + S3Source.csv("s3://b/data.csv", delimiter="BAR") + + +def test_csv_headers_duplicates_rejected(): + with pytest.raises(ValidationError, match="duplicates"): + S3Source.csv("s3://b/k.csv", headers=["a", "a"]) + + +def test_json_lines_requires_key(): + with pytest.raises(ValidationError, match="object key"): + S3Source.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.batch("") + + +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.batch(ref) + + +def test_function_name_over_max_length_rejected(): + with pytest.raises(ValidationError, match="at most 170 characters"): + DistributedMapProcessor.batch("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_sample_size_requires_percentage(): + with pytest.raises(ValidationError, match="minimum_sample_size"): + DistributedMapCompletionConfig(minimum_sample_size=10) + + +def test_completion_negative_count_rejected(): + with pytest.raises(ValidationError, match="non-negative"): + DistributedMapCompletionConfig(tolerated_failure_count=-1) + + +def test_completion_percentage_out_of_range_rejected(): + with pytest.raises(ValidationError, match="between 0 and 100"): + DistributedMapCompletionConfig(tolerated_failure_percentage=150) + + +def test_completion_sample_size_below_one_rejected(): + with pytest.raises(ValidationError, match="at least 1"): + DistributedMapCompletionConfig( + tolerated_failure_percentage=5, minimum_sample_size=0 + ) + + +def test_completion_failure_count_factory(): + assert DistributedMapCompletionConfig.failure_count(3).tolerated_failure_count == 3 + + +def test_retry_duration_below_minimum_rejected(): + with pytest.raises(ValidationError, match="1 minute and 6 hours"): + DistributedMapProcessor.batch("p", max_retry_duration=Duration.from_seconds(30)) + + +def test_max_items_below_one_rejected(): + with pytest.raises(ValidationError, match="at least 1"): + S3Source.json_lines("s3://b/k.jsonl", max_items=0) + + +def test_csv_requires_object_key(): + with pytest.raises(ValidationError, match="csv requires an S3 object key"): + S3Source.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=ReaderSourceConfig(function_name="r"), + ) + + +def test_success_destination_all_false_rejected(): + with pytest.raises(ValidationError, match="success destination must include"): + S3Destination.successes( + "s3://out/ok", include_input=False, include_output=False + ) + + +def test_failure_destination_all_false_rejected(): + with pytest.raises(ValidationError, match="failure destination must include"): + S3Destination.failures("s3://out/bad", include_input=False, include_error=False) + + +def test_invalid_s3_uri_rejected(): + with pytest.raises(ValidationError, match="Invalid S3 URI"): + S3Source.json_lines("s3:///key.jsonl") + + +def test_non_s3_scheme_uri_rejected(): + with pytest.raises(ValidationError, match="must start with s3://"): + S3Source.json_lines("http://foo/bar") + + +def test_csv_empty_headers_rejected(): + with pytest.raises(ValidationError, match="must be non-empty"): + S3Source.csv("s3://b/f.csv", headers=[]) + + +def test_negative_retry_attempts_rejected(): + with pytest.raises(ValidationError, match="non-negative"): + DistributedMapProcessor.batch("p", max_retry_attempts=-5) + + +def test_batch_size_out_of_range_rejected(): + with pytest.raises(ValidationError, match="between 1 and 10000"): + DistributedMapProcessor.batch("p", batch_size=0) + + +def test_durable_execution_name_prefix_too_long_rejected(): + with pytest.raises(ValidationError, match="1 to 36 characters"): + DistributedMapProcessor.batch("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..6c01127f5 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, 10001, 50000]) +def test_map_run_rejects_out_of_range_max_concurrency(max_concurrency: int): + """Test distributed_map rejects a max_concurrency outside the service range.""" + context = create_test_context() + + with pytest.raises( + ValidationError, match="max_concurrency must be between 1 and 10000" + ): + context.distributed_map( + ["a"], + DistributedMapProcessor.batch("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.batch("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.batch("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.batch("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.batch("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.batch("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.batch("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.batch("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..1f4675947 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/dmap/handlers_test.py @@ -0,0 +1,253 @@ +"""Unit tests for the non-durable distributed map authoring wrappers.""" + +from __future__ import annotations + +import threading +import time +from unittest.mock import MagicMock + +import pytest + +from aws_durable_execution_sdk_python.config import ( + DistributedMapConfig, + DistributedMapProcessor, + DistributedMapStatus, +) +from aws_durable_execution_sdk_python.dmap.handlers import ( + 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 ExecutionError, ValidationError +from aws_durable_execution_sdk_python.lambda_service import ( + DistributedMapCompletionReason, + DistributedMapDetails, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.operation.dmap import ( + DistributedMapOperationExecutor, + _distributed_map_status_from_operation, +) + + +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 test_status_derived_from_operation_when_details_omit_status(): + """Status comes from the operation, not from details.""" + details = DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=3, + total_count=3, + ) + op = _terminal_operation(OperationStatus.SUCCEEDED, details) + assert _distributed_map_status_from_operation(op) is DistributedMapStatus.SUCCEEDED + + +def test_status_from_operation_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"): + _distributed_map_status_from_operation(op) + + +def test_resolved_summary_uses_operation_status_not_details(): + """_resolve_summary carries the derived status when details omit it.""" + details = DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=2, + failure_count=0, + unprocessed_count=0, + total_count=2, + ) + executor = DistributedMapOperationExecutor( + source=["a", "b"], + processor=DistributedMapProcessor.batch("proc"), + max_concurrency=1, + state=MagicMock(), + operation_identifier=MagicMock(), + config=DistributedMapConfig(), + ) + summary = executor._resolve_summary( + _terminal_operation(OperationStatus.SUCCEEDED, details) + ) + assert summary.status is DistributedMapStatus.SUCCEEDED + assert summary.completion_reason is DistributedMapCompletionReason.ALL_COMPLETED + assert summary.success_count == 2 + + +def test_item_handler_invalid_report_rejected(): + with pytest.raises(ValidationError, match="report must be"): + distributed_map_item_handler(lambda x: x, report="bogus") + + +def test_durable_item_handler_invalid_report_rejected(): + with pytest.raises(ValidationError, match="report must be"): + durable_distributed_map_item_handler(lambda _ctx, item: item, report="bogus") + + +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_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_failures_form_reports_only_failures(): + handler = distributed_map_item_handler(lambda x: x, report="failures") + resp = handler({"records": [{"itemId": "0", "body": "1"}]}) + assert resp == {"batchItemFailures": []} + + +def test_batch_handler_success_and_propagates_error(): + seen: list = [] + handler = distributed_map_batch_handler(lambda items: seen.extend(items) or "done") + assert ( + handler( + {"records": [{"itemId": "0", "body": "1"}, {"itemId": "1", "body": "2"}]} + ) + == "done" + ) + assert seen == [1, 2] + + def boom(_items): + msg = "batch failed" + raise RuntimeError(msg) + + failing = distributed_map_batch_handler(boom) + with pytest.raises(RuntimeError, match="batch failed"): + failing({"records": [{"itemId": "0", "body": "1"}]}) + + +def test_reader_returns_items_and_next_state_then_exhausts(): + def read(state): + if state is None: + return ReaderPage(items=[1, 2], next_state={"page": 1}) + return ReaderPage(items=[3]) + + handler = distributed_map_reader(read) + first = handler({"state": None, "maxItems": 10}) + assert first["items"] == [1, 2] + assert first["nextState"] == '{"page": 1}' + + second = handler({"state": '{"page": 1}', "maxItems": 10}) + assert second["items"] == [3] + assert "nextState" not in second + + +def test_reader_rejects_page_over_max_items(): + handler = distributed_map_reader(lambda _s: ReaderPage(items=[1, 2, 3])) + with pytest.raises(ValidationError, match="exceeding maxItems"): + handler({"state": None, "maxItems": 2}) + + +def test_reader_rejects_oversized_next_state(): + handler = distributed_map_reader( + lambda _s: ReaderPage(items=[1], next_state="x" * 40_000) + ) + with pytest.raises(ValidationError, match="32 KB limit"): + handler({"state": None, "maxItems": 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: 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..9048ae462 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/dmap/models_test.py @@ -0,0 +1,173 @@ +"""Unit tests for the distributed map result types.""" + +from __future__ import annotations + +import pytest + +from aws_durable_execution_sdk_python.config import ( + DistributedMapCompletionReason, + DistributedMapItemStatus, + DistributedMapStatus, +) +from aws_durable_execution_sdk_python.dmap.models import ( + DistributedMapItemError, + DistributedMapResult, + DistributedMapResultItem, + DistributedMapSummary, +) +from aws_durable_execution_sdk_python.exceptions import DistributedMapError + + +def _result(status, completion_reason, failure_count, items): + return DistributedMapResult( + status=status, + completion_reason=completion_reason, + success_count=len(items) - failure_count, + failure_count=failure_count, + unprocessed_count=0, + all=items, + ) + + +def test_summary_throw_if_error(): + ok = DistributedMapSummary( + status=DistributedMapStatus.SUCCEEDED, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=2, + failure_count=0, + unprocessed_count=0, + ) + ok.throw_if_error() # no raise + + failed = DistributedMapSummary( + status=DistributedMapStatus.FAILED, + completion_reason=DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED, + success_count=0, + failure_count=1, + unprocessed_count=0, + ) + with pytest.raises(DistributedMapError): + failed.throw_if_error() + + succeeded_with_failures = DistributedMapSummary( + status=DistributedMapStatus.SUCCEEDED, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=1, + failure_count=1, + unprocessed_count=0, + ) + with pytest.raises(DistributedMapError): + succeeded_with_failures.throw_if_error() + + +def test_map_run_result_succeeded_failed_filters(): + items = [ + DistributedMapResultItem( + item_id="0", status=DistributedMapItemStatus.SUCCEEDED, output=1 + ), + DistributedMapResultItem( + item_id="1", + status=DistributedMapItemStatus.FAILED, + error=DistributedMapItemError(error_type="E", error_message="boom"), + ), + ] + result = DistributedMapResult( + status=DistributedMapStatus.SUCCEEDED, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=1, + failure_count=1, + unprocessed_count=0, + all=items, + ) + assert [i.item_id for i in result.succeeded()] == ["0"] + assert [i.item_id for i in result.failed()] == ["1"] + assert result.get_results() == [1] + assert result.get_errors()[0].error_message == "boom" + + +def test_summary_has_failure_false_without_failures(): + summary = DistributedMapSummary( + status=DistributedMapStatus.SUCCEEDED, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=0, + failure_count=0, + unprocessed_count=0, + ) + assert summary.has_failure is False + + +def test_result_throw_if_error_raises_first_item_error(): + result = _result( + DistributedMapStatus.SUCCEEDED, + DistributedMapCompletionReason.ALL_COMPLETED, + 1, + [ + DistributedMapResultItem( + item_id="0", status=DistributedMapItemStatus.SUCCEEDED, output=1 + ), + DistributedMapResultItem( + item_id="1", + status=DistributedMapItemStatus.FAILED, + error=DistributedMapItemError(error_type="E", error_message="boom"), + ), + ], + ) + with pytest.raises(DistributedMapError, match="E: boom"): + result.throw_if_error() + + +def test_result_throw_if_error_names_item_without_detail(): + result = _result( + DistributedMapStatus.SUCCEEDED, + DistributedMapCompletionReason.ALL_COMPLETED, + 1, + [ + DistributedMapResultItem( + item_id="7", status=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, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=0, + failure_count=2, + unprocessed_count=0, + all=[], + ) + with pytest.raises(DistributedMapError, match="2 item"): + result.throw_if_error() + + +def test_result_throw_if_error_run_level_failure(): + result = DistributedMapResult( + status=DistributedMapStatus.FAILED, + completion_reason=DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED, + success_count=0, + failure_count=1, + unprocessed_count=0, + all=[], + ) + with pytest.raises(DistributedMapError, match="Map run ended FAILED"): + result.throw_if_error() + + +def test_result_throw_if_error_clean_success_does_not_raise(): + result = DistributedMapResult( + status=DistributedMapStatus.SUCCEEDED, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=1, + failure_count=0, + unprocessed_count=0, + all=[ + DistributedMapResultItem( + item_id="0", status=DistributedMapItemStatus.SUCCEEDED, output=1 + ) + ], + ) + result.throw_if_error() 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..f6336d610 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/dmap_helpers_int_test.py @@ -0,0 +1,146 @@ +"""Integration tests for the durable distributed map authoring wrappers. + +The durable variants are ``@durable_execution`` handlers; they run through the +full invocation harness with a mocked checkpoint backend. +""" + +from __future__ import annotations + +import json +from unittest.mock import Mock, patch + +from aws_durable_execution_sdk_python.dmap.handlers import ( + durable_distributed_map_batch_handler, + durable_distributed_map_item_handler, +) +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, # noqa: F401 (ensures decorator import path is valid) +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + Operation, + OperationStatus, + OperationType, +) + +_ARN = "arn:aws:lambda:us-east-1:123456789012:function:proc:1/durable-execution/execution-1" + + +def _lambda_context(): + ctx = Mock() + ctx.aws_request_id = "test-request-id" + ctx.client_context = None + ctx.identity = None + ctx._epoch_deadline_time_in_ms = 0 # noqa: SLF001 + ctx.invoked_function_arn = "test-arn" + ctx.tenant_id = None + return ctx + + +def _event(records: list[dict]): + return { + "DurableExecutionArn": _ARN, + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [ + { + "Id": "execution-1", + "Type": "EXECUTION", + "Status": "STARTED", + "ExecutionDetails": { + "InputPayload": json.dumps({"records": records}) + }, + } + ], + "NextMarker": "", + }, + "LocalRunner": True, + } + + +def _run(handler, records: list[dict]): + operations = [ + Operation( + operation_id="execution-1", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + ) + ] + + def mock_checkpoint( + durable_execution_arn, checkpoint_token, updates, client_token="token" + ): # noqa: S107 + for update in updates: + operations.append( + Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + ) + ) + return CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState( + operations=operations.copy() + ), + ) + + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client_class.initialize_client.return_value = mock_client + mock_client.checkpoint = mock_checkpoint + return handler(_event(records), _lambda_context()) + + +def test_durable_item_handler_reports_results(): + handler = 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_batch_handler_returns_value(): + handler = durable_distributed_map_batch_handler( + lambda _ctx, items: {"count": len(items)} + ) + result = _run(handler, [{"itemId": "0", "body": "1"}, {"itemId": "1", "body": "2"}]) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + assert json.loads(result["Result"]) == {"count": 2} + + +def test_durable_item_handler_failures_form(): + handler = durable_distributed_map_item_handler( + lambda _ctx, item: item, report="failures" + ) + result = _run(handler, [{"itemId": "0", "body": "1"}]) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + data = json.loads(result["Result"]) + assert data == {"batchItemFailures": []} diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/dmap_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/dmap_int_test.py new file mode 100644 index 000000000..b07917bb7 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/dmap_int_test.py @@ -0,0 +1,246 @@ +"""Integration tests for ctx.distributed_map through a full durable_execution invocation. + +Drives the real suspend/resume flow: the first invocation starts the DISTRIBUTED_MAP +operation and suspends (PENDING); a replay invocation with the operation +completed (carrying DistributedMapDetails) resolves to the summary/result. The backend +is mocked; no emulator or real service is involved. +""" + +from __future__ import annotations + +import json +from typing import Any +from unittest.mock import Mock, patch + +from aws_durable_execution_sdk_python.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, + 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.batch("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( + { + "Status": "SUCCEEDED", + "CompletionReason": "ALL_COMPLETED", + "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.item_results("proc"), + max_concurrency=2, + config=DistributedMapResultConfig(), + ) + assert isinstance(result, DistributedMapResult) + return { + "results": result.get_results(), + "errors": [e.error_message for e in result.get_errors()], + } + + _map_run_id, replay_event = _replay_event( + { + "Status": "SUCCEEDED", + "CompletionReason": "ALL_COMPLETED", + "SuccessCount": 1, + "FailureCount": 1, + "UnprocessedCount": 0, + "TotalCount": 2, + "Results": [ + {"ItemId": "0", "Status": "SUCCEEDED", "Output": "5"}, + { + "ItemId": "1", + "Status": "FAILED", + "Error": {"ErrorType": "E", "ErrorMessage": "boom"}, + }, + ], + } + ) + replay = _run(handler, replay_event) + assert replay["Status"] == InvocationStatus.SUCCEEDED.value + data = json.loads(replay["Result"]) + assert data == {"results": [5], "errors": ["boom"]} + + +def test_map_run_throw_if_error_fails_execution(): + @durable_execution + def handler(event, context: DurableContext) -> dict[str, Any]: + summary = context.distributed_map( + ["a"], + DistributedMapProcessor.batch("proc"), + max_concurrency=1, + ) + summary.throw_if_error() + return {"status": summary.status.value} + + _map_run_id, replay_event = _replay_event( + { + "Status": "FAILED", + "CompletionReason": "FAILURE_TOLERANCE_EXCEEDED", + "SuccessCount": 0, + "FailureCount": 1, + "UnprocessedCount": 0, + "TotalCount": 1, + } + ) + replay = _run(handler, replay_event) + assert replay["Status"] == InvocationStatus.FAILED.value diff --git a/packages/aws-durable-execution-sdk-python/tests/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..aa5a10e03 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 @@ -42,6 +42,10 @@ WaitOptions, _is_in_var_dir, ) +from aws_durable_execution_sdk_python.lambda_service import ( + DistributedMapDetails, + DistributedMapResultItem, +) # ============================================================================= @@ -1271,6 +1275,41 @@ 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 Status/CompletionReason) parses without raising, leaving both None.""" + 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.completion_reason 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": { + "CompletionReason": "ALL_COMPLETED", + "TotalCount": 3, + "SuccessCount": 3, + }, + } + operation = Operation.from_dict(data) + assert operation.distributed_map_details is not None + assert operation.distributed_map_details.completion_reason.value == "ALL_COMPLETED" + + def test_operation_from_dict_individual_options(): """Test Operation.from_dict with each option type individually.""" # Test with just ContextOptions @@ -3029,3 +3068,60 @@ 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", + "CompletionReason": "ALL_COMPLETED", + "SuccessCount": 1, + "FailureCount": 1, + "UnprocessedCount": 0, + "TotalCount": 2, + "DistributedMapRunArn": "arn:aws:lambda:us-east-1:123456789012:map-run:x", + "Results": [ + {"ItemId": "0", "Status": "SUCCEEDED", "Output": 5}, + { + "ItemId": "1", + "Status": "FAILED", + "Error": {"ErrorType": "E", "ErrorMessage": "boom"}, + }, + ], + } + details = DistributedMapDetails.from_dict(data) + assert details.total_count == 2 + assert details.results is not None + assert details.results[0].item_id == "0" + assert details.results[1].error is not None + assert details.results[1].error.type == "E" + + +def test_details_missing_completion_reason_parses_with_none(): + # In-flight runs omit CompletionReason too; parsing must succeed with None. + details = DistributedMapDetails.from_dict({"SuccessCount": 2}) + assert details.completion_reason is None + assert details.success_count == 2 + + +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 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..305a507b1 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dmap_test.py @@ -0,0 +1,1388 @@ +"""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, + DistributedMapCompletionReason, + DistributedMapItemStatus, + DistributedMapStatus, + DistributedMapCompletionConfig, + DistributedMapConfig, + DistributedMapResultConfig, + DistributedMapCsvDelimiter, + DistributedMapDestinationConfig, + InlineSource, + ReaderSource, + S3Destination, + S3Source, + DistributedMapProcessor, +) +from aws_durable_execution_sdk_python.exceptions import ( + ExecutionError, + DistributedMapError, + SuspendExecution, + ValidationError, +) +from aws_durable_execution_sdk_python.identifier import OperationIdentifier +from aws_durable_execution_sdk_python.lambda_service import ( + ErrorObject, + DistributedMapDetails, + 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.batch(processor) + executor = DistributedMapOperationExecutor( + source=source, + processor=processor, + max_concurrency=max_concurrency, + state=state, + operation_identifier=operation_identifier, + config=config, + ) + return executor.process() + + +def _identifier( + operation_id: str, name: str | None = "test_map_run" +) -> OperationIdentifier: + return OperationIdentifier( + operation_id, OperationSubType.DISTRIBUTED_MAP, None, name + ) + + +def test_map_run_handler_already_succeeded(): + """Test distributed_map_handler returns a summary when the operation already succeeded.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + operation = Operation( + operation_id="mr1", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=5, + failure_count=0, + unprocessed_count=0, + total_count=5, + distributed_map_run_arn="arn:aws:lambda:us-east-1:123456789012: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.completion_reason is DistributedMapCompletionReason.ALL_COMPLETED + assert result.success_count == 5 + assert result.failure_count == 0 + assert result.total_count == 5 + assert result.distributed_map_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, + status=OperationStatus.FAILED, + distributed_map_details=DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED, + success_count=3, + failure_count=2, + unprocessed_count=1, + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + result = distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr2"), + ) + + assert result.status is DistributedMapStatus.FAILED + assert ( + result.completion_reason + is DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED + ) + assert result.failure_count == 2 + assert result.has_failure is True + + +def test_map_run_handler_succeeded_no_details_raises(): + """Test a succeeded operation carrying no DistributedMapDetails raises ExecutionError.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + operation = Operation( + operation_id="mr3", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=None, + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + with pytest.raises(ExecutionError): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr3"), + ) + + +def test_map_run_handler_already_started(): + """Test distributed_map_handler suspends when the operation is already started.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + operation = Operation( + operation_id="mr5", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + with pytest.raises( + SuspendExecution, match="Map run mr5 started, suspending for completion" + ): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr5"), + ) + + +def test_map_run_handler_new_operation(): + """Test distributed_map_handler creates a START checkpoint for a new operation.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + started_op = Operation( + operation_id="mr6", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + started = CheckpointedResult.create_from_operation(started_op) + mock_state.get_checkpoint_result.side_effect = [not_found, started] + + with pytest.raises(SuspendExecution): + distributed_map_handler( + source=["a", "b", "c"], + processor="test_processor", + max_concurrency=42, + state=mock_state, + operation_identifier=_identifier("mr6"), + ) + + mock_state.create_checkpoint.assert_called_once() + operation_update = mock_state.create_checkpoint.call_args[1]["operation_update"] + assert operation_update.operation_id == "mr6" + assert operation_update.operation_type == OperationType.DISTRIBUTED_MAP + assert operation_update.action == OperationAction.START + assert operation_update.name == "test_map_run" + + distributed_map_options = operation_update.to_dict()["DistributedMapOptions"] + assert distributed_map_options["MaxConcurrency"] == 42 + assert distributed_map_options["Processor"]["FunctionName"] == "test_processor" + assert distributed_map_options["Source"]["InlineSourceConfig"]["Items"] == [ + '"a"', + '"b"', + '"c"', + ] + + +def test_map_run_handler_no_config(): + """Test distributed_map_handler uses a default config when none is provided.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + started_op = Operation( + operation_id="mr7", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + started = CheckpointedResult.create_from_operation(started_op) + mock_state.get_checkpoint_result.side_effect = [not_found, started] + + with pytest.raises(SuspendExecution): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr7"), + config=None, + ) + + mock_state.create_checkpoint.assert_called_once() + + +# Immediate Response Handling Tests +# ============================================================================ + + +def test_map_run_immediate_response_get_checkpoint_result_called_twice(): + """Test get_checkpoint_result is called twice when a checkpoint is created.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + started_op = Operation( + operation_id="mr8", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + started = CheckpointedResult.create_from_operation(started_op) + mock_state.get_checkpoint_result.side_effect = [not_found, started] + + with pytest.raises(SuspendExecution): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr8"), + ) + + assert mock_state.get_checkpoint_result.call_count == 2 + + +def test_map_run_immediate_response_create_checkpoint_is_sync_true(): + """Test create_checkpoint is called with is_sync=True.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + started_op = Operation( + operation_id="mr9", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + started = CheckpointedResult.create_from_operation(started_op) + mock_state.get_checkpoint_result.side_effect = [not_found, started] + + with pytest.raises(SuspendExecution): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr9"), + ) + + mock_state.create_checkpoint.assert_called_once() + assert mock_state.create_checkpoint.call_args[1]["is_sync"] is True + + +def test_map_run_immediate_response_immediate_success(): + """Test immediate success: second check returns SUCCEEDED, summary returned.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + succeeded_op = Operation( + operation_id="mr10", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=1, + ), + ) + succeeded = CheckpointedResult.create_from_operation(succeeded_op) + mock_state.get_checkpoint_result.side_effect = [not_found, succeeded] + + result = distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr10"), + ) + + assert result.status is DistributedMapStatus.SUCCEEDED + assert result.success_count == 1 + mock_state.create_checkpoint.assert_called_once() + assert mock_state.get_checkpoint_result.call_count == 2 + + +def test_map_run_immediate_response_no_immediate_response(): + """Test no immediate response: second check returns STARTED, suspends.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + not_found = CheckpointedResult.create_not_found() + started_op = Operation( + operation_id="mr12", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.STARTED, + ) + started = CheckpointedResult.create_from_operation(started_op) + mock_state.get_checkpoint_result.side_effect = [not_found, started] + + with pytest.raises(SuspendExecution): + distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr12"), + ) + + mock_state.create_checkpoint.assert_called_once() + assert mock_state.get_checkpoint_result.call_count == 2 + + +def test_map_run_immediate_response_already_completed(): + """Test already completed: first check is SUCCEEDED, no checkpoint created.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + + succeeded_op = Operation( + operation_id="mr13", + operation_type=OperationType.DISTRIBUTED_MAP, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(succeeded_op) + ) + + result = distributed_map_handler( + source=["a"], + processor="test_processor", + max_concurrency=10, + state=mock_state, + operation_identifier=_identifier("mr13"), + ) + + assert result.status is DistributedMapStatus.SUCCEEDED + mock_state.create_checkpoint.assert_not_called() + assert mock_state.get_checkpoint_result.call_count == 1 + + +@patch( + "aws_durable_execution_sdk_python.operation.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, + 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, + status=OperationStatus.STARTED, + ) + ) + return [not_found, started] + + +def test_processor_item_failures_sets_response_types(): + """item_failures serializes FunctionResponseTypes=REPORT_BATCH_ITEM_FAILURES.""" + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.item_failures("proc", batch_size=25), + max_concurrency=4, + config=DistributedMapConfig(), + ) + processor = options["Processor"] + assert processor["FunctionName"] == "proc" + assert processor["FunctionResponseTypes"] == ["REPORT_BATCH_ITEM_FAILURES"] + assert processor["BatchSize"] == 25 + + +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.item_results( + "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_RESULTS"] + 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.batch("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=S3Source.json_lines("s3://bucket/data.jsonl", max_items=500), + processor=DistributedMapProcessor.batch("proc"), + max_concurrency=2, + config=DistributedMapConfig(), + ) + source = options["Source"] + assert source["Type"] == "S3" + assert source["MaxItemsToRead"] == 500 + assert source["S3SourceConfig"]["Bucket"] == "bucket" + assert source["S3SourceConfig"]["Key"] == "data.jsonl" + assert source["S3SourceConfig"]["Format"] == "JSON_LINES" + + +def test_full_config_serializes_all_blocks(): + """Completion, destination, timeout, and result-collection blocks all serialize.""" + config = DistributedMapResultConfig( + completion_config=DistributedMapCompletionConfig.failure_percentage( + 5, minimum_sample_size=200 + ), + destination=DistributedMapDestinationConfig( + on_success=S3Destination.successes("s3://out/ok"), + on_failure=S3Destination.failures("s3://out/bad"), + ), + timeout=Duration.from_minutes(30), + ) + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.item_results("proc"), + max_concurrency=8, + config=config, + ) + assert options["CompletionConfig"] == { + "ToleratedFailurePercentage": 5, + "MinimumSampleSize": 200, + } + on_success = options["Destination"]["OnSuccess"] + assert on_success["Type"] == "S3" + assert on_success["S3DestinationConfig"]["Bucket"] == "out" + assert on_success["S3DestinationConfig"]["KeyPrefix"] == "ok" + on_failure = options["Destination"]["OnFailure"] + assert on_failure["Include"] == ["INPUT", "ERROR"] + assert on_failure["S3DestinationConfig"]["Bucket"] == "out" + assert options["TimeoutSeconds"] == 1800 + assert options["ResultCollection"] == {"Mode": "INLINE"} + + +def test_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, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + 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.item_results("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].error_type == "ItemError" + assert errors[0].error_message == "boom" + + +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, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=2, + ), + ) + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(operation) + ) + + result = distributed_map_handler( + source=["a", "b"], + processor=DistributedMapProcessor.batch("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(): + """CSV headers map to GIVEN. expected_columns stays client-side under FIRST_ROW.""" + # headers -> HeaderLocation GIVEN, headers sent + given = _start_options( + _new_op_state_calls(), + source=S3Source.csv("s3://b/data.csv", headers=["a", "b"]), + processor=DistributedMapProcessor.batch("proc"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + csv_opts = given["Source"]["S3SourceConfig"]["CsvFormatOptions"] + assert csv_opts["HeaderLocation"] == "GIVEN" + assert csv_opts["Headers"] == ["a", "b"] + assert csv_opts["Delimiter"] == "COMMA" + + # no headers -> HeaderLocation FIRST_ROW + first_row = _start_options( + _new_op_state_calls(), + source=S3Source.csv("s3://b/data.csv"), + processor=DistributedMapProcessor.batch("proc"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + s3_cfg = first_row["Source"]["S3SourceConfig"] + assert s3_cfg["CsvFormatOptions"]["HeaderLocation"] == "FIRST_ROW" + assert "Headers" not in s3_cfg["CsvFormatOptions"] + + +# Call-site validation and serdes tests +# ============================================================================ + + +def test_csv_delimiter_accepts_enum(): + opts = _start_options( + _new_op_state_calls(), + source=S3Source.csv( + "s3://b/data.csv", delimiter=DistributedMapCsvDelimiter.PIPE + ), + processor=DistributedMapProcessor.batch("proc"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + assert opts["Source"]["S3SourceConfig"]["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.batch("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=ReaderSource.from_function("reader", initial_state={"page": 0}), + processor=DistributedMapProcessor.batch("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=ReaderSource.from_function("reader", initial_state="x" * 40_000), + processor=DistributedMapProcessor.batch("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=InlineSource.of(["a", "b"], serdes=_UpperSerDes()), + processor=DistributedMapProcessor.batch("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.batch("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.batch("p"), + max_concurrency=1, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_success=S3Destination.successes("s3://out/ok") + ) + ), + ) + dest = opts["Destination"] + assert "OnSuccess" in dest + assert "OnFailure" not in dest + + +def test_destination_only_failure(): + opts = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.batch("p"), + max_concurrency=1, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_failure=S3Destination.failures("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=S3Source.objects("s3://b/prefix/"), + processor=DistributedMapProcessor.batch("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + s3 = opts["Source"]["S3SourceConfig"] + assert s3["Transform"] == "NONE" + assert s3["KeyPrefix"] == "prefix/" + assert "Format" not in s3 + + +def test_s3_flattened_json_lines_source_config(): + opts = _start_options( + _new_op_state_calls(), + source=S3Source.flattened_json_lines("s3://b/prefix/"), + processor=DistributedMapProcessor.batch("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + s3 = opts["Source"]["S3SourceConfig"] + assert s3["Transform"] == "LOAD_AND_FLATTEN" + assert s3["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_empty_completion_config_omitted(): + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.batch("p"), + max_concurrency=1, + config=DistributedMapConfig(completion_config=DistributedMapCompletionConfig()), + ) + assert "CompletionConfig" not in options + + +# --- Source translation variants --- + + +def test_objects_whole_bucket_prefix_config(): + options = _start_options( + _new_op_state_calls(), + source=S3Source.objects("s3://bucket"), + processor=DistributedMapProcessor.batch("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + s3 = options["Source"]["S3SourceConfig"] + assert s3["KeyPrefix"] == "" + assert s3["Transform"] == "NONE" + assert "Key" not in s3 + + +def test_flattened_csv_source_config(): + options = _start_options( + _new_op_state_calls(), + source=S3Source.flattened_csv("s3://b/prefix", headers=["a", "b"]), + processor=DistributedMapProcessor.batch("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + s3 = options["Source"]["S3SourceConfig"] + assert s3["Transform"] == "LOAD_AND_FLATTEN" + assert s3["Format"] == "CSV" + assert s3["CsvFormatOptions"]["HeaderLocation"] == "GIVEN" + assert s3["CsvFormatOptions"]["Headers"] == ["a", "b"] + + +def test_reader_source_without_initial_state_omits_state(): + options = _start_options( + _new_op_state_calls(), + source=ReaderSource.from_function("reader"), + processor=DistributedMapProcessor.batch("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=InlineSource.of([1, 2], serdes=_PlainTextSerDes()), + processor=DistributedMapProcessor.batch("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.batch("p"), + max_concurrency=1, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_success=S3Destination.successes( + "s3://out/ok", + include_input=True, + include_output=True, + expected_bucket_owner="123456789012", + ) + ) + ), + ) + on_success = options["Destination"]["OnSuccess"] + assert on_success["Include"] == ["INPUT", "OUTPUT"] + assert on_success["S3DestinationConfig"]["ExpectedBucketOwner"] == "123456789012" + assert "OnFailure" not in options["Destination"] + + +def test_failure_destination_error_only_and_owner(): + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.batch("p"), + max_concurrency=1, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_failure=S3Destination.failures( + "s3://out/bad", + include_input=False, + include_error=True, + expected_bucket_owner="123456789012", + ) + ) + ), + ) + on_failure = options["Destination"]["OnFailure"] + assert on_failure["Include"] == ["ERROR"] + assert on_failure["S3DestinationConfig"]["ExpectedBucketOwner"] == "123456789012" + + +# --- 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, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=1, + failure_count=1, + unprocessed_count=0, + total_count=2, + distributed_map_run_arn="arn:aws:lambda:us-east-1:123456789012:function:fn:$LATEST/durable-execution/exec1/invoke1/distributed-map-run/z", + completion_details="done", + 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["CompletionDetails"] == "done" + assert block["TotalCount"] == 2 + + parsed = Operation.from_dict(op.to_dict()) + assert parsed.distributed_map_details is not None + assert parsed.distributed_map_details.results[0].output == 5 + assert parsed.distributed_map_details.results[1].error.type == "E" + + +def test_options_full_round_trip(): + options_dict = _start_options( + _new_op_state_calls(), + source=S3Source.csv("s3://b/f.csv", headers=["a"]), + processor=DistributedMapProcessor.item_results( + "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=DistributedMapDestinationConfig( + on_success=S3Destination.successes("s3://o/ok"), + on_failure=S3Destination.failures("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_RESULTS, + ) + 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=S3Source.csv("s3://b/f.csv"), + processor=DistributedMapProcessor.batch("p"), + max_concurrency=1, + config=DistributedMapConfig(), + ) + csv_opts = options["Source"]["S3SourceConfig"]["CsvFormatOptions"] + assert csv_opts["HeaderLocation"] == "FIRST_ROW" + assert "Headers" not in csv_opts + assert csv_opts["Delimiter"] == "COMMA" + + +def test_s3_source_expected_bucket_owner(): + options = _start_options( + _new_op_state_calls(), + source=S3Source.json_lines( + "s3://b/k.jsonl", expected_bucket_owner="123456789012" + ), + processor=DistributedMapProcessor.batch("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.batch("p"), + max_concurrency=1, + config=DistributedMapConfig(destination=DistributedMapDestinationConfig()), + ) + assert "Destination" not in options + + +def test_success_destination_input_only(): + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.batch("p"), + max_concurrency=1, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_success=S3Destination.successes( + "s3://out/ok", include_input=True, include_output=False + ) + ) + ), + ) + assert options["Destination"]["OnSuccess"]["Include"] == ["INPUT"] + + +def test_failure_destination_input_only(): + options = _start_options( + _new_op_state_calls(), + source=["a"], + processor=DistributedMapProcessor.batch("p"), + max_concurrency=1, + config=DistributedMapConfig( + destination=DistributedMapDestinationConfig( + on_failure=S3Destination.failures( + "s3://out/bad", include_input=True, include_error=False + ) + ) + ), + ) + assert options["Destination"]["OnFailure"]["Include"] == ["INPUT"] + + +def test_reader_source_options_round_trip(): + options_dict = _start_options( + _new_op_state_calls(), + source=ReaderSource.from_function("reader", initial_state={"page": 0}), + processor=DistributedMapProcessor.batch("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, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + 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, + 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, + status=status, + distributed_map_details=DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED, + 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_completion_reason_raises(): + """A terminal operation missing its completion reason 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, + status=OperationStatus.STOPPED, + distributed_map_details=DistributedMapDetails( + 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 CompletionReason"): + 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, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + 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.item_results("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, + status=OperationStatus.SUCCEEDED, + distributed_map_details=DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + 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.item_results("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.batch( + "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 From b2f76ca1ac418ac80af850ca6466729c32120f64 Mon Sep 17 00:00:00 2001 From: nvasiu Date: Fri, 2 Oct 2026 22:03:13 +0000 Subject: [PATCH 2/3] refactor(dmap): Align with SDK conventions - The wire types translate their config-layer counterparts through from_* and create_* classmethods, so operation/dmap.py holds the executor alone and its serialization reads the operation id and execution ARN off self. - dmap/models.py rebuilds the summary and the result from a terminal operation, and parses the run ARN to expose the run id. - The source and destination factories are classmethods on the type they return, and the resolved types that are not a config= argument drop the Config suffix. - Collection fields are list rather than tuple, matching the rest of the SDK. - The item handler decorators take a typed response_mode in place of a report string. - The concurrency bound and the wire sentinels sit beside the types they apply to, and each shared type is imported from the module that defines it. - The executor fetches its checkpoint through the base class helper, which raises NonDeterministicExecutionError on an identity mismatch. - Each test file covers the module it is named for. --- .../__init__.py | 2 + .../config.py | 195 ++++---- .../context.py | 3 +- .../dmap/handlers.py | 119 +++-- .../dmap/models.py | 128 ++++- .../lambda_service.py | 219 ++++++++- .../operation/dmap.py | 453 ++++-------------- .../tests/config_test.py | 34 +- .../tests/dmap/handlers_test.py | 132 ++--- .../tests/dmap/models_test.py | 106 +++- .../tests/e2e/dmap_helpers_int_test.py | 51 +- .../tests/e2e/dmap_int_test.py | 5 +- .../tests/operation/dmap_test.py | 208 ++++++-- 13 files changed, 979 insertions(+), 676 deletions(-) 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 bd0199539..dc38cc342 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 @@ -20,6 +20,7 @@ DistributedMapItemStatus, DistributedMapStatus, ParallelBranch, + ProcessorResponseMode, complete_batch, continue_batch, ) @@ -101,6 +102,7 @@ "InvocationError", "InvokeError", "ParallelBranch", + "ProcessorResponseMode", "PluginLoadError", "RetryableSerDesError", "SerDesError", 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 111af6309..f6c951162 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 @@ -576,8 +576,8 @@ def _validate_bucket_owner(value: str | None) -> None: raise ValidationError(msg) -def _validate_columns(name: str, columns: tuple[str, ...] | None) -> None: - """Validate a CSV columns/headers tuple is non-empty with no duplicates.""" +def _validate_columns(name: str, columns: list[str] | None) -> None: + """Validate a CSV columns/headers list is non-empty with no duplicates.""" if columns is None: return if len(columns) == 0: @@ -786,15 +786,15 @@ def _parse_delimiter( @dataclass(frozen=True) -class S3SourceConfig: - """Resolved S3 source configuration.""" +class S3Source: + """Resolved S3 source.""" bucket: str key: str | None = None prefix: str | None = None fmt: DistributedMapSourceFormat | None = None delimiter: DistributedMapCsvDelimiter | None = None - headers: tuple[str, ...] | None = None + headers: list[str] | None = None expected_bucket_owner: str | None = None def __post_init__(self) -> None: @@ -803,8 +803,8 @@ def __post_init__(self) -> None: @dataclass(frozen=True) -class ReaderSourceConfig: - """Resolved reader-function source configuration.""" +class ReaderSource: + """Resolved reader-function source.""" function_name: str initial_state: Any = None @@ -819,10 +819,10 @@ class DistributedMapSource: """Source configuration for a map run.""" max_items: int | None = None - inline_items: tuple[Any, ...] | None = None + inline_items: list[Any] | None = None inline_serdes: SerDes | None = None # None = DEFAULT_JSON_SERDES - s3: S3SourceConfig | None = None - reader: ReaderSourceConfig | None = None + s3: S3Source | None = None + reader: ReaderSource | None = None def __post_init__(self) -> None: if self.max_items is not None and self.max_items < 1: @@ -837,30 +837,24 @@ def __post_init__(self) -> None: msg = "exactly one of inline_items, s3 or reader must be set" raise ValidationError(msg) - -class InlineSource: - """Inline source factory.""" - - @staticmethod - def of( + @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 DistributedMapSource( - inline_items=tuple(items), + return cls( + inline_items=list(items), inline_serdes=serdes, max_items=max_items, ) - -class S3Source: - """S3 source factories.""" - - @staticmethod - def json_lines( + @classmethod + def s3_json_lines( + cls, uri: str, *, expected_bucket_owner: str | None = None, @@ -870,11 +864,11 @@ def json_lines( parsed_uri = S3Uri.parse(uri) bucket, key = parsed_uri.bucket, parsed_uri.path if key is None: - msg = "json_lines requires an S3 object key" + msg = "s3_json_lines requires an S3 object key" raise ValidationError(msg) - return DistributedMapSource( + return cls( max_items=max_items, - s3=S3SourceConfig( + s3=S3Source( bucket=bucket, key=key, fmt=DistributedMapSourceFormat.JSON_LINES, @@ -882,8 +876,9 @@ def json_lines( ), ) - @staticmethod - def json_array( + @classmethod + def s3_json_array( + cls, uri: str, *, expected_bucket_owner: str | None = None, @@ -893,11 +888,11 @@ def json_array( parsed_uri = S3Uri.parse(uri) bucket, key = parsed_uri.bucket, parsed_uri.path if key is None: - msg = "json_array requires an S3 object key" + msg = "s3_json_array requires an S3 object key" raise ValidationError(msg) - return DistributedMapSource( + return cls( max_items=max_items, - s3=S3SourceConfig( + s3=S3Source( bucket=bucket, key=key, fmt=DistributedMapSourceFormat.JSON_ARRAY, @@ -905,8 +900,9 @@ def json_array( ), ) - @staticmethod - def csv( + @classmethod + def s3_csv( + cls, uri: str, *, headers: Sequence[str] | None = None, @@ -918,22 +914,23 @@ def csv( parsed_uri = S3Uri.parse(uri) bucket, key = parsed_uri.bucket, parsed_uri.path if key is None: - msg = "csv requires an S3 object key" + msg = "s3_csv requires an S3 object key" raise ValidationError(msg) - return DistributedMapSource( + return cls( max_items=max_items, - s3=S3SourceConfig( + s3=S3Source( bucket=bucket, key=key, fmt=DistributedMapSourceFormat.CSV, delimiter=_parse_delimiter(delimiter), - headers=tuple(headers) if headers is not None else None, + headers=list(headers) if headers is not None else None, expected_bucket_owner=expected_bucket_owner, ), ) - @staticmethod - def objects( + @classmethod + def s3_objects( + cls, prefix_uri: str, *, expected_bucket_owner: str | None = None, @@ -942,17 +939,18 @@ def objects( """Read each object under a prefix as one item.""" parsed_uri = S3Uri.parse(prefix_uri) bucket, prefix = parsed_uri.bucket, parsed_uri.path - return DistributedMapSource( + return cls( max_items=max_items, - s3=S3SourceConfig( + s3=S3Source( bucket=bucket, prefix=prefix or "", expected_bucket_owner=expected_bucket_owner, ), ) - @staticmethod - def flattened_json_lines( + @classmethod + def s3_flattened_json_lines( + cls, prefix_uri: str, *, expected_bucket_owner: str | None = None, @@ -961,9 +959,9 @@ def flattened_json_lines( """Read a prefix, flattening each object's lines into items.""" parsed_uri = S3Uri.parse(prefix_uri) bucket, prefix = parsed_uri.bucket, parsed_uri.path - return DistributedMapSource( + return cls( max_items=max_items, - s3=S3SourceConfig( + s3=S3Source( bucket=bucket, prefix=prefix or "", fmt=DistributedMapSourceFormat.JSON_LINES, @@ -971,8 +969,9 @@ def flattened_json_lines( ), ) - @staticmethod - def flattened_json_array( + @classmethod + def s3_flattened_json_array( + cls, prefix_uri: str, *, expected_bucket_owner: str | None = None, @@ -981,9 +980,9 @@ def flattened_json_array( """Read a prefix, flattening each object's JSON array elements into items.""" parsed_uri = S3Uri.parse(prefix_uri) bucket, prefix = parsed_uri.bucket, parsed_uri.path - return DistributedMapSource( + return cls( max_items=max_items, - s3=S3SourceConfig( + s3=S3Source( bucket=bucket, prefix=prefix or "", fmt=DistributedMapSourceFormat.JSON_ARRAY, @@ -991,8 +990,9 @@ def flattened_json_array( ), ) - @staticmethod - def flattened_csv( + @classmethod + def s3_flattened_csv( + cls, prefix_uri: str, *, headers: Sequence[str] | None = None, @@ -1003,24 +1003,21 @@ def flattened_csv( """Read a prefix, flattening each object's records into items.""" parsed_uri = S3Uri.parse(prefix_uri) bucket, prefix = parsed_uri.bucket, parsed_uri.path - return DistributedMapSource( + return cls( max_items=max_items, - s3=S3SourceConfig( + s3=S3Source( bucket=bucket, prefix=prefix or "", fmt=DistributedMapSourceFormat.CSV, delimiter=_parse_delimiter(delimiter), - headers=tuple(headers) if headers is not None else None, + headers=list(headers) if headers is not None else None, expected_bucket_owner=expected_bucket_owner, ), ) - -class ReaderSource: - """Reader-function source factories.""" - - @staticmethod - def from_function( + @classmethod + def reader_function( + cls, name: str, *, initial_state: Any = None, @@ -1028,9 +1025,9 @@ def from_function( max_items: int | None = None, ) -> DistributedMapSource: """Page items from a customer-supplied reader Lambda function.""" - return DistributedMapSource( + return cls( max_items=max_items, - reader=ReaderSourceConfig( + reader=ReaderSource( function_name=name, initial_state=initial_state, state_serdes=state_serdes, @@ -1054,6 +1051,25 @@ def __post_init__(self) -> None: 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: @@ -1071,39 +1087,9 @@ def __post_init__(self) -> None: msg = "failure destination must include input or error" raise ValidationError(msg) - -@dataclass(frozen=True) -class DistributedMapDestinationConfig: - """Destination routing for map run results.""" - - on_success: SuccessDestination | None = None - on_failure: FailureDestination | None = None - - -class S3Destination: - """S3 destination factories.""" - - @staticmethod - def successes( - prefix_uri: str, - *, - include_input: bool = False, - include_output: bool = True, - expected_bucket_owner: str | None = None, - ) -> SuccessDestination: - """Route succeeded item records to an S3 prefix.""" - parsed_uri = S3Uri.parse(prefix_uri) - bucket, prefix = parsed_uri.bucket, parsed_uri.path - return SuccessDestination( - bucket=bucket, - prefix=prefix or "", - include_input=include_input, - include_output=include_output, - expected_bucket_owner=expected_bucket_owner, - ) - - @staticmethod - def failures( + @classmethod + def from_uri( + cls, prefix_uri: str, *, include_input: bool = True, @@ -1112,21 +1098,28 @@ def failures( ) -> FailureDestination: """Route permanently-failed item records to an S3 prefix.""" parsed_uri = S3Uri.parse(prefix_uri) - bucket, prefix = parsed_uri.bucket, parsed_uri.path - return FailureDestination( - bucket=bucket, - prefix=prefix or "", + 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: DistributedMapDestinationConfig | None = None + destination: DistributedMapDestination | None = None completion_config: DistributedMapCompletionConfig | None = None timeout: Duration | None = None 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 e7de3162b..e21e85ba0 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 @@ -52,6 +52,7 @@ OperationIdNamespace, ) from aws_durable_execution_sdk_python.lambda_service import ( + _MAX_CONCURRENCY_LIMIT, OperationSubType, OperationType, ) @@ -114,8 +115,6 @@ PASS_THROUGH_SERDES: SerDes[Any] = PassThroughSerDes() -_MAX_CONCURRENCY_LIMIT = 10000 - @dataclass(frozen=True) class ExecutionContext: 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 index 952744f71..924085918 100644 --- 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 @@ -11,10 +11,14 @@ from collections.abc import MutableMapping from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field -from typing import Any, Callable, Literal +from typing import Any, Callable -from aws_durable_execution_sdk_python.concurrency.models import BatchItemStatus -from aws_durable_execution_sdk_python.config import CompletionConfig, MapConfig +from aws_durable_execution_sdk_python.config import ( + BatchItemStatus, + CompletionConfig, + MapConfig, + ProcessorResponseMode, +) from aws_durable_execution_sdk_python.exceptions import ValidationError from aws_durable_execution_sdk_python.execution import durable_execution from aws_durable_execution_sdk_python.serdes import ( @@ -53,7 +57,7 @@ def to_dict(self) -> MutableMapping[str, Any]: class ProcessorEvent: """The event a processor function is invoked with, one batch of items.""" - records: tuple[ProcessorRecord, ...] = () + records: list[ProcessorRecord] = field(default_factory=list) @classmethod def from_dict(cls, data: MutableMapping[str, Any]) -> ProcessorEvent: @@ -61,7 +65,7 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> ProcessorEvent: if not isinstance(records, list): msg = "expected a distributed map processor envelope with a 'records' list" raise ValidationError(msg) - return cls(records=tuple(ProcessorRecord.from_dict(r) for r in records)) + 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]} @@ -123,17 +127,17 @@ def to_dict(self) -> MutableMapping[str, Any]: class ItemHandlerResponse: """What an item handler returns. A None ``results`` reports failures only.""" - failures: tuple[ItemFailure, ...] = () - results: tuple[ItemResult, ...] | None = None + failures: list[ItemFailure] = field(default_factory=list) + results: list[ItemResult] | None = None @classmethod def from_dict(cls, data: MutableMapping[str, Any]) -> ItemHandlerResponse: results = data.get("batchItemResults") return cls( - failures=tuple( + failures=[ ItemFailure.from_dict(f) for f in data.get("batchItemFailures", []) - ), - results=tuple(ItemResult.from_dict(r) for r in results) + ], + results=[ItemResult.from_dict(r) for r in results] if results is not None else None, ) @@ -151,21 +155,22 @@ def to_dict(self) -> MutableMapping[str, Any]: class ReaderEvent: """The event a reader function is invoked with.""" - max_items: int + max_items_per_page: int state: str | None = None @classmethod def from_dict(cls, data: MutableMapping[str, Any]) -> ReaderEvent: - max_items = data.get("maxItems") - if not isinstance(max_items, int): + 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 'maxItems'" + "expected a distributed map reader envelope with an integer " + "'maxItemsPerPage'" ) raise ValidationError(msg) - return cls(max_items=max_items, state=data.get("state")) + 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] = {"maxItems": self.max_items} + result: MutableMapping[str, Any] = {"maxItemsPerPage": self.max_items_per_page} if self.state is not None: result["state"] = self.state return result @@ -175,15 +180,15 @@ def to_dict(self) -> MutableMapping[str, Any]: class ReaderResponse: """What a reader function returns. A None ``next_state`` exhausts the source.""" - items: tuple[Any, ...] = () + items: list[Any] = field(default_factory=list) next_state: str | None = None @classmethod def from_dict(cls, data: MutableMapping[str, Any]) -> ReaderResponse: - return cls(items=tuple(data.get("items", ())), next_state=data.get("nextState")) + return cls(items=list(data.get("items", [])), next_state=data.get("nextState")) def to_dict(self) -> MutableMapping[str, Any]: - result: MutableMapping[str, Any] = {"items": list(self.items)} + result: MutableMapping[str, Any] = {"items": self.items} if self.next_state is not None: result["nextState"] = self.next_state return result @@ -199,9 +204,12 @@ def _to_output(serdes: SerDes[Any], value: Any, ctx: SerDesContext) -> str: return serdes.serialize(value, ctx) -def _validate_report(report: str) -> None: - if report not in ("results", "failures"): - msg = f"report must be 'results' or 'failures', got: {report!r}" +def _validate_response_mode(response_mode: ProcessorResponseMode) -> None: + if response_mode is ProcessorResponseMode.BATCH: + msg = ( + "response_mode must be ProcessorResponseMode.ITEM_RESULTS or " + "ProcessorResponseMode.ITEM_FAILURES, got: BATCH" + ) raise ValidationError(msg) @@ -217,13 +225,14 @@ def distributed_map_item_handler( item_serdes: SerDes[Any] | None = None, result_serdes: SerDes[Any] | None = None, concurrency: int = 1, - report: Literal["results", "failures"] = "results", + response_mode: ProcessorResponseMode = ProcessorResponseMode.ITEM_RESULTS, ) -> Callable[..., Any]: """Wrap a Lambda to be used as an item_results processor. - Pass ``report="failures"`` for an item_failures 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``. + Pass ``response_mode=ProcessorResponseMode.ITEM_FAILURES`` for an + item_failures 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``. 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 @@ -235,9 +244,9 @@ def distributed_map_item_handler( item_serdes=item_serdes, result_serdes=result_serdes, concurrency=concurrency, - report=report, + response_mode=response_mode, ) - _validate_report(report) + _validate_response_mode(response_mode) _validate_concurrency(concurrency) serdes = item_serdes or DEFAULT_JSON_SERDES out_serdes = result_serdes or DEFAULT_JSON_SERDES @@ -265,7 +274,7 @@ def run(record: ProcessorRecord) -> Any: err = errors[i] if err is not None: failures.append(ItemFailure.from_exception(record.item_id, err)) - elif report == "results": + elif response_mode is ProcessorResponseMode.ITEM_RESULTS: results.append( ItemResult( item_identifier=record.item_id, @@ -274,8 +283,10 @@ def run(record: ProcessorRecord) -> Any: ) response = ItemHandlerResponse( - failures=tuple(failures), - results=tuple(results) if report == "results" else None, + failures=failures, + results=results + if response_mode is ProcessorResponseMode.ITEM_RESULTS + else None, ) return dict(response.to_dict()) @@ -310,15 +321,16 @@ def handler(event: dict[str, Any], _context: Any = None) -> Any: def distributed_map_reader( - func: Callable[[Any], ReaderPage] | None = None, + func: Callable[[Any, int], ReaderPage] | 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 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``. + ``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``. """ if func is None: return functools.partial(distributed_map_reader, state_serdes=state_serdes) @@ -333,11 +345,11 @@ def handler(event: dict[str, Any], _context: Any = None) -> dict[str, Any]: else None ) - page = func(state) - if len(page.items) > reader_event.max_items: + 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 maxItems " - f"{reader_event.max_items}" + f"reader returned {len(page.items)} items, exceeding " + f"maxItemsPerPage {reader_event.max_items_per_page}" ) raise ValidationError(msg) @@ -348,9 +360,7 @@ def handler(event: dict[str, Any], _context: Any = None) -> dict[str, Any]: msg = f"reader next_state exceeds the {_READER_STATE_LIMIT // 1024} KB limit" raise ValidationError(msg) - return dict( - ReaderResponse(items=tuple(page.items), next_state=next_state).to_dict() - ) + return dict(ReaderResponse(items=page.items, next_state=next_state).to_dict()) return handler @@ -360,21 +370,27 @@ def durable_distributed_map_item_handler( *, item_serdes: SerDes[Any] | None = None, result_serdes: SerDes[Any] | None = None, - report: Literal["results", "failures"] = "results", + concurrency: int = 1, + response_mode: ProcessorResponseMode = ProcessorResponseMode.ITEM_RESULTS, ) -> 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``. + + 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, - report=report, + concurrency=concurrency, + response_mode=response_mode, ) - _validate_report(report) + _validate_response_mode(response_mode) + _validate_concurrency(concurrency) serdes = item_serdes or DEFAULT_JSON_SERDES out_serdes = result_serdes or DEFAULT_JSON_SERDES @@ -389,7 +405,10 @@ def per_item(inner_ctx: Any, body: Any, _index: int, _inputs: Any) -> Any: batch = context.map( [record.body for record in records], per_item, - config=MapConfig(completion_config=CompletionConfig.all_completed()), + config=MapConfig( + max_concurrency=concurrency, + completion_config=CompletionConfig.all_completed(), + ), ) results: list[ItemResult] = [] @@ -397,7 +416,7 @@ def per_item(inner_ctx: Any, body: Any, _index: int, _inputs: Any) -> Any: for bi in batch.all: item_id = records[bi.index].item_id if bi.status is BatchItemStatus.SUCCEEDED: - if report == "results": + if response_mode is ProcessorResponseMode.ITEM_RESULTS: results.append( ItemResult( item_identifier=item_id, @@ -415,8 +434,10 @@ def per_item(inner_ctx: Any, body: Any, _index: int, _inputs: Any) -> Any: ) response = ItemHandlerResponse( - failures=tuple(failures), - results=tuple(results) if report == "results" else None, + failures=failures, + results=results + if response_mode is ProcessorResponseMode.ITEM_RESULTS + else None, ) return dict(response.to_dict()) 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 index 7d9ad80c6..feeb4e95f 100644 --- 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 @@ -3,16 +3,90 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any +from typing import TYPE_CHECKING, Any, NoReturn from aws_durable_execution_sdk_python.exceptions import ( DistributedMapError, + ExecutionError, + ValidationError, ) from aws_durable_execution_sdk_python.config import ( DistributedMapCompletionReason, DistributedMapItemStatus, DistributedMapStatus, ) +from aws_durable_execution_sdk_python.lambda_service import OperationStatus +from aws_durable_execution_sdk_python.types import DurableExecutionArn + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python.lambda_service import ( + DistributedMapDetails, + 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 + + +def _resolve_terminal( + operation: Operation | None, +) -> tuple[DistributedMapDetails, DistributedMapStatus, DistributedMapCompletionReason]: + """Read the details, status and completion reason a terminal operation must carry.""" + details = operation.distributed_map_details if operation else None + if details is None: + msg = "DISTRIBUTED_MAP operation succeeded but carried no DistributedMapDetails" + raise ExecutionError(msg) + status = _to_distributed_map_status(operation) + if details.completion_reason is None: + msg = ( + f"DISTRIBUTED_MAP operation ended {status.value} but carried no " + f"CompletionReason" + ) + raise ExecutionError(msg) + return details, status, details.completion_reason + + +_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) @@ -32,6 +106,29 @@ class DistributedMapSummary: completion_details: str | None = None total_count: int | None = None + @classmethod + def from_operation(cls, operation: Operation | None) -> DistributedMapSummary: + """Rebuild the summary from a terminal DISTRIBUTED_MAP operation.""" + details, status, completion_reason = _resolve_terminal(operation) + return cls( + status=status, + completion_reason=completion_reason, + success_count=details.success_count, + failure_count=details.failure_count, + unprocessed_count=details.unprocessed_count, + distributed_map_run_arn=details.distributed_map_run_arn, + completion_details=details.completion_details, + total_count=details.total_count, + ) + + @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.""" @@ -82,6 +179,35 @@ class DistributedMapResult(DistributedMapSummary): all: list[DistributedMapResultItem] = field(default_factory=list) + @classmethod + def from_operation(cls, operation: Operation | None) -> NoReturn: + """Always raises. Use :meth:`from_operation_and_items` instead.""" + msg = ( + "DistributedMapResult carries per-item results, which from_operation " + "cannot supply. Call DistributedMapResult.from_operation_and_items " + "with the deserialized items, or DistributedMapSummary.from_operation " + "for the counts alone." + ) + raise ValidationError(msg) + + @classmethod + def from_operation_and_items( + cls, operation: Operation | None, items: list[DistributedMapResultItem] + ) -> DistributedMapResult: + """Rebuild the result from a terminal operation and its deserialized items.""" + details, status, completion_reason = _resolve_terminal(operation) + return cls( + status=status, + completion_reason=completion_reason, + success_count=details.success_count, + failure_count=details.failure_count, + unprocessed_count=details.unprocessed_count, + distributed_map_run_arn=details.distributed_map_run_arn, + completion_details=details.completion_details, + total_count=details.total_count, + all=items, + ) + def succeeded(self) -> list[DistributedMapResultItem]: """Return the items that succeeded.""" return [ 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 a51568d16..2bd80c38a 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 @@ -22,11 +22,20 @@ 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 ( DistributedMapCompletionReason, DistributedMapCsvDelimiter, + DistributedMapDestination, DistributedMapItemStatus, + DistributedMapProcessor, DistributedMapSourceFormat, + FailureDestination, + ProcessorResponseMode, + S3Source, + SuccessDestination, ) from aws_durable_execution_sdk_python.exceptions import ( CheckpointError, @@ -501,7 +510,7 @@ class DistributedMapDetails: success_count: int = 0 failure_count: int = 0 unprocessed_count: int = 0 - results: tuple[DistributedMapResultItem, ...] | None = None + results: list[DistributedMapResultItem] | None = None @classmethod def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapDetails: @@ -526,9 +535,7 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapDetails: success_count=data.get("SuccessCount", 0), failure_count=data.get("FailureCount", 0), unprocessed_count=data.get("UnprocessedCount", 0), - results=tuple( - DistributedMapResultItem.from_dict(item) for item in results_raw - ) + results=[DistributedMapResultItem.from_dict(item) for item in results_raw] if results_raw is not None else None, ) @@ -634,7 +641,7 @@ class DistributedMapCsvFormatOptions: """Represent CSV format options for an S3 source.""" header_location: DistributedMapCsvHeaderLocation - headers: tuple[str, ...] | None = None + headers: list[str] | None = None delimiter: DistributedMapCsvDelimiter | None = None def to_dict(self) -> MutableMapping[str, Any]: @@ -642,7 +649,7 @@ def to_dict(self) -> MutableMapping[str, Any]: "HeaderLocation": self.header_location.value } if self.headers is not None: - result["Headers"] = list(self.headers) + result["Headers"] = self.headers if self.delimiter is not None: result["Delimiter"] = self.delimiter.value return result @@ -657,12 +664,27 @@ def from_dict( header_location=DistributedMapCsvHeaderLocation( data.get("HeaderLocation", "FIRST_ROW") ), - headers=tuple(headers) if headers is not None else None, + 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 DistributedMapS3SourceConfig: @@ -711,6 +733,28 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapS3SourceConf else None, ) + @staticmethod + def _transform_for(s3: S3Source) -> DistributedMapS3SourceTransform | None: + """A prefix source flattens when a format is set, and lists keys when it is not.""" + if s3.prefix is None: + return None + if s3.fmt is None: + return DistributedMapS3SourceTransform.NONE + return DistributedMapS3SourceTransform.LOAD_AND_FLATTEN + + @classmethod + def from_source(cls, s3: S3Source) -> Self: + """Translate the S3 source, deriving the transform.""" + return cls( + bucket=s3.bucket, + key=s3.key, + key_prefix=s3.prefix, + transform=cls._transform_for(s3), + expected_bucket_owner=s3.expected_bucket_owner, + fmt=s3.fmt, + csv_format_options=DistributedMapCsvFormatOptions.from_source(s3), + ) + @dataclass(frozen=True) class DistributedMapReaderFunctionSourceConfig: @@ -739,16 +783,16 @@ def from_dict( class DistributedMapInlineSourceConfig: """Represent the items an inline map run source reads from.""" - items: tuple[str, ...] = () + items: list[str] = field(default_factory=list) def to_dict(self) -> MutableMapping[str, Any]: - return {"Items": list(self.items)} + return {"Items": self.items} @classmethod def from_dict( cls, data: MutableMapping[str, Any] ) -> DistributedMapInlineSourceConfig: - return cls(items=tuple(data.get("Items", ()))) + return cls(items=list(data.get("Items", []))) @dataclass(frozen=True) @@ -800,15 +844,60 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapSourceConfig 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_FAILURES: DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_FAILURES, + ProcessorResponseMode.ITEM_RESULTS: DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_RESULTS, +} + @dataclass(frozen=True) class DistributedMapProcessorConfig: """Represent a map run processor.""" function_name: str - function_response_types: tuple[DistributedMapFunctionResponseType, ...] | None = ( - None - ) + 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 @@ -835,9 +924,9 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapProcessorCon response_types = data.get("FunctionResponseTypes") return cls( function_name=data.get("FunctionName", ""), - function_response_types=tuple( + function_response_types=[ DistributedMapFunctionResponseType(t) for t in response_types - ) + ] if response_types else None, batch_size=data.get("BatchSize"), @@ -846,6 +935,27 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapProcessorCon 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: @@ -875,6 +985,21 @@ def from_dict( minimum_sample_size=data.get("MinimumSampleSize"), ) + @classmethod + def from_config(cls, completion_config: ConfigCompletionConfig) -> Self | None: + """Translate the completion config, or None when it carries no threshold.""" + if ( + completion_config.tolerated_failure_count is None + and completion_config.tolerated_failure_percentage is None + and completion_config.minimum_sample_size is None + ): + return None + 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: @@ -909,7 +1034,7 @@ class _DistributedMapDestinationEntry: """Hold the members the OnSuccess and OnFailure shapes have in common.""" type: DistributedMapDestinationType - include: tuple[DistributedMapDestinationInclude, ...] + include: list[DistributedMapDestinationInclude] s3_destination_config: DistributedMapS3DestinationConfig def to_dict(self) -> MutableMapping[str, Any]: @@ -924,22 +1049,59 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> Self: s3_raw = data.get("S3DestinationConfig") or {} return cls( type=DistributedMapDestinationType(data.get("Type", "S3")), - include=tuple( + 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: @@ -971,6 +1133,24 @@ def from_dict( else None, ) + @classmethod + def from_config(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: @@ -988,6 +1168,9 @@ def from_dict( return cls(mode=DistributedMapResultCollectionMode(data.get("Mode", "NONE"))) +_MAX_CONCURRENCY_LIMIT = 10000 + + @dataclass(frozen=True) class DistributedMapOptions: """Configuration options for starting a map run.""" 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 index 33ed11e15..eab7f3baa 100644 --- 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 @@ -7,12 +7,8 @@ from typing import TYPE_CHECKING, Any from aws_durable_execution_sdk_python.config import ( - DistributedMapProcessor, DistributedMapResultConfig, DistributedMapSource, - DistributedMapSourceFormat, - DistributedMapStatus, - ProcessorResponseMode, ) from aws_durable_execution_sdk_python.dmap.models import ( DistributedMapItemError, @@ -26,26 +22,12 @@ ) from aws_durable_execution_sdk_python.lambda_service import ( DistributedMapCompletionConfig, - DistributedMapCsvFormatOptions, - DistributedMapCsvHeaderLocation, DistributedMapDestinationConfig, - DistributedMapDestinationInclude, - DistributedMapDestinationType, - DistributedMapFunctionResponseType, - DistributedMapInlineSourceConfig, - DistributedMapOnFailureConfig, - DistributedMapOnSuccessConfig, DistributedMapOptions, DistributedMapProcessorConfig, - DistributedMapReaderFunctionSourceConfig, DistributedMapResultCollectionConfig, DistributedMapResultCollectionMode, - DistributedMapS3DestinationConfig, - DistributedMapS3SourceConfig, - DistributedMapS3SourceTransform, DistributedMapSourceConfig, - DistributedMapSourceType, - OperationStatus, OperationUpdate, ) from aws_durable_execution_sdk_python.operation.base import ( @@ -64,12 +46,15 @@ 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, @@ -80,293 +65,6 @@ # 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 -_UNLIMITED_RETRY_WIRE = -1 - -_RESPONSE_TYPE_FOR_MODE = { - ProcessorResponseMode.ITEM_FAILURES: DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_FAILURES, - ProcessorResponseMode.ITEM_RESULTS: DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_RESULTS, -} - - -def _build_inline_items( - items: tuple[Any, ...], - serdes: Any, - operation_id: str, - durable_execution_arn: str, -) -> tuple[str, ...]: - """Serialize each inline item, enforcing the 1 MB cap on the whole list.""" - serialized_items: list[str] = [] - for item in items: - serialized_items.append( - serialize( - serdes=serdes, - value=item, - operation_id=operation_id, - durable_execution_arn=durable_execution_arn, - ) - ) - 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 tuple(serialized_items) - - -def _build_completion_config( - completion_config: Any, -) -> DistributedMapCompletionConfig | None: - """Translate the completion config, or None when it carries no threshold.""" - if ( - completion_config.tolerated_failure_count is None - and completion_config.tolerated_failure_percentage is None - and completion_config.minimum_sample_size is None - ): - return None - return DistributedMapCompletionConfig( - tolerated_failure_count=completion_config.tolerated_failure_count, - tolerated_failure_percentage=completion_config.tolerated_failure_percentage, - minimum_sample_size=completion_config.minimum_sample_size, - ) - - -def _build_processor_config( - processor: DistributedMapProcessor, -) -> DistributedMapProcessorConfig: - """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 - max_retry_duration_seconds: 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 - if processor.max_retry_duration is not None: - max_retry_duration_seconds = processor.max_retry_duration.to_seconds() - return DistributedMapProcessorConfig( - 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=max_retry_duration_seconds, - durable_execution_name_prefix=processor.durable_execution_name_prefix, - ) - - -def _transform_for(s3: Any) -> DistributedMapS3SourceTransform | None: - """A prefix source flattens when a format is set, and lists keys when it is not.""" - if s3.prefix is None: - return None - if s3.fmt is None: - return DistributedMapS3SourceTransform.NONE - return DistributedMapS3SourceTransform.LOAD_AND_FLATTEN - - -def _build_s3_source_config(s3: Any) -> DistributedMapS3SourceConfig: - """Translate the S3 source, deriving the transform and CSV header location.""" - csv_format_options: DistributedMapCsvFormatOptions | None = None - if s3.fmt is DistributedMapSourceFormat.CSV: - csv_format_options = DistributedMapCsvFormatOptions( - header_location=( - DistributedMapCsvHeaderLocation.GIVEN - if s3.headers is not None - else DistributedMapCsvHeaderLocation.FIRST_ROW - ), - headers=s3.headers, - delimiter=s3.delimiter, - ) - return DistributedMapS3SourceConfig( - bucket=s3.bucket, - key=s3.key, - key_prefix=s3.prefix, - transform=_transform_for(s3), - expected_bucket_owner=s3.expected_bucket_owner, - fmt=s3.fmt, - csv_format_options=csv_format_options, - ) - - -def _destination_entry_fields( - destination: Any, *, second_include: DistributedMapDestinationInclude -) -> dict[str, Any]: - """Build the members the OnSuccess and OnFailure shapes share.""" - include: list[DistributedMapDestinationInclude] = [] - if destination.include_input: - include.append(DistributedMapDestinationInclude.INPUT) - if second_include is DistributedMapDestinationInclude.OUTPUT: - if destination.include_output: - include.append(second_include) - elif destination.include_error: - include.append(second_include) - return { - "type": DistributedMapDestinationType.S3, - "include": tuple(include), - "s3_destination_config": DistributedMapS3DestinationConfig( - bucket=destination.bucket, - key_prefix=destination.prefix, - expected_bucket_owner=destination.expected_bucket_owner, - ), - } - - -def _build_destination_config( - destination: Any, -) -> DistributedMapDestinationConfig | None: - """Translate the destination config, or None when neither side is set.""" - on_success = ( - DistributedMapOnSuccessConfig( - **_destination_entry_fields( - destination.on_success, - second_include=DistributedMapDestinationInclude.OUTPUT, - ) - ) - if destination.on_success is not None - else None - ) - on_failure = ( - DistributedMapOnFailureConfig( - **_destination_entry_fields( - destination.on_failure, - second_include=DistributedMapDestinationInclude.ERROR, - ) - ) - if destination.on_failure is not None - else None - ) - if on_success is None and on_failure is None: - return None - return DistributedMapDestinationConfig(on_success=on_success, on_failure=on_failure) - - -def _build_source_config( - source: DistributedMapSource | Sequence[Any], - operation_id: str, - durable_execution_arn: str, -) -> DistributedMapSourceConfig: - """Translate a source (typed or plain-list shorthand) into its source config.""" - if not isinstance(source, DistributedMapSource): - # A plain list is treated as an inline source with the default serializer. - serialized_items = _build_inline_items( - tuple(source), DEFAULT_JSON_SERDES, operation_id, durable_execution_arn - ) - return DistributedMapSourceConfig( - source_type=DistributedMapSourceType.INLINE, - inline_source_config=DistributedMapInlineSourceConfig( - items=serialized_items - ), - ) - - if source.inline_items is not None: - serialized_items = _build_inline_items( - source.inline_items, - source.inline_serdes or DEFAULT_JSON_SERDES, - operation_id, - durable_execution_arn, - ) - return DistributedMapSourceConfig( - source_type=DistributedMapSourceType.INLINE, - inline_source_config=DistributedMapInlineSourceConfig( - items=serialized_items - ), - max_items=source.max_items, - ) - if source.s3 is not None: - return DistributedMapSourceConfig( - source_type=DistributedMapSourceType.S3, - max_items=source.max_items, - s3_config=_build_s3_source_config(source.s3), - ) - if source.reader is not None: - reader_config = DistributedMapReaderFunctionSourceConfig( - function_name=source.reader.function_name - ) - if source.reader.initial_state is not None: - state = serialize( - serdes=source.reader.state_serdes or DEFAULT_JSON_SERDES, - value=source.reader.initial_state, - operation_id=operation_id, - durable_execution_arn=durable_execution_arn, - ) - if len(state.encode("utf-8")) > _READER_STATE_LIMIT: - msg = ( - f"reader initial_state exceeds the " - f"{_READER_STATE_LIMIT // 1024} KB limit" - ) - raise ValidationError(msg) - reader_config = DistributedMapReaderFunctionSourceConfig( - function_name=source.reader.function_name, initial_state=state - ) - return DistributedMapSourceConfig( - source_type=DistributedMapSourceType.READER_FUNCTION, - max_items=source.max_items, - reader_config=reader_config, - ) - msg = "Distributed map source has no configured items" - raise ExecutionError(msg) - - -def _build_distributed_map_options( - source: DistributedMapSource | Sequence[Any], - processor: DistributedMapProcessor, - max_concurrency: int, - config: DistributedMapConfig, - operation_id: str, - durable_execution_arn: str, -) -> DistributedMapOptions: - """Assemble the DistributedMapOptions payload from the operands and config.""" - result_collection = ( - DistributedMapResultCollectionConfig( - mode=DistributedMapResultCollectionMode.INLINE - ) - if isinstance(config, DistributedMapResultConfig) - else None - ) - return DistributedMapOptions( - max_concurrency=max_concurrency, - source=_build_source_config(source, operation_id, durable_execution_arn), - processor=_build_processor_config(processor), - destination=( - _build_destination_config(config.destination) - if config.destination is not None - else None - ), - completion_config=( - _build_completion_config(config.completion_config) - if config.completion_config is not None - else None - ), - result_collection=result_collection, - timeout_seconds=config.timeout.to_seconds() - if config.timeout is not None - else None, - ) - - -def _distributed_map_status_from_operation( - operation: Operation | None, -) -> DistributedMapStatus: - """Derive the map run status from the operation's terminal status.""" - op_status = operation.status if operation else None - resolved = ( - { - OperationStatus.SUCCEEDED: DistributedMapStatus.SUCCEEDED, - OperationStatus.FAILED: DistributedMapStatus.FAILED, - OperationStatus.STOPPED: DistributedMapStatus.STOPPED, - OperationStatus.TIMED_OUT: DistributedMapStatus.TIMED_OUT, - }.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 DistributedMapOperationExecutor(OperationExecutor[DistributedMapSummary]): @@ -404,45 +102,111 @@ def __init__( self.operation_identifier = operation_identifier self.config = config - def _resolve_summary(self, operation: Operation | None) -> DistributedMapSummary: - """Reconstruct the resolved summary/result from the terminal operation.""" - details = operation.distributed_map_details if operation else None - if details is None: - msg = "DISTRIBUTED_MAP operation succeeded but carried no DistributedMapDetails" - raise ExecutionError(msg) - status = _distributed_map_status_from_operation(operation) - completion_reason = details.completion_reason - if completion_reason is None: + 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"DISTRIBUTED_MAP operation ended {status.value} but carried no " - f"CompletionReason" + f"inline source exceeds the {_INLINE_SIZE_LIMIT // 1024 // 1024} MB limit " + f"(serialized size: {total} bytes)" ) - raise ExecutionError(msg) - if not isinstance(self.config, DistributedMapResultConfig): - return DistributedMapSummary( - status=status, - completion_reason=completion_reason, - success_count=details.success_count, - failure_count=details.failure_count, - unprocessed_count=details.unprocessed_count, - distributed_map_run_arn=details.distributed_map_run_arn, - completion_details=details.completion_details, - total_count=details.total_count, + 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 ) - return DistributedMapResult( - status=status, - completion_reason=completion_reason, - success_count=details.success_count, - failure_count=details.failure_count, - unprocessed_count=details.unprocessed_count, - distributed_map_run_arn=details.distributed_map_run_arn, - completion_details=details.completion_details, - total_count=details.total_count, - all=self._deserialize_items(details, self.config.result_serdes), + 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_config(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: Any + self, details: DistributedMapDetails, result_serdes: SerDes | None ) -> list[DistributedMapResultItem]: """Deserialize the service result items into customer result items.""" items: list[DistributedMapResultItem] = [] @@ -485,9 +249,7 @@ def check_result_status(self) -> CheckResult[DistributedMapSummary]: Raises: SuspendExecution: For STARTED operations waiting for completion """ - checkpointed_result: CheckpointedResult = self.state.get_checkpoint_result( - self.operation_identifier.operation_id - ) + checkpointed_result: CheckpointedResult = self._get_checkpoint_result() # Terminal success - build the summary/result from the operation if checkpointed_result.is_succeeded(): @@ -530,14 +292,7 @@ def check_result_status(self) -> CheckResult[DistributedMapSummary]: start_operation: OperationUpdate = ( OperationUpdate.create_distributed_map_start( identifier=self.operation_identifier, - distributed_map_options=_build_distributed_map_options( - source=self.source, - processor=self.processor, - max_concurrency=self.max_concurrency, - config=self.config, - operation_id=self.operation_identifier.operation_id, - durable_execution_arn=self.state.durable_execution_arn, - ), + distributed_map_options=self._build_options(), ) ) # Checkpoint map run START with blocking (is_sync=True). 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 9b70a9b01..33396bb37 100644 --- a/packages/aws-durable-execution-sdk-python/tests/config_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/config_test.py @@ -25,11 +25,11 @@ from aws_durable_execution_sdk_python.config import ( DistributedMapCompletionConfig, DistributedMapConfig, - S3Destination, - S3Source, + SuccessDestination, + FailureDestination, DistributedMapProcessor, DistributedMapSource, - ReaderSourceConfig, + ReaderSource, ) @@ -304,22 +304,24 @@ def test_retry_duration_out_of_range_rejected(): def test_expected_bucket_owner_must_be_12_digits(): with pytest.raises(ValidationError, match="12-digit"): - S3Source.json_lines("s3://b/k.jsonl", expected_bucket_owner="123") + 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"): - S3Source.csv("s3://b/data.csv", delimiter="BAR") + DistributedMapSource.s3_csv("s3://b/data.csv", delimiter="BAR") def test_csv_headers_duplicates_rejected(): with pytest.raises(ValidationError, match="duplicates"): - S3Source.csv("s3://b/k.csv", headers=["a", "a"]) + DistributedMapSource.s3_csv("s3://b/k.csv", headers=["a", "a"]) def test_json_lines_requires_key(): with pytest.raises(ValidationError, match="object key"): - S3Source.json_lines("s3://bucket-only") + DistributedMapSource.s3_json_lines("s3://bucket-only") def test_timeout_out_of_range_rejected(): @@ -389,12 +391,12 @@ def test_retry_duration_below_minimum_rejected(): def test_max_items_below_one_rejected(): with pytest.raises(ValidationError, match="at least 1"): - S3Source.json_lines("s3://b/k.jsonl", max_items=0) + DistributedMapSource.s3_json_lines("s3://b/k.jsonl", max_items=0) def test_csv_requires_object_key(): with pytest.raises(ValidationError, match="csv requires an S3 object key"): - S3Source.csv("s3://bucket") + DistributedMapSource.s3_csv("s3://bucket") def test_source_without_any_kind_rejected(): @@ -406,35 +408,37 @@ def test_source_with_two_kinds_rejected(): with pytest.raises(ValidationError, match="exactly one of"): DistributedMapSource( inline_items=("a",), - reader=ReaderSourceConfig(function_name="r"), + reader=ReaderSource(function_name="r"), ) def test_success_destination_all_false_rejected(): with pytest.raises(ValidationError, match="success destination must include"): - S3Destination.successes( + 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"): - S3Destination.failures("s3://out/bad", include_input=False, include_error=False) + 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"): - S3Source.json_lines("s3:///key.jsonl") + DistributedMapSource.s3_json_lines("s3:///key.jsonl") def test_non_s3_scheme_uri_rejected(): with pytest.raises(ValidationError, match="must start with s3://"): - S3Source.json_lines("http://foo/bar") + DistributedMapSource.s3_json_lines("http://foo/bar") def test_csv_empty_headers_rejected(): with pytest.raises(ValidationError, match="must be non-empty"): - S3Source.csv("s3://b/f.csv", headers=[]) + DistributedMapSource.s3_csv("s3://b/f.csv", headers=[]) def test_negative_retry_attempts_rejected(): 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 index 1f4675947..d03e4acdd 100644 --- a/packages/aws-durable-execution-sdk-python/tests/dmap/handlers_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/dmap/handlers_test.py @@ -4,14 +4,11 @@ import threading import time -from unittest.mock import MagicMock import pytest from aws_durable_execution_sdk_python.config import ( - DistributedMapConfig, - DistributedMapProcessor, - DistributedMapStatus, + ProcessorResponseMode, ) from aws_durable_execution_sdk_python.dmap.handlers import ( ReaderPage, @@ -20,85 +17,21 @@ durable_distributed_map_item_handler, distributed_map_reader, ) -from aws_durable_execution_sdk_python.exceptions import ExecutionError, ValidationError -from aws_durable_execution_sdk_python.lambda_service import ( - DistributedMapCompletionReason, - DistributedMapDetails, - Operation, - OperationStatus, - OperationType, -) -from aws_durable_execution_sdk_python.operation.dmap import ( - DistributedMapOperationExecutor, - _distributed_map_status_from_operation, -) +from aws_durable_execution_sdk_python.exceptions import ValidationError -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 test_status_derived_from_operation_when_details_omit_status(): - """Status comes from the operation, not from details.""" - details = DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, - success_count=3, - total_count=3, - ) - op = _terminal_operation(OperationStatus.SUCCEEDED, details) - assert _distributed_map_status_from_operation(op) is DistributedMapStatus.SUCCEEDED - - -def test_status_from_operation_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"): - _distributed_map_status_from_operation(op) - - -def test_resolved_summary_uses_operation_status_not_details(): - """_resolve_summary carries the derived status when details omit it.""" - details = DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, - success_count=2, - failure_count=0, - unprocessed_count=0, - total_count=2, - ) - executor = DistributedMapOperationExecutor( - source=["a", "b"], - processor=DistributedMapProcessor.batch("proc"), - max_concurrency=1, - state=MagicMock(), - operation_identifier=MagicMock(), - config=DistributedMapConfig(), - ) - summary = executor._resolve_summary( - _terminal_operation(OperationStatus.SUCCEEDED, details) - ) - assert summary.status is DistributedMapStatus.SUCCEEDED - assert summary.completion_reason is DistributedMapCompletionReason.ALL_COMPLETED - assert summary.success_count == 2 - - -def test_item_handler_invalid_report_rejected(): - with pytest.raises(ValidationError, match="report must be"): - distributed_map_item_handler(lambda x: x, report="bogus") +def test_item_handler_rejects_batch_response_mode(): + with pytest.raises(ValidationError, match="response_mode must be"): + distributed_map_item_handler( + lambda x: x, response_mode=ProcessorResponseMode.BATCH + ) -def test_durable_item_handler_invalid_report_rejected(): - with pytest.raises(ValidationError, match="report must be"): - durable_distributed_map_item_handler(lambda _ctx, item: item, report="bogus") +def test_durable_item_handler_rejects_batch_response_mode(): + with pytest.raises(ValidationError, match="response_mode must be"): + durable_distributed_map_item_handler( + lambda _ctx, item: item, response_mode=ProcessorResponseMode.BATCH + ) def test_item_handler_rejects_concurrency_below_one(): @@ -108,6 +41,13 @@ def test_item_handler_rejects_concurrency_below_one(): 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 @@ -186,7 +126,9 @@ def process(x): def test_item_handler_failures_form_reports_only_failures(): - handler = distributed_map_item_handler(lambda x: x, report="failures") + handler = distributed_map_item_handler( + lambda x: x, response_mode=ProcessorResponseMode.ITEM_FAILURES + ) resp = handler({"records": [{"itemId": "0", "body": "1"}]}) assert resp == {"batchItemFailures": []} @@ -212,33 +154,41 @@ def boom(_items): def test_reader_returns_items_and_next_state_then_exhausts(): - def read(state): + def read(state, max_items_per_page): if state is None: - return ReaderPage(items=[1, 2], next_state={"page": 1}) + 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, "maxItems": 10}) + first = handler({"state": None, "maxItemsPerPage": 10}) assert first["items"] == [1, 2] assert first["nextState"] == '{"page": 1}' - second = handler({"state": '{"page": 1}', "maxItems": 10}) + second = handler({"state": '{"page": 1}', "maxItemsPerPage": 10}) assert second["items"] == [3] assert "nextState" not in second -def test_reader_rejects_page_over_max_items(): - handler = distributed_map_reader(lambda _s: ReaderPage(items=[1, 2, 3])) - with pytest.raises(ValidationError, match="exceeding maxItems"): - handler({"state": None, "maxItems": 2}) +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: ReaderPage(items=[1], next_state="x" * 40_000) + lambda _s, _max: ReaderPage(items=[1], next_state="x" * 40_000) ) with pytest.raises(ValidationError, match="32 KB limit"): - handler({"state": None, "maxItems": 10}) + handler({"state": None, "maxItemsPerPage": 10}) def test_item_handler_rejects_non_processor_envelope(): @@ -248,6 +198,6 @@ def test_item_handler_rejects_non_processor_envelope(): def test_reader_rejects_non_reader_envelope(): - handler = distributed_map_reader(lambda _state: ReaderPage(items=[])) + 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 index 9048ae462..4eb215537 100644 --- a/packages/aws-durable-execution-sdk-python/tests/dmap/models_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/dmap/models_test.py @@ -13,9 +13,32 @@ DistributedMapItemError, DistributedMapResult, DistributedMapResultItem, + DistributedMapRunArn, DistributedMapSummary, + _to_distributed_map_status, ) -from aws_durable_execution_sdk_python.exceptions import DistributedMapError +from aws_durable_execution_sdk_python.exceptions import ( + DistributedMapError, + ExecutionError, + ValidationError, +) +from aws_durable_execution_sdk_python.lambda_service import ( + DistributedMapDetails, + 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, completion_reason, failure_count, items): @@ -96,6 +119,60 @@ def test_summary_has_failure_false_without_failures(): 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, + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=0, + failure_count=0, + unprocessed_count=0, + distributed_map_run_arn=arn, + ) + + +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_result_throw_if_error_raises_first_item_error(): result = _result( DistributedMapStatus.SUCCEEDED, @@ -171,3 +248,30 @@ def test_result_throw_if_error_clean_success_does_not_raise(): ], ) result.throw_if_error() + + +def test_result_from_operation_rejects_the_inherited_factory(): + """The summary factory cannot supply items, so it refuses rather than returning an empty result.""" + with pytest.raises(ValidationError, match="from_operation_and_items"): + DistributedMapResult.from_operation(None) + + +def test_status_derived_from_operation_when_details_omit_status(): + """Status comes from the operation, not from details.""" + details = DistributedMapDetails( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + 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 index f6336d610..15e058a5c 100644 --- 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 @@ -7,8 +7,11 @@ from __future__ import annotations import json +import threading +import time from unittest.mock import Mock, patch +from aws_durable_execution_sdk_python.config import ProcessorResponseMode from aws_durable_execution_sdk_python.dmap.handlers import ( durable_distributed_map_batch_handler, durable_distributed_map_item_handler, @@ -138,9 +141,55 @@ def test_durable_batch_handler_returns_value(): def test_durable_item_handler_failures_form(): handler = durable_distributed_map_item_handler( - lambda _ctx, item: item, report="failures" + lambda _ctx, item: item, response_mode=ProcessorResponseMode.ITEM_FAILURES ) result = _run(handler, [{"itemId": "0", "body": "1"}]) assert result["Status"] == InvocationStatus.SUCCEEDED.value data = json.loads(result["Result"]) assert data == {"batchItemFailures": []} + + +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 index b07917bb7..fda8326e5 100644 --- 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 @@ -115,6 +115,8 @@ def mock_checkpoint( 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, ) @@ -162,7 +164,6 @@ def handler(event, context: DurableContext) -> dict[str, Any]: # Replay with the run completed. _map_run_id, replay_event = _replay_event( { - "Status": "SUCCEEDED", "CompletionReason": "ALL_COMPLETED", "SuccessCount": 2, "FailureCount": 0, @@ -199,7 +200,6 @@ def handler(event, context: DurableContext) -> dict[str, Any]: _map_run_id, replay_event = _replay_event( { - "Status": "SUCCEEDED", "CompletionReason": "ALL_COMPLETED", "SuccessCount": 1, "FailureCount": 1, @@ -234,7 +234,6 @@ def handler(event, context: DurableContext) -> dict[str, Any]: _map_run_id, replay_event = _replay_event( { - "Status": "FAILED", "CompletionReason": "FAILURE_TOLERANCE_EXCEEDED", "SuccessCount": 0, "FailureCount": 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 index 305a507b1..2b5adcfc1 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/dmap_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dmap_test.py @@ -20,16 +20,16 @@ DistributedMapConfig, DistributedMapResultConfig, DistributedMapCsvDelimiter, - DistributedMapDestinationConfig, - InlineSource, - ReaderSource, - S3Destination, - S3Source, + DistributedMapDestination, + DistributedMapSource, + SuccessDestination, + FailureDestination, DistributedMapProcessor, ) from aws_durable_execution_sdk_python.exceptions import ( ExecutionError, DistributedMapError, + NonDeterministicExecutionError, SuspendExecution, ValidationError, ) @@ -94,6 +94,8 @@ def test_map_run_handler_already_succeeded(): 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( completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, @@ -137,6 +139,8 @@ def test_map_run_handler_resolves_non_success_without_raising(): 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( completion_reason=DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED, @@ -174,6 +178,8 @@ def test_map_run_handler_succeeded_no_details_raises(): 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, ) @@ -199,6 +205,8 @@ def test_map_run_handler_already_started(): 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 = ( @@ -226,6 +234,8 @@ def test_map_run_handler_new_operation(): 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) @@ -266,6 +276,8 @@ def test_map_run_handler_no_config(): 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) @@ -297,6 +309,8 @@ def test_map_run_immediate_response_get_checkpoint_result_called_twice(): 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) @@ -323,6 +337,8 @@ def test_map_run_immediate_response_create_checkpoint_is_sync_true(): 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) @@ -350,6 +366,8 @@ def test_map_run_immediate_response_immediate_success(): 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( completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, @@ -382,6 +400,8 @@ def test_map_run_immediate_response_no_immediate_response(): 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) @@ -408,6 +428,8 @@ def test_map_run_immediate_response_already_completed(): 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( completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, @@ -442,6 +464,8 @@ def test_map_run_handler_suspend_does_not_raise(mock_suspend): 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) @@ -494,6 +518,8 @@ def _new_op_state_calls(): Operation( operation_id="mrw", operation_type=OperationType.DISTRIBUTED_MAP, + sub_type=OperationSubType.DISTRIBUTED_MAP, + name="test_map_run", status=OperationStatus.STARTED, ) ) @@ -552,7 +578,9 @@ def test_s3_source_serializes_config(): """An S3 json_lines source serializes to an S3SourceConfig block.""" options = _start_options( _new_op_state_calls(), - source=S3Source.json_lines("s3://bucket/data.jsonl", max_items=500), + source=DistributedMapSource.s3_json_lines( + "s3://bucket/data.jsonl", max_items=500 + ), processor=DistributedMapProcessor.batch("proc"), max_concurrency=2, config=DistributedMapConfig(), @@ -571,9 +599,9 @@ def test_full_config_serializes_all_blocks(): completion_config=DistributedMapCompletionConfig.failure_percentage( 5, minimum_sample_size=200 ), - destination=DistributedMapDestinationConfig( - on_success=S3Destination.successes("s3://out/ok"), - on_failure=S3Destination.failures("s3://out/bad"), + destination=DistributedMapDestination( + on_success=SuccessDestination.from_uri("s3://out/ok"), + on_failure=FailureDestination.from_uri("s3://out/bad"), ), timeout=Duration.from_minutes(30), ) @@ -607,6 +635,8 @@ def test_result_config_returns_map_run_result_with_items(): 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( completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, @@ -653,6 +683,8 @@ def test_plain_config_returns_plain_summary(): 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( completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, @@ -677,11 +709,11 @@ def test_plain_config_returns_plain_summary(): def test_csv_source_header_location(): - """CSV headers map to GIVEN. expected_columns stays client-side under FIRST_ROW.""" + """Explicit headers map to GIVEN, omitting them maps to FIRST_ROW.""" # headers -> HeaderLocation GIVEN, headers sent given = _start_options( _new_op_state_calls(), - source=S3Source.csv("s3://b/data.csv", headers=["a", "b"]), + source=DistributedMapSource.s3_csv("s3://b/data.csv", headers=["a", "b"]), processor=DistributedMapProcessor.batch("proc"), max_concurrency=1, config=DistributedMapConfig(), @@ -694,7 +726,7 @@ def test_csv_source_header_location(): # no headers -> HeaderLocation FIRST_ROW first_row = _start_options( _new_op_state_calls(), - source=S3Source.csv("s3://b/data.csv"), + source=DistributedMapSource.s3_csv("s3://b/data.csv"), processor=DistributedMapProcessor.batch("proc"), max_concurrency=1, config=DistributedMapConfig(), @@ -711,7 +743,7 @@ def test_csv_source_header_location(): def test_csv_delimiter_accepts_enum(): opts = _start_options( _new_op_state_calls(), - source=S3Source.csv( + source=DistributedMapSource.s3_csv( "s3://b/data.csv", delimiter=DistributedMapCsvDelimiter.PIPE ), processor=DistributedMapProcessor.batch("proc"), @@ -737,7 +769,9 @@ 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=ReaderSource.from_function("reader", initial_state={"page": 0}), + source=DistributedMapSource.reader_function( + "reader", initial_state={"page": 0} + ), processor=DistributedMapProcessor.batch("proc"), max_concurrency=1, config=DistributedMapConfig(), @@ -750,7 +784,9 @@ def test_reader_state_serialized_and_capped(): with pytest.raises(ValidationError, match="32 KB limit"): _start_options( _new_op_state_calls(), - source=ReaderSource.from_function("reader", initial_state="x" * 40_000), + source=DistributedMapSource.reader_function( + "reader", initial_state="x" * 40_000 + ), processor=DistributedMapProcessor.batch("proc"), max_concurrency=1, config=DistributedMapConfig(), @@ -770,7 +806,7 @@ def deserialize(self, data, _serdes_context): # noqa: ANN001, ANN201 options = _start_options( _new_op_state_calls(), - source=InlineSource.of(["a", "b"], serdes=_UpperSerDes()), + source=DistributedMapSource.inline(["a", "b"], serdes=_UpperSerDes()), processor=DistributedMapProcessor.batch("proc"), max_concurrency=1, config=DistributedMapConfig(), @@ -803,8 +839,8 @@ def test_destination_only_success(): processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig( - destination=DistributedMapDestinationConfig( - on_success=S3Destination.successes("s3://out/ok") + destination=DistributedMapDestination( + on_success=SuccessDestination.from_uri("s3://out/ok") ) ), ) @@ -820,8 +856,8 @@ def test_destination_only_failure(): processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig( - destination=DistributedMapDestinationConfig( - on_failure=S3Destination.failures("s3://out/bad") + destination=DistributedMapDestination( + on_failure=FailureDestination.from_uri("s3://out/bad") ) ), ) @@ -833,7 +869,7 @@ def test_destination_only_failure(): def test_s3_objects_source_config(): opts = _start_options( _new_op_state_calls(), - source=S3Source.objects("s3://b/prefix/"), + source=DistributedMapSource.s3_objects("s3://b/prefix/"), processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig(), @@ -847,7 +883,7 @@ def test_s3_objects_source_config(): def test_s3_flattened_json_lines_source_config(): opts = _start_options( _new_op_state_calls(), - source=S3Source.flattened_json_lines("s3://b/prefix/"), + source=DistributedMapSource.s3_flattened_json_lines("s3://b/prefix/"), processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig(), @@ -896,7 +932,7 @@ def test_empty_completion_config_omitted(): def test_objects_whole_bucket_prefix_config(): options = _start_options( _new_op_state_calls(), - source=S3Source.objects("s3://bucket"), + source=DistributedMapSource.s3_objects("s3://bucket"), processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig(), @@ -910,7 +946,9 @@ def test_objects_whole_bucket_prefix_config(): def test_flattened_csv_source_config(): options = _start_options( _new_op_state_calls(), - source=S3Source.flattened_csv("s3://b/prefix", headers=["a", "b"]), + source=DistributedMapSource.s3_flattened_csv( + "s3://b/prefix", headers=["a", "b"] + ), processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig(), @@ -925,7 +963,7 @@ def test_flattened_csv_source_config(): def test_reader_source_without_initial_state_omits_state(): options = _start_options( _new_op_state_calls(), - source=ReaderSource.from_function("reader"), + source=DistributedMapSource.reader_function("reader"), processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig(), @@ -948,7 +986,7 @@ def deserialize(self, data, _serdes_context): # noqa: ANN001, ANN201, ARG002 options = _start_options( _new_op_state_calls(), - source=InlineSource.of([1, 2], serdes=_PlainTextSerDes()), + source=DistributedMapSource.inline([1, 2], serdes=_PlainTextSerDes()), processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig(), @@ -966,8 +1004,8 @@ def test_success_destination_include_input_and_owner(): processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig( - destination=DistributedMapDestinationConfig( - on_success=S3Destination.successes( + destination=DistributedMapDestination( + on_success=SuccessDestination.from_uri( "s3://out/ok", include_input=True, include_output=True, @@ -989,8 +1027,8 @@ def test_failure_destination_error_only_and_owner(): processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig( - destination=DistributedMapDestinationConfig( - on_failure=S3Destination.failures( + destination=DistributedMapDestination( + on_failure=FailureDestination.from_uri( "s3://out/bad", include_input=False, include_error=True, @@ -1011,6 +1049,8 @@ 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( completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, @@ -1049,7 +1089,7 @@ def test_operation_to_dict_round_trip_preserves_results(): def test_options_full_round_trip(): options_dict = _start_options( _new_op_state_calls(), - source=S3Source.csv("s3://b/f.csv", headers=["a"]), + source=DistributedMapSource.s3_csv("s3://b/f.csv", headers=["a"]), processor=DistributedMapProcessor.item_results( "proc", batch_size=5, @@ -1059,9 +1099,9 @@ def test_options_full_round_trip(): ), max_concurrency=3, config=DistributedMapResultConfig( - destination=DistributedMapDestinationConfig( - on_success=S3Destination.successes("s3://o/ok"), - on_failure=S3Destination.failures("s3://o/bad"), + 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), @@ -1070,9 +1110,9 @@ def test_options_full_round_trip(): 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_RESULTS, - ) + assert parsed.processor.function_response_types == [ + DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_RESULTS + ] assert parsed.processor.max_retry_attempts == -1 assert parsed.processor.durable_execution_name_prefix == "pfx" assert parsed.destination is not None @@ -1089,7 +1129,7 @@ def test_options_full_round_trip(): def test_csv_first_row_no_headers(): options = _start_options( _new_op_state_calls(), - source=S3Source.csv("s3://b/f.csv"), + source=DistributedMapSource.s3_csv("s3://b/f.csv"), processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig(), @@ -1103,7 +1143,7 @@ def test_csv_first_row_no_headers(): def test_s3_source_expected_bucket_owner(): options = _start_options( _new_op_state_calls(), - source=S3Source.json_lines( + source=DistributedMapSource.s3_json_lines( "s3://b/k.jsonl", expected_bucket_owner="123456789012" ), processor=DistributedMapProcessor.batch("p"), @@ -1119,7 +1159,7 @@ def test_empty_destination_config_omitted(): source=["a"], processor=DistributedMapProcessor.batch("p"), max_concurrency=1, - config=DistributedMapConfig(destination=DistributedMapDestinationConfig()), + config=DistributedMapConfig(destination=DistributedMapDestination()), ) assert "Destination" not in options @@ -1131,8 +1171,8 @@ def test_success_destination_input_only(): processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig( - destination=DistributedMapDestinationConfig( - on_success=S3Destination.successes( + destination=DistributedMapDestination( + on_success=SuccessDestination.from_uri( "s3://out/ok", include_input=True, include_output=False ) ) @@ -1148,8 +1188,8 @@ def test_failure_destination_input_only(): processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig( - destination=DistributedMapDestinationConfig( - on_failure=S3Destination.failures( + destination=DistributedMapDestination( + on_failure=FailureDestination.from_uri( "s3://out/bad", include_input=True, include_error=False ) ) @@ -1161,7 +1201,9 @@ def test_failure_destination_input_only(): def test_reader_source_options_round_trip(): options_dict = _start_options( _new_op_state_calls(), - source=ReaderSource.from_function("reader", initial_state={"page": 0}), + source=DistributedMapSource.reader_function( + "reader", initial_state={"page": 0} + ), processor=DistributedMapProcessor.batch("p"), max_concurrency=1, config=DistributedMapConfig(), @@ -1176,6 +1218,8 @@ 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( completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, @@ -1202,6 +1246,8 @@ def test_operation_level_terminal_failure_without_details_raises(status): 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 = ( @@ -1232,6 +1278,8 @@ def test_terminal_failure_with_details_resolves_with_summary(status, expected): 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( completion_reason=DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED, @@ -1264,6 +1312,8 @@ def test_terminal_without_completion_reason_raises(): 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( success_count=0, @@ -1304,6 +1354,8 @@ def deserialize(self, data, _serdes_context): # noqa: ANN001, ANN201 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( completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, @@ -1348,6 +1400,8 @@ def test_falsy_output_round_trips(value): 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( completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, @@ -1386,3 +1440,67 @@ def test_processor_retry_duration_only(): 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( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + success_count=2, + failure_count=0, + unprocessed_count=0, + total_count=2, + ), + ) + executor = DistributedMapOperationExecutor( + source=["a", "b"], + processor=DistributedMapProcessor.batch("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.completion_reason is DistributedMapCompletionReason.ALL_COMPLETED + 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( + completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + 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.batch("proc"), + max_concurrency=1, + state=mock_state, + operation_identifier=_identifier("dmap-id"), + ) From 2f2f73ae1af82e0e01a1028f3016da64fb51ae86 Mon Sep 17 00:00:00 2001 From: nvasiu Date: Thu, 15 Oct 2026 22:32:04 +0000 Subject: [PATCH 3/3] refactor(dmap): align with functional spec - Merge the processor response types into REPORT_BATCH_ITEM_OUTCOMES - Split the S3 source config into the three blocks selected by Type - Replace CompletionReason and CompletionDetails with SourceExhausted and Error - Carry the shared ErrorObject on per-item errors - Return one entry per item from get_results and get_errors - Succeed an item whose handler returns None with no output - Fail only the item whose output cannot be serialized - Serialize reader items with an item serdes - Raise ValidationError for a configuration the service rejects - Reject an empty completion config and a bad max_retry_attempts - Replace _resolve_terminal with _summary_fields and drop the override - Name the operation's status in the missing-details error - Rename from_config to from_destination --- .../__init__.py | 4 - .../config.py | 170 +++++--- .../context.py | 11 +- .../dmap/handlers.py | 214 +++++----- .../dmap/models.py | 151 ++++--- .../lambda_service.py | 268 ++++++++---- .../operation/dmap.py | 37 +- .../tests/config_test.py | 83 +++- .../tests/context_test.py | 24 +- .../tests/dmap/handlers_test.py | 174 ++++++-- .../tests/dmap/models_test.py | 194 +++++++-- .../tests/e2e/dmap_helpers_int_test.py | 82 +++- .../tests/e2e/dmap_int_test.py | 20 +- .../tests/lambda_service_test.py | 71 +++- .../tests/operation/dmap_test.py | 399 +++++++++++++----- 15 files changed, 1357 insertions(+), 545 deletions(-) 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 dc38cc342..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,6 @@ CompletionItemStatus, CompletionOutcome, CompletionStatus, - DistributedMapCompletionReason, DistributedMapItemStatus, DistributedMapStatus, ParallelBranch, @@ -39,7 +38,6 @@ durable_distributed_map_item_handler, ) from aws_durable_execution_sdk_python.dmap.models import ( - DistributedMapItemError, DistributedMapResult, DistributedMapResultItem, DistributedMapSummary, @@ -87,9 +85,7 @@ "CompletionItemStatus", "CompletionOutcome", "CompletionStatus", - "DistributedMapCompletionReason", "DistributedMapError", - "DistributedMapItemError", "DistributedMapItemStatus", "DistributedMapResult", "DistributedMapResultItem", 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 f6c951162..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 @@ -23,21 +23,6 @@ class DistributedMapItemStatus(Enum): SUCCEEDED = "SUCCEEDED" FAILED = "FAILED" - - -class DistributedMapCompletionReason(Enum): - ALL_COMPLETED = "ALL_COMPLETED" - ITEM_LIMIT_REACHED = "ITEM_LIMIT_REACHED" - STOPPED = "STOPPED" - TIMED_OUT = "TIMED_OUT" - FAILURE_TOLERANCE_EXCEEDED = "FAILURE_TOLERANCE_EXCEEDED" - SOURCE_FAILED = "SOURCE_FAILED" - DESTINATION_FAILED = "DESTINATION_FAILED" - INLINE_RESULT_LIMIT_EXCEEDED = "INLINE_RESULT_LIMIT_EXCEEDED" - INVALID_CONFIGURATION = "INVALID_CONFIGURATION" - QUOTA_EXCEEDED = "QUOTA_EXCEEDED" - KMS_ACCESS_DENIED = "KMS_ACCESS_DENIED" - INTERNAL_ERROR = "INTERNAL_ERROR" UNKNOWN_TO_SDK_VERSION = "UNKNOWN_TO_SDK_VERSION" @@ -47,12 +32,19 @@ class DistributedMapSourceFormat(Enum): 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 per-item outcomes.""" + """How a processor function reports item outcomes.""" BATCH = "BATCH" - ITEM_FAILURES = "ITEM_FAILURES" - ITEM_RESULTS = "ITEM_RESULTS" + ITEM_OUTCOMES = "ITEM_OUTCOMES" class DistributedMapCsvDelimiter(Enum): @@ -576,16 +568,31 @@ def _validate_bucket_owner(value: str | None) -> None: 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 is non-empty with no duplicates.""" + """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. @@ -607,7 +614,7 @@ class DistributedMapCompletionConfig: """Failure-tolerance configuration for a map run.""" tolerated_failure_count: int | None = None - tolerated_failure_percentage: float | None = None + tolerated_failure_percentage: int | None = None minimum_sample_size: int | None = None def __post_init__(self) -> None: @@ -620,6 +627,17 @@ def __post_init__(self) -> None: "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 @@ -635,6 +653,14 @@ def __post_init__(self) -> None: 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 ): @@ -657,7 +683,7 @@ def failure_count(count: int) -> DistributedMapCompletionConfig: @staticmethod def failure_percentage( - percentage: float, *, minimum_sample_size: int | None = None + percentage: int, *, minimum_sample_size: int | None = None ) -> DistributedMapCompletionConfig: """Abort once the failure rate exceeds this percentage.""" return DistributedMapCompletionConfig( @@ -686,11 +712,19 @@ def __post_init__(self) -> None: ): msg = f"batch_size must be between 1 and 10000, got: {self.batch_size}" raise ValidationError(msg) - if isinstance(self.max_retry_attempts, int) and self.max_retry_attempts < 0: + 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 non-negative or " + "max_retry_attempts must be a non-negative integer or " "DistributedMapProcessor.UNLIMITED, " - f"got: {self.max_retry_attempts}" + f"got: {attempts!r}" ) raise ValidationError(msg) if self.max_retry_duration is not None and not ( @@ -711,7 +745,7 @@ def __post_init__(self) -> None: raise ValidationError(msg) @classmethod - def batch( + def report_batch_outcome( cls, name: str, *, @@ -720,7 +754,7 @@ def batch( max_retry_duration: Duration | None = None, durable_execution_name_prefix: str | None = None, ) -> DistributedMapProcessor: - """Processor that reports a single pass/fail outcome for the whole batch, with no per-item results.""" + """Processor that reports one outcome for the whole batch. Items have no outputs.""" return cls( function_name=name, response_mode=ProcessorResponseMode.BATCH, @@ -731,7 +765,7 @@ def batch( ) @classmethod - def item_failures( + def report_item_outcomes( cls, name: str, *, @@ -740,30 +774,10 @@ def item_failures( max_retry_duration: Duration | None = None, durable_execution_name_prefix: str | None = None, ) -> DistributedMapProcessor: - """Processor that reports the ids of failed items, with all others marked succeeded.""" + """Processor that may report a failure or an output per item. Unlisted items succeed.""" return cls( function_name=name, - response_mode=ProcessorResponseMode.ITEM_FAILURES, - 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 item_results( - 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 the results (output or error) for every item.""" - return cls( - function_name=name, - response_mode=ProcessorResponseMode.ITEM_RESULTS, + response_mode=ProcessorResponseMode.ITEM_OUTCOMES, batch_size=batch_size, max_retry_attempts=max_retry_attempts, max_retry_duration=max_retry_duration, @@ -789,9 +803,10 @@ def _parse_delimiter( class S3Source: """Resolved S3 source.""" + source_type: DistributedMapS3SourceType bucket: str key: str | None = None - prefix: str | None = None + key_prefix: str | None = None fmt: DistributedMapSourceFormat | None = None delimiter: DistributedMapCsvDelimiter | None = None headers: list[str] | None = None @@ -800,6 +815,29 @@ class S3Source: 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) @@ -825,8 +863,13 @@ class DistributedMapSource: reader: ReaderSource | None = None def __post_init__(self) -> None: - if self.max_items is not None and self.max_items < 1: - msg = f"max_items must be at least 1, got: {self.max_items}" + 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 @@ -869,6 +912,7 @@ def s3_json_lines( return cls( max_items=max_items, s3=S3Source( + source_type=DistributedMapS3SourceType.OBJECT, bucket=bucket, key=key, fmt=DistributedMapSourceFormat.JSON_LINES, @@ -893,6 +937,7 @@ def s3_json_array( return cls( max_items=max_items, s3=S3Source( + source_type=DistributedMapS3SourceType.OBJECT, bucket=bucket, key=key, fmt=DistributedMapSourceFormat.JSON_ARRAY, @@ -919,6 +964,7 @@ def s3_csv( return cls( max_items=max_items, s3=S3Source( + source_type=DistributedMapS3SourceType.OBJECT, bucket=bucket, key=key, fmt=DistributedMapSourceFormat.CSV, @@ -936,14 +982,15 @@ def s3_objects( expected_bucket_owner: str | None = None, max_items: int | None = None, ) -> DistributedMapSource: - """Read each object under a prefix as one item.""" + """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, - prefix=prefix or "", + key_prefix=prefix or "", expected_bucket_owner=expected_bucket_owner, ), ) @@ -956,14 +1003,15 @@ def s3_flattened_json_lines( expected_bucket_owner: str | None = None, max_items: int | None = None, ) -> DistributedMapSource: - """Read a prefix, flattening each object's lines into items.""" + """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, - prefix=prefix or "", + key_prefix=prefix or "", fmt=DistributedMapSourceFormat.JSON_LINES, expected_bucket_owner=expected_bucket_owner, ), @@ -977,14 +1025,15 @@ def s3_flattened_json_array( expected_bucket_owner: str | None = None, max_items: int | None = None, ) -> DistributedMapSource: - """Read a prefix, flattening each object's JSON array elements into items.""" + """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, - prefix=prefix or "", + key_prefix=prefix or "", fmt=DistributedMapSourceFormat.JSON_ARRAY, expected_bucket_owner=expected_bucket_owner, ), @@ -1000,14 +1049,15 @@ def s3_flattened_csv( expected_bucket_owner: str | None = None, max_items: int | None = None, ) -> DistributedMapSource: - """Read a prefix, flattening each object's records into items.""" + """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, - prefix=prefix or "", + key_prefix=prefix or "", fmt=DistributedMapSourceFormat.CSV, delimiter=_parse_delimiter(delimiter), headers=list(headers) if headers is not None else None, 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 e21e85ba0..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 @@ -52,7 +52,6 @@ OperationIdNamespace, ) from aws_durable_execution_sdk_python.lambda_service import ( - _MAX_CONCURRENCY_LIMIT, OperationSubType, OperationType, ) @@ -732,7 +731,8 @@ def distributed_map( 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 + max_concurrency: Maximum concurrent processor invocations. Required, + with no service default. name: Optional name for the operation config: Optional run-level configuration @@ -743,11 +743,8 @@ def distributed_map( if not isinstance(source, (DistributedMapSource, list, tuple)): msg = "source must be a DistributedMapSource or a list/tuple of items" raise ValidationError(msg) - if not 1 <= max_concurrency <= _MAX_CONCURRENCY_LIMIT: - msg = ( - f"max_concurrency must be between 1 and {_MAX_CONCURRENCY_LIMIT}, " - f"got: {max_concurrency}" - ) + 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() 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 index 924085918..6d14edab5 100644 --- 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 @@ -17,10 +17,10 @@ BatchItemStatus, CompletionConfig, MapConfig, - ProcessorResponseMode, ) 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, @@ -73,10 +73,10 @@ def to_dict(self) -> MutableMapping[str, Any]: @dataclass(frozen=True) class ItemResult: - """One item's output, reported by an item_results processor.""" + """One item's output, reported by an item_outcomes processor.""" item_identifier: str - output: Any = None + output: str | None = None @classmethod def from_dict(cls, data: MutableMapping[str, Any]) -> ItemResult: @@ -85,70 +85,82 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> ItemResult: ) def to_dict(self) -> MutableMapping[str, Any]: - return {"itemIdentifier": self.item_identifier, "output": self.output} + 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_failures or item_results processor.""" + """One item's failure, reported by an item_outcomes processor.""" item_identifier: str - error_type: str = "" - error_message: str = "" + error: ErrorObject | None = None @classmethod - def from_exception(cls, item_identifier: str, exc: BaseException) -> ItemFailure: + def from_exception(cls, item_identifier: str, exc: Exception) -> ItemFailure: return cls( - item_identifier=item_identifier, - error_type=type(exc).__name__, - error_message=str(exc), + item_identifier=item_identifier, error=ErrorObject.from_exception(exc) ) @classmethod def from_dict(cls, data: MutableMapping[str, Any]) -> ItemFailure: - error = data.get("error") or {} + error = data.get("error") return cls( item_identifier=data.get("itemIdentifier", ""), - error_type=error.get("errorType", ""), - error_message=error.get("errorMessage", ""), + 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]: - return { - "itemIdentifier": self.item_identifier, - "error": { - "errorType": self.error_type, - "errorMessage": self.error_message, - }, - } + 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 None ``results`` reports failures only.""" + """What an item handler returns: a failure or a result for each item.""" failures: list[ItemFailure] = field(default_factory=list) - results: list[ItemResult] | None = None + results: list[ItemResult] = field(default_factory=list) @classmethod def from_dict(cls, data: MutableMapping[str, Any]) -> ItemHandlerResponse: - results = data.get("batchItemResults") return cls( failures=[ ItemFailure.from_dict(f) for f in data.get("batchItemFailures", []) ], - results=[ItemResult.from_dict(r) for r in results] - if results is not None - else None, + results=[ItemResult.from_dict(r) for r in data.get("batchItemResults", [])], ) def to_dict(self) -> MutableMapping[str, Any]: - result: MutableMapping[str, Any] = { - "batchItemFailures": [f.to_dict() for f in self.failures] + return { + "batchItemFailures": [f.to_dict() for f in self.failures], + "batchItemResults": [r.to_dict() for r in self.results], } - if self.results is not None: - result["batchItemResults"] = [r.to_dict() for r in self.results] - return result @dataclass(frozen=True) @@ -180,7 +192,7 @@ def to_dict(self) -> MutableMapping[str, Any]: class ReaderResponse: """What a reader function returns. A None ``next_state`` exhausts the source.""" - items: list[Any] = field(default_factory=list) + items: list[str] = field(default_factory=list) next_state: str | None = None @classmethod @@ -199,20 +211,13 @@ def _to_item(serdes: SerDes[Any], body: str, ctx: SerDesContext) -> Any: return serdes.deserialize(body, ctx) -def _to_output(serdes: SerDes[Any], value: Any, ctx: SerDesContext) -> str: - """Serialize an item's output for the Output field.""" +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_response_mode(response_mode: ProcessorResponseMode) -> None: - if response_mode is ProcessorResponseMode.BATCH: - msg = ( - "response_mode must be ProcessorResponseMode.ITEM_RESULTS or " - "ProcessorResponseMode.ITEM_FAILURES, got: BATCH" - ) - raise ValidationError(msg) - - def _validate_concurrency(concurrency: int) -> None: if concurrency < 1: msg = f"concurrency must be at least 1, got: {concurrency}" @@ -225,14 +230,17 @@ def distributed_map_item_handler( item_serdes: SerDes[Any] | None = None, result_serdes: SerDes[Any] | None = None, concurrency: int = 1, - response_mode: ProcessorResponseMode = ProcessorResponseMode.ITEM_RESULTS, ) -> Callable[..., Any]: - """Wrap a Lambda to be used as an item_results processor. + """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. - Pass ``response_mode=ProcessorResponseMode.ITEM_FAILURES`` for an - item_failures 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``. + 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 @@ -244,9 +252,7 @@ def distributed_map_item_handler( item_serdes=item_serdes, result_serdes=result_serdes, concurrency=concurrency, - response_mode=response_mode, ) - _validate_response_mode(response_mode) _validate_concurrency(concurrency) serdes = item_serdes or DEFAULT_JSON_SERDES out_serdes = result_serdes or DEFAULT_JSON_SERDES @@ -255,11 +261,13 @@ def handler(event: dict[str, Any], _context: Any = None) -> dict[str, Any]: ctx = SerDesContext() records = ProcessorEvent.from_dict(event).records - def run(record: ProcessorRecord) -> Any: - return func(_to_item(serdes, record.body, ctx)) + 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[Any] = [None] * len(records) - errors: list[BaseException | None] = [None] * len(records) + 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(): @@ -274,27 +282,19 @@ def run(record: ProcessorRecord) -> Any: err = errors[i] if err is not None: failures.append(ItemFailure.from_exception(record.item_id, err)) - elif response_mode is ProcessorResponseMode.ITEM_RESULTS: + else: results.append( - ItemResult( - item_identifier=record.item_id, - output=_to_output(out_serdes, outputs[i], ctx), - ) + ItemResult(item_identifier=record.item_id, output=outputs[i]) ) - response = ItemHandlerResponse( - failures=failures, - results=results - if response_mode is ProcessorResponseMode.ITEM_RESULTS - else None, - ) + response = ItemHandlerResponse(failures=failures, results=results) return dict(response.to_dict()) return handler def distributed_map_batch_handler( - func: Callable[[list[Any]], Any] | None = None, + func: Callable[[list[Any]], None] | None = None, *, item_serdes: SerDes[Any] | None = None, ) -> Callable[..., Any]: @@ -303,14 +303,17 @@ def distributed_map_batch_handler( ``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) -> Any: + def handler(event: dict[str, Any], _context: Any = None) -> None: ctx = SerDesContext() - return func( + func( [ _to_item(serdes, record.body, ctx) for record in ProcessorEvent.from_dict(event).records @@ -323,6 +326,7 @@ def handler(event: dict[str, Any], _context: Any = None) -> Any: 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. @@ -331,16 +335,25 @@ def distributed_map_reader( 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, state_serdes=state_serdes) - serdes = state_serdes or DEFAULT_JSON_SERDES + 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 = ( - serdes.deserialize(reader_event.state, ctx) + state_sd.deserialize(reader_event.state, ctx) if reader_event.state is not None else None ) @@ -353,14 +366,16 @@ def handler(event: dict[str, Any], _context: Any = None) -> dict[str, Any]: ) 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 = serdes.serialize(page.next_state, ctx) + 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=page.items, next_state=next_state).to_dict()) + return dict(ReaderResponse(items=items, next_state=next_state).to_dict()) return handler @@ -371,12 +386,17 @@ def durable_distributed_map_item_handler( item_serdes: SerDes[Any] | None = None, result_serdes: SerDes[Any] | None = None, concurrency: int = 1, - response_mode: ProcessorResponseMode = ProcessorResponseMode.ITEM_RESULTS, ) -> 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. @@ -387,9 +407,7 @@ def durable_distributed_map_item_handler( item_serdes=item_serdes, result_serdes=result_serdes, concurrency=concurrency, - response_mode=response_mode, ) - _validate_response_mode(response_mode) _validate_concurrency(concurrency) serdes = item_serdes or DEFAULT_JSON_SERDES out_serdes = result_serdes or DEFAULT_JSON_SERDES @@ -416,36 +434,25 @@ def per_item(inner_ctx: Any, body: Any, _index: int, _inputs: Any) -> Any: for bi in batch.all: item_id = records[bi.index].item_id if bi.status is BatchItemStatus.SUCCEEDED: - if response_mode is ProcessorResponseMode.ITEM_RESULTS: - results.append( - ItemResult( - item_identifier=item_id, - output=_to_output(out_serdes, bi.result, ctx), - ) - ) + # 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: - err = bi.error - failures.append( - ItemFailure( - item_identifier=item_id, - error_type=(err.type or "") if err else "", - error_message=(err.message or "") if err else "", - ) - ) + failures.append(ItemFailure(item_identifier=item_id, error=bi.error)) - response = ItemHandlerResponse( - failures=failures, - results=results - if response_mode is ProcessorResponseMode.ITEM_RESULTS - else None, - ) + response = ItemHandlerResponse(failures=failures, results=results) return dict(response.to_dict()) return handler def durable_distributed_map_batch_handler( - func: Callable[..., Any] | None = None, + func: Callable[..., None] | None = None, *, item_serdes: SerDes[Any] | None = None, ) -> Callable[..., Any]: @@ -453,6 +460,9 @@ def durable_distributed_map_batch_handler( 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( @@ -461,9 +471,9 @@ def durable_distributed_map_batch_handler( serdes = item_serdes or DEFAULT_JSON_SERDES @durable_execution - def handler(event: dict[str, Any], context: Any) -> Any: + def handler(event: dict[str, Any], context: Any) -> None: ctx = SerDesContext() - return func( + func( context, [ _to_item(serdes, record.body, ctx) 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 index feeb4e95f..b39919726 100644 --- 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 @@ -3,26 +3,21 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, NoReturn +from typing import TYPE_CHECKING, Any, Self, TypedDict from aws_durable_execution_sdk_python.exceptions import ( DistributedMapError, ExecutionError, - ValidationError, ) from aws_durable_execution_sdk_python.config import ( - DistributedMapCompletionReason, DistributedMapItemStatus, DistributedMapStatus, ) -from aws_durable_execution_sdk_python.lambda_service import OperationStatus +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 ( - DistributedMapDetails, - Operation, - ) + from aws_durable_execution_sdk_python.lambda_service import Operation _STATUS_FOR_OPERATION_STATUS = { OperationStatus.SUCCEEDED: DistributedMapStatus.SUCCEEDED, @@ -47,22 +42,52 @@ def _to_distributed_map_status(operation: Operation | None) -> DistributedMapSta return resolved -def _resolve_terminal( - operation: Operation | None, -) -> tuple[DistributedMapDetails, DistributedMapStatus, DistributedMapCompletionReason]: - """Read the details, status and completion reason a terminal operation must carry.""" +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: - msg = "DISTRIBUTED_MAP operation succeeded but carried no DistributedMapDetails" + # 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) - if details.completion_reason is None: + # 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"CompletionReason" + f"SourceExhausted" ) raise ExecutionError(msg) - return details, status, details.completion_reason + 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/" @@ -98,28 +123,22 @@ class DistributedMapSummary: """ status: DistributedMapStatus - completion_reason: DistributedMapCompletionReason success_count: int failure_count: int unprocessed_count: int distributed_map_run_arn: str | None = None - completion_details: str | None = None + source_exhausted: bool | None = None + error: ErrorObject | None = None total_count: int | None = None @classmethod - def from_operation(cls, operation: Operation | None) -> DistributedMapSummary: - """Rebuild the summary from a terminal DISTRIBUTED_MAP operation.""" - details, status, completion_reason = _resolve_terminal(operation) - return cls( - status=status, - completion_reason=completion_reason, - success_count=details.success_count, - failure_count=details.failure_count, - unprocessed_count=details.unprocessed_count, - distributed_map_run_arn=details.distributed_map_run_arn, - completion_details=details.completion_details, - total_count=details.total_count, - ) + 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: @@ -137,12 +156,16 @@ def has_failure(self) -> bool: def throw_if_error(self) -> None: """Raise :class:`DistributedMapError` on any non-success outcome.""" if self.status is not DistributedMapStatus.SUCCEEDED: - detail = f", {self.completion_details}" if self.completion_details else "" - msg = ( - f"Map run ended {self.status.value} " - f"(reason: {self.completion_reason.value}{detail})" + 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, ) - raise DistributedMapError(msg) if self.failure_count > 0: msg = ( f"Map run succeeded but {self.failure_count} item(s) permanently failed" @@ -150,14 +173,6 @@ def throw_if_error(self) -> None: raise DistributedMapError(msg) -@dataclass(frozen=True) -class DistributedMapItemError: - """Error for a single failed map run item.""" - - error_type: str - error_message: str - - @dataclass(frozen=True) class DistributedMapResultItem: """Outcome of a single map run item.""" @@ -165,7 +180,7 @@ class DistributedMapResultItem: item_id: str status: DistributedMapItemStatus output: Any | None = None - error: DistributedMapItemError | None = None + error: ErrorObject | None = None @dataclass(frozen=True) @@ -179,34 +194,12 @@ class DistributedMapResult(DistributedMapSummary): all: list[DistributedMapResultItem] = field(default_factory=list) - @classmethod - def from_operation(cls, operation: Operation | None) -> NoReturn: - """Always raises. Use :meth:`from_operation_and_items` instead.""" - msg = ( - "DistributedMapResult carries per-item results, which from_operation " - "cannot supply. Call DistributedMapResult.from_operation_and_items " - "with the deserialized items, or DistributedMapSummary.from_operation " - "for the counts alone." - ) - raise ValidationError(msg) - @classmethod def from_operation_and_items( cls, operation: Operation | None, items: list[DistributedMapResultItem] - ) -> DistributedMapResult: + ) -> Self: """Rebuild the result from a terminal operation and its deserialized items.""" - details, status, completion_reason = _resolve_terminal(operation) - return cls( - status=status, - completion_reason=completion_reason, - success_count=details.success_count, - failure_count=details.failure_count, - unprocessed_count=details.unprocessed_count, - distributed_map_run_arn=details.distributed_map_run_arn, - completion_details=details.completion_details, - total_count=details.total_count, - all=items, - ) + return cls(**_summary_fields(operation), all=items) def succeeded(self) -> list[DistributedMapResultItem]: """Return the items that succeeded.""" @@ -223,20 +216,19 @@ def failed(self) -> list[DistributedMapResultItem]: ] def get_results(self) -> list[Any]: - """Return the outputs of the succeeded items.""" + """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 - and item.output is not None ] - def get_errors(self) -> list[DistributedMapItemError]: - """Return the errors of the failed items.""" + 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 and item.error is not None + if item.status is DistributedMapItemStatus.FAILED ] def throw_if_error(self) -> None: @@ -245,9 +237,14 @@ def throw_if_error(self) -> None: failed = self.failed() if failed: first = failed[0] - if first.error is not None: - msg = f"{first.error.error_type}: {first.error.error_message}" - raise DistributedMapError(msg) + 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/lambda_service.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py index 2bd80c38a..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 @@ -26,11 +26,11 @@ DistributedMapCompletionConfig as ConfigCompletionConfig, ) from aws_durable_execution_sdk_python.config import ( - DistributedMapCompletionReason, DistributedMapCsvDelimiter, DistributedMapDestination, DistributedMapItemStatus, DistributedMapProcessor, + DistributedMapS3SourceType, DistributedMapSourceFormat, FailureDestination, ProcessorResponseMode, @@ -138,19 +138,13 @@ class DistributedMapSourceType(Enum): READER_FUNCTION = "READER_FUNCTION" -class DistributedMapS3SourceTransform(Enum): - NONE = "NONE" - LOAD_AND_FLATTEN = "LOAD_AND_FLATTEN" - - class DistributedMapCsvHeaderLocation(Enum): FIRST_ROW = "FIRST_ROW" GIVEN = "GIVEN" class DistributedMapFunctionResponseType(Enum): - REPORT_BATCH_ITEM_FAILURES = "REPORT_BATCH_ITEM_FAILURES" - REPORT_BATCH_ITEM_RESULTS = "REPORT_BATCH_ITEM_RESULTS" + REPORT_BATCH_ITEM_OUTCOMES = "REPORT_BATCH_ITEM_OUTCOMES" class DistributedMapDestinationType(Enum): @@ -480,10 +474,22 @@ class DistributedMapResultItem: @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, data.get("Status", ""), "item status" + 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, @@ -503,9 +509,9 @@ def to_dict(self) -> MutableMapping[str, Any]: @dataclass(frozen=True) class DistributedMapDetails: - completion_reason: DistributedMapCompletionReason | None = None distributed_map_run_arn: str | None = None - completion_details: 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 @@ -515,22 +521,11 @@ class DistributedMapDetails: @classmethod def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapDetails: results_raw = data.get("Results") - reason_raw = data.get("CompletionReason") - # CompletionReason is absent while a run is in flight, so parse it only - # when present rather than requiring it. + error_raw = data.get("Error") return cls( - completion_reason=( - _parse_enum( - DistributedMapCompletionReason, - reason_raw, - "completion reason", - unknown_fallback=DistributedMapCompletionReason.UNKNOWN_TO_SDK_VERSION, - ) - if reason_raw is not None - else None - ), distributed_map_run_arn=data.get("DistributedMapRunArn"), - completion_details=data.get("CompletionDetails"), + 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), @@ -687,72 +682,175 @@ def from_source(cls, s3: S3Source) -> Self | None: @dataclass(frozen=True) -class DistributedMapS3SourceConfig: - """Represent an S3 distributed map source config.""" +class DistributedMapS3ObjectConfig: + """Represent the object an OBJECT S3 source reads.""" - bucket: str - key: str | None = None - key_prefix: str | None = None - transform: DistributedMapS3SourceTransform | None = None - expected_bucket_owner: str | None = None - fmt: DistributedMapSourceFormat | None = None + key: str + fmt: DistributedMapSourceFormat csv_format_options: DistributedMapCsvFormatOptions | None = None def to_dict(self) -> MutableMapping[str, Any]: - result: MutableMapping[str, Any] = {"Bucket": self.bucket} - if self.key is not None: - result["Key"] = self.key - if self.key_prefix is not None: - result["KeyPrefix"] = self.key_prefix - if self.transform is not None: - result["Transform"] = self.transform.value - if self.expected_bucket_owner is not None: - result["ExpectedBucketOwner"] = self.expected_bucket_owner - if self.fmt is not None: - result["Format"] = self.fmt.value + 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]) -> DistributedMapS3SourceConfig: - transform = data.get("Transform") - fmt = data.get("Format") + def from_dict(cls, data: MutableMapping[str, Any]) -> DistributedMapS3ObjectConfig: csv_raw = data.get("CsvFormatOptions") return cls( - bucket=data.get("Bucket", ""), - key=data.get("Key"), - key_prefix=data.get("KeyPrefix"), - transform=DistributedMapS3SourceTransform(transform) - if transform is not None + key=data.get("Key", ""), + fmt=DistributedMapSourceFormat(data["Format"]), + csv_format_options=DistributedMapCsvFormatOptions.from_dict(csv_raw) + if csv_raw is not None else None, - expected_bucket_owner=data.get("ExpectedBucketOwner"), - fmt=DistributedMapSourceFormat(fmt) if fmt 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, ) - @staticmethod - def _transform_for(s3: S3Source) -> DistributedMapS3SourceTransform | None: - """A prefix source flattens when a format is set, and lists keys when it is not.""" - if s3.prefix is None: - return None - if s3.fmt is None: - return DistributedMapS3SourceTransform.NONE - return DistributedMapS3SourceTransform.LOAD_AND_FLATTEN + +@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, deriving the transform.""" + """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, - key=s3.key, - key_prefix=s3.prefix, - transform=cls._transform_for(s3), expected_bucket_owner=s3.expected_bucket_owner, - fmt=s3.fmt, - csv_format_options=DistributedMapCsvFormatOptions.from_source(s3), + listed_objects_config=DistributedMapS3ListedObjectsConfig( + key_prefix=s3.key_prefix or "", + fmt=s3.fmt or DistributedMapSourceFormat.JSON_LINES, + csv_format_options=csv_options, + ), ) @@ -887,8 +985,7 @@ def create_reader( _UNLIMITED_RETRY_WIRE = -1 _RESPONSE_TYPE_FOR_MODE = { - ProcessorResponseMode.ITEM_FAILURES: DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_FAILURES, - ProcessorResponseMode.ITEM_RESULTS: DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_RESULTS, + ProcessorResponseMode.ITEM_OUTCOMES: DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_OUTCOMES, } @@ -962,7 +1059,7 @@ class DistributedMapCompletionConfig: """Represent a map run completion (failure-tolerance) config.""" tolerated_failure_count: int | None = None - tolerated_failure_percentage: float | None = None + tolerated_failure_percentage: int | None = None minimum_sample_size: int | None = None def to_dict(self) -> MutableMapping[str, Any]: @@ -986,14 +1083,8 @@ def from_dict( ) @classmethod - def from_config(cls, completion_config: ConfigCompletionConfig) -> Self | None: - """Translate the completion config, or None when it carries no threshold.""" - if ( - completion_config.tolerated_failure_count is None - and completion_config.tolerated_failure_percentage is None - and completion_config.minimum_sample_size is None - ): - return None + 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, @@ -1134,7 +1225,7 @@ def from_dict( ) @classmethod - def from_config(cls, destination: DistributedMapDestination) -> Self | None: + 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 @@ -1168,9 +1259,6 @@ def from_dict( return cls(mode=DistributedMapResultCollectionMode(data.get("Mode", "NONE"))) -_MAX_CONCURRENCY_LIMIT = 10000 - - @dataclass(frozen=True) class DistributedMapOptions: """Configuration options for starting a map run.""" @@ -1800,17 +1888,17 @@ def to_dict(self) -> MutableMapping[str, Any]: "FailureCount": self.distributed_map_details.failure_count, "UnprocessedCount": self.distributed_map_details.unprocessed_count, } - if self.distributed_map_details.completion_reason is not None: - distributed_map_details_dict["CompletionReason"] = ( - self.distributed_map_details.completion_reason.value - ) if self.distributed_map_details.distributed_map_run_arn: distributed_map_details_dict["DistributedMapRunArn"] = ( self.distributed_map_details.distributed_map_run_arn ) - if self.distributed_map_details.completion_details: - distributed_map_details_dict["CompletionDetails"] = ( - self.distributed_map_details.completion_details + if self.distributed_map_details.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"] = ( 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 index eab7f3baa..097cdd628 100644 --- 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 @@ -11,12 +11,13 @@ DistributedMapSource, ) from aws_durable_execution_sdk_python.dmap.models import ( - DistributedMapItemError, DistributedMapResult, DistributedMapResultItem, DistributedMapSummary, ) from aws_durable_execution_sdk_python.exceptions import ( + INVALID_PARAMETER_VALUE_EXCEPTION, + CheckpointError, ExecutionError, ValidationError, ) @@ -67,6 +68,21 @@ _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. @@ -172,7 +188,7 @@ def _build_options(self) -> DistributedMapOptions: source=self._build_source_config(), processor=DistributedMapProcessorConfig.from_processor(self.processor), destination=( - DistributedMapDestinationConfig.from_config(config.destination) + DistributedMapDestinationConfig.from_destination(config.destination) if config.destination is not None else None ), @@ -219,20 +235,12 @@ def _deserialize_items( operation_id=self.operation_identifier.operation_id, durable_execution_arn=self.state.durable_execution_arn, ) - error = ( - DistributedMapItemError( - error_type=entry.error.type or "", - error_message=entry.error.message or "", - ) - if entry.error is not None - else None - ) items.append( DistributedMapResultItem( item_id=entry.item_id, status=entry.status, output=output, - error=error, + error=entry.error, ) ) return items @@ -297,7 +305,12 @@ def check_result_status(self) -> CheckResult[DistributedMapSummary]: ) # Checkpoint map run START with blocking (is_sync=True). # Must ensure the map run is recorded before suspending execution. - self.state.create_checkpoint(operation_update=start_operation, is_sync=True) + 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", 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 33396bb37..a5183478d 100644 --- a/packages/aws-durable-execution-sdk-python/tests/config_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/config_test.py @@ -297,9 +297,13 @@ def test_completion_decision_rejects_outcome_when_not_complete(): def test_retry_duration_out_of_range_rejected(): with pytest.raises(ValidationError, match="between 1 minute and 6 hours"): - DistributedMapProcessor.batch("p", max_retry_duration=Duration.from_seconds(30)) + DistributedMapProcessor.report_batch_outcome( + "p", max_retry_duration=Duration.from_seconds(30) + ) with pytest.raises(ValidationError, match="between 1 minute and 6 hours"): - DistributedMapProcessor.batch("p", max_retry_duration=Duration.from_hours(7)) + DistributedMapProcessor.report_batch_outcome( + "p", max_retry_duration=Duration.from_hours(7) + ) def test_expected_bucket_owner_must_be_12_digits(): @@ -331,7 +335,7 @@ def test_timeout_out_of_range_rejected(): def test_empty_function_name_rejected(): with pytest.raises(ValidationError, match="non-empty"): - DistributedMapProcessor.batch("") + DistributedMapProcessor.report_batch_outcome("") def test_valid_function_references_accepted(): @@ -343,12 +347,12 @@ def test_valid_function_references_accepted(): "arn:aws:lambda:us-east-1:123456789012:function:my-func:1", ): # Should not raise. - DistributedMapProcessor.batch(ref) + DistributedMapProcessor.report_batch_outcome(ref) def test_function_name_over_max_length_rejected(): with pytest.raises(ValidationError, match="at most 170 characters"): - DistributedMapProcessor.batch("f" * 171) + DistributedMapProcessor.report_batch_outcome("f" * 171) def test_completion_count_and_percentage_mutually_exclusive(): @@ -358,9 +362,17 @@ def test_completion_count_and_percentage_mutually_exclusive(): ) +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(minimum_sample_size=10) + DistributedMapCompletionConfig( + tolerated_failure_count=1, minimum_sample_size=10 + ) def test_completion_negative_count_rejected(): @@ -373,6 +385,12 @@ def test_completion_percentage_out_of_range_rejected(): 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( @@ -386,14 +404,42 @@ def test_completion_failure_count_factory(): def test_retry_duration_below_minimum_rejected(): with pytest.raises(ValidationError, match="1 minute and 6 hours"): - DistributedMapProcessor.batch("p", max_retry_duration=Duration.from_seconds(30)) + DistributedMapProcessor.report_batch_outcome( + "p", max_retry_duration=Duration.from_seconds(30) + ) def test_max_items_below_one_rejected(): - with pytest.raises(ValidationError, match="at least 1"): + 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") @@ -443,14 +489,29 @@ def test_csv_empty_headers_rejected(): def test_negative_retry_attempts_rejected(): with pytest.raises(ValidationError, match="non-negative"): - DistributedMapProcessor.batch("p", max_retry_attempts=-5) + 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.batch("p", batch_size=0) + 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.batch("p", durable_execution_name_prefix="x" * 37) + 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 6c01127f5..cb4128914 100644 --- a/packages/aws-durable-execution-sdk-python/tests/context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/context_test.py @@ -1096,17 +1096,17 @@ def test_wait_with_time_less_than_one(mock_executor_class): # region distributed_map -@pytest.mark.parametrize("max_concurrency", [0, -1, 10001, 50000]) -def test_map_run_rejects_out_of_range_max_concurrency(max_concurrency: int): - """Test distributed_map rejects a max_concurrency outside the service range.""" +@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 between 1 and 10000" + ValidationError, match="max_concurrency must be greater than zero" ): context.distributed_map( ["a"], - DistributedMapProcessor.batch("test_processor"), + DistributedMapProcessor.report_batch_outcome("test_processor"), max_concurrency=max_concurrency, ) @@ -1122,7 +1122,7 @@ def test_distributed_map_rejects_non_list_source(source): with pytest.raises(ValidationError, match="list/tuple"): context.distributed_map( source, - DistributedMapProcessor.batch("test_processor"), + DistributedMapProcessor.report_batch_outcome("test_processor"), max_concurrency=1, ) @@ -1137,7 +1137,7 @@ def test_distributed_map_basic(mock_executor_class): mock_state.durable_execution_arn = ( "arn:aws:durable:us-east-1:123456789012:execution/test" ) - processor = DistributedMapProcessor.batch("test_processor") + processor = DistributedMapProcessor.report_batch_outcome("test_processor") context = create_test_context(state=mock_state) expected_operation_id = next(operation_id_sequence()) @@ -1168,7 +1168,7 @@ def test_distributed_map_with_name_and_config(mock_executor_class): mock_state.durable_execution_arn = ( "arn:aws:durable:us-east-1:123456789012:execution/test" ) - processor = DistributedMapProcessor.batch("test_processor") + processor = DistributedMapProcessor.report_batch_outcome("test_processor") config = DistributedMapConfig() context = create_test_context(state=mock_state) @@ -1206,7 +1206,7 @@ def test_distributed_map_with_parent_id(mock_executor_class): mock_state.durable_execution_arn = ( "arn:aws:durable:us-east-1:123456789012:execution/test" ) - processor = DistributedMapProcessor.batch("test_processor") + 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 @@ -1240,7 +1240,7 @@ def test_distributed_map_increments_counter(mock_executor_class): mock_state.durable_execution_arn = ( "arn:aws:durable:us-east-1:123456789012:execution/test" ) - processor = DistributedMapProcessor.batch("test_processor") + 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 @@ -1272,7 +1272,7 @@ def test_distributed_map_defaults_config_when_none(mock_executor_class): mock_state.durable_execution_arn = ( "arn:aws:durable:us-east-1:123456789012:execution/test" ) - processor = DistributedMapProcessor.batch("test_processor") + processor = DistributedMapProcessor.report_batch_outcome("test_processor") context = create_test_context(state=mock_state) context.distributed_map(["a"], processor, max_concurrency=1) @@ -1292,7 +1292,7 @@ def test_distributed_map_returns_process_result(mock_executor_class): mock_state.durable_execution_arn = ( "arn:aws:durable:us-east-1:123456789012:execution/test" ) - processor = DistributedMapProcessor.batch("test_processor") + processor = DistributedMapProcessor.report_batch_outcome("test_processor") context = create_test_context(state=mock_state) result = context.distributed_map(["a"], processor, max_concurrency=1) 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 index d03e4acdd..02c892cd7 100644 --- a/packages/aws-durable-execution-sdk-python/tests/dmap/handlers_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/dmap/handlers_test.py @@ -2,36 +2,24 @@ from __future__ import annotations +import json import threading import time +from typing import Any import pytest -from aws_durable_execution_sdk_python.config import ( - ProcessorResponseMode, -) 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 ValidationError - - -def test_item_handler_rejects_batch_response_mode(): - with pytest.raises(ValidationError, match="response_mode must be"): - distributed_map_item_handler( - lambda x: x, response_mode=ProcessorResponseMode.BATCH - ) - - -def test_durable_item_handler_rejects_batch_response_mode(): - with pytest.raises(ValidationError, match="response_mode must be"): - durable_distributed_map_item_handler( - lambda _ctx, item: item, response_mode=ProcessorResponseMode.BATCH - ) +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(): @@ -125,22 +113,102 @@ def process(x): ] -def test_item_handler_failures_form_reports_only_failures(): - handler = distributed_map_item_handler( - lambda x: x, response_mode=ProcessorResponseMode.ITEM_FAILURES +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": []} + 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 = [] - handler = distributed_map_batch_handler(lambda items: seen.extend(items) or "done") + + 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"}]} ) - == "done" + is None ) assert seen == [1, 2] @@ -161,20 +229,72 @@ def read(state, max_items_per_page): handler = distributed_map_reader(read) first = handler({"state": None, "maxItemsPerPage": 10}) - assert first["items"] == [1, 2] + assert first["items"] == ["1", "2"] assert first["nextState"] == '{"page": 1}' second = handler({"state": '{"page": 1}', "maxItemsPerPage": 10}) - assert second["items"] == [3] + 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] + assert handler({"state": None, "maxItemsPerPage": 3})["items"] == ["0", "1", "2"] def test_reader_rejects_page_over_max_items_per_page(): 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 index 4eb215537..edbca77e3 100644 --- a/packages/aws-durable-execution-sdk-python/tests/dmap/models_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/dmap/models_test.py @@ -5,12 +5,10 @@ import pytest from aws_durable_execution_sdk_python.config import ( - DistributedMapCompletionReason, DistributedMapItemStatus, DistributedMapStatus, ) from aws_durable_execution_sdk_python.dmap.models import ( - DistributedMapItemError, DistributedMapResult, DistributedMapResultItem, DistributedMapRunArn, @@ -20,10 +18,10 @@ from aws_durable_execution_sdk_python.exceptions import ( DistributedMapError, ExecutionError, - ValidationError, ) from aws_durable_execution_sdk_python.lambda_service import ( DistributedMapDetails, + ErrorObject, Operation, OperationStatus, OperationType, @@ -41,10 +39,9 @@ def _terminal_operation( ) -def _result(status, completion_reason, failure_count, items): +def _result(status, failure_count, items): return DistributedMapResult( status=status, - completion_reason=completion_reason, success_count=len(items) - failure_count, failure_count=failure_count, unprocessed_count=0, @@ -55,7 +52,6 @@ def _result(status, completion_reason, failure_count, items): def test_summary_throw_if_error(): ok = DistributedMapSummary( status=DistributedMapStatus.SUCCEEDED, - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, success_count=2, failure_count=0, unprocessed_count=0, @@ -64,7 +60,6 @@ def test_summary_throw_if_error(): failed = DistributedMapSummary( status=DistributedMapStatus.FAILED, - completion_reason=DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED, success_count=0, failure_count=1, unprocessed_count=0, @@ -74,7 +69,6 @@ def test_summary_throw_if_error(): succeeded_with_failures = DistributedMapSummary( status=DistributedMapStatus.SUCCEEDED, - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, success_count=1, failure_count=1, unprocessed_count=0, @@ -91,12 +85,11 @@ def test_map_run_result_succeeded_failed_filters(): DistributedMapResultItem( item_id="1", status=DistributedMapItemStatus.FAILED, - error=DistributedMapItemError(error_type="E", error_message="boom"), + error=ErrorObject(message="boom", type="E", data=None, stack_trace=None), ), ] result = DistributedMapResult( status=DistributedMapStatus.SUCCEEDED, - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, success_count=1, failure_count=1, unprocessed_count=0, @@ -105,13 +98,78 @@ def test_map_run_result_succeeded_failed_filters(): assert [i.item_id for i in result.succeeded()] == ["0"] assert [i.item_id for i in result.failed()] == ["1"] assert result.get_results() == [1] - assert result.get_errors()[0].error_message == "boom" + 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, - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, success_count=0, failure_count=0, unprocessed_count=0, @@ -128,11 +186,11 @@ def test_summary_has_failure_false_without_failures(): def _summary_with_arn(arn): return DistributedMapSummary( status=DistributedMapStatus.SUCCEEDED, - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, success_count=0, failure_count=0, unprocessed_count=0, distributed_map_run_arn=arn, + source_exhausted=True, ) @@ -173,10 +231,90 @@ 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, - DistributedMapCompletionReason.ALL_COMPLETED, 1, [ DistributedMapResultItem( @@ -185,7 +323,9 @@ def test_result_throw_if_error_raises_first_item_error(): DistributedMapResultItem( item_id="1", status=DistributedMapItemStatus.FAILED, - error=DistributedMapItemError(error_type="E", error_message="boom"), + error=ErrorObject( + message="boom", type="E", data=None, stack_trace=None + ), ), ], ) @@ -196,7 +336,6 @@ def test_result_throw_if_error_raises_first_item_error(): def test_result_throw_if_error_names_item_without_detail(): result = _result( DistributedMapStatus.SUCCEEDED, - DistributedMapCompletionReason.ALL_COMPLETED, 1, [ DistributedMapResultItem( @@ -211,7 +350,6 @@ def test_result_throw_if_error_names_item_without_detail(): def test_result_throw_if_error_falls_back_to_summary_when_no_items(): result = DistributedMapResult( status=DistributedMapStatus.SUCCEEDED, - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, success_count=0, failure_count=2, unprocessed_count=0, @@ -224,7 +362,6 @@ def test_result_throw_if_error_falls_back_to_summary_when_no_items(): def test_result_throw_if_error_run_level_failure(): result = DistributedMapResult( status=DistributedMapStatus.FAILED, - completion_reason=DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED, success_count=0, failure_count=1, unprocessed_count=0, @@ -237,7 +374,6 @@ def test_result_throw_if_error_run_level_failure(): def test_result_throw_if_error_clean_success_does_not_raise(): result = DistributedMapResult( status=DistributedMapStatus.SUCCEEDED, - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, success_count=1, failure_count=0, unprocessed_count=0, @@ -250,16 +386,26 @@ def test_result_throw_if_error_clean_success_does_not_raise(): result.throw_if_error() -def test_result_from_operation_rejects_the_inherited_factory(): - """The summary factory cannot supply items, so it refuses rather than returning an empty result.""" - with pytest.raises(ValidationError, match="from_operation_and_items"): - DistributedMapResult.from_operation(None) +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( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, success_count=3, total_count=3, ) 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 index 15e058a5c..114c8a29f 100644 --- 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 @@ -9,13 +9,14 @@ import json import threading import time +from typing import Any from unittest.mock import Mock, patch -from aws_durable_execution_sdk_python.config import ProcessorResponseMode 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) @@ -27,6 +28,7 @@ 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" @@ -130,23 +132,87 @@ def process(_ctx, item): assert data["batchItemFailures"][0]["error"]["errorType"] == "ValueError" -def test_durable_batch_handler_returns_value(): - handler = durable_distributed_map_batch_handler( - lambda _ctx, items: {"count": len(items)} +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"]) == {"count": 2} + 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) -def test_durable_item_handler_failures_form(): handler = durable_distributed_map_item_handler( - lambda _ctx, item: item, response_mode=ProcessorResponseMode.ITEM_FAILURES + 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": []} + assert data == { + "batchItemFailures": [], + "batchItemResults": [{"itemIdentifier": "0"}], + } def test_durable_item_handler_runs_items_one_at_a_time_by_default(): 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 index fda8326e5..3c2d9f59e 100644 --- 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 @@ -147,7 +147,7 @@ def test_map_run_suspends_then_resumes_with_summary(): def handler(event, context: DurableContext) -> dict[str, Any]: summary = context.distributed_map( ["a", "b"], - DistributedMapProcessor.batch("proc"), + DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=2, ) return { @@ -164,7 +164,7 @@ def handler(event, context: DurableContext) -> dict[str, Any]: # Replay with the run completed. _map_run_id, replay_event = _replay_event( { - "CompletionReason": "ALL_COMPLETED", + "SourceExhausted": True, "SuccessCount": 2, "FailureCount": 0, "UnprocessedCount": 0, @@ -188,19 +188,21 @@ def test_map_run_result_config_returns_items(): def handler(event, context: DurableContext) -> dict[str, Any]: result = context.distributed_map( ["a", "b"], - DistributedMapProcessor.item_results("proc"), + DistributedMapProcessor.report_item_outcomes("proc"), max_concurrency=2, config=DistributedMapResultConfig(), ) assert isinstance(result, DistributedMapResult) return { "results": result.get_results(), - "errors": [e.error_message for e in result.get_errors()], + "errors": [ + e.message if e is not None else None for e in result.get_errors() + ], } _map_run_id, replay_event = _replay_event( { - "CompletionReason": "ALL_COMPLETED", + "SourceExhausted": True, "SuccessCount": 1, "FailureCount": 1, "UnprocessedCount": 0, @@ -226,7 +228,7 @@ def test_map_run_throw_if_error_fails_execution(): def handler(event, context: DurableContext) -> dict[str, Any]: summary = context.distributed_map( ["a"], - DistributedMapProcessor.batch("proc"), + DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=1, ) summary.throw_if_error() @@ -234,7 +236,11 @@ def handler(event, context: DurableContext) -> dict[str, Any]: _map_run_id, replay_event = _replay_event( { - "CompletionReason": "FAILURE_TOLERANCE_EXCEEDED", + "SourceExhausted": False, + "Error": { + "ErrorType": "DistributedMapFailure", + "ErrorMessage": "failure tolerance exceeded", + }, "SuccessCount": 0, "FailureCount": 1, "UnprocessedCount": 0, 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 aa5a10e03..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, @@ -1276,7 +1278,7 @@ def test_operation_from_dict_no_options(): def test_operation_from_dict_in_flight_distributed_map_details_parsed_without_status(): - """In-flight DistributedMapDetails (no Status/CompletionReason) parses without raising, leaving both None.""" + """In-flight DistributedMapDetails (no SourceExhausted) parses without raising.""" data = { "Id": "dmap-id", "Type": "DISTRIBUTED_MAP", @@ -1288,7 +1290,8 @@ def test_operation_from_dict_in_flight_distributed_map_details_parsed_without_st } operation = Operation.from_dict(data) assert operation.distributed_map_details is not None - assert operation.distributed_map_details.completion_reason is 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 @@ -1300,14 +1303,14 @@ def test_operation_from_dict_terminal_distributed_map_details_parsed(): "Type": "DISTRIBUTED_MAP", "Status": "SUCCEEDED", "DistributedMapDetails": { - "CompletionReason": "ALL_COMPLETED", + "SourceExhausted": True, "TotalCount": 3, "SuccessCount": 3, }, } operation = Operation.from_dict(data) assert operation.distributed_map_details is not None - assert operation.distributed_map_details.completion_reason.value == "ALL_COMPLETED" + assert operation.distributed_map_details.source_exhausted is True def test_operation_from_dict_individual_options(): @@ -3078,7 +3081,7 @@ def test_operation_to_dict_omits_absent_context_error_and_replay_children(): def test_map_run_details_from_dict_parses_results(): data = { "Status": "SUCCEEDED", - "CompletionReason": "ALL_COMPLETED", + "SourceExhausted": True, "SuccessCount": 1, "FailureCount": 1, "UnprocessedCount": 0, @@ -3095,19 +3098,40 @@ def test_map_run_details_from_dict_parses_results(): } 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_completion_reason_parses_with_none(): - # In-flight runs omit CompletionReason too; parsing must succeed with None. +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.completion_reason is None + 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}} @@ -3125,3 +3149,34 @@ def test_result_item_wire_round_trip(): 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 index 2b5adcfc1..4d1b7f319 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/dmap_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dmap_test.py @@ -13,7 +13,6 @@ ) from aws_durable_execution_sdk_python.config import ( Duration, - DistributedMapCompletionReason, DistributedMapItemStatus, DistributedMapStatus, DistributedMapCompletionConfig, @@ -28,7 +27,9 @@ ) from aws_durable_execution_sdk_python.exceptions import ( ExecutionError, + CheckpointError, DistributedMapError, + DurableApiErrorCategory, NonDeterministicExecutionError, SuspendExecution, ValidationError, @@ -66,7 +67,7 @@ def distributed_map_handler( if not config: config = DistributedMapConfig() if isinstance(processor, str): - processor = DistributedMapProcessor.batch(processor) + processor = DistributedMapProcessor.report_batch_outcome(processor) executor = DistributedMapOperationExecutor( source=source, processor=processor, @@ -98,7 +99,7 @@ def test_map_run_handler_already_succeeded(): name="test_map_run", status=OperationStatus.SUCCEEDED, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + source_exhausted=True, success_count=5, failure_count=0, unprocessed_count=0, @@ -119,7 +120,8 @@ def test_map_run_handler_already_succeeded(): ) assert result.status is DistributedMapStatus.SUCCEEDED - assert result.completion_reason is DistributedMapCompletionReason.ALL_COMPLETED + 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 @@ -143,10 +145,17 @@ def test_map_run_handler_resolves_non_success_without_raising(): name="test_map_run", status=OperationStatus.FAILED, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED, + 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 = ( @@ -162,10 +171,9 @@ def test_map_run_handler_resolves_non_success_without_raising(): ) assert result.status is DistributedMapStatus.FAILED - assert ( - result.completion_reason - is DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED - ) + 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 @@ -370,7 +378,7 @@ def test_map_run_immediate_response_immediate_success(): name="test_map_run", status=OperationStatus.SUCCEEDED, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + source_exhausted=True, success_count=1, ), ) @@ -432,7 +440,7 @@ def test_map_run_immediate_response_already_completed(): name="test_map_run", status=OperationStatus.SUCCEEDED, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + source_exhausted=True, ), ) mock_state.get_checkpoint_result.return_value = ( @@ -452,6 +460,63 @@ def test_map_run_immediate_response_already_completed(): 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" ) @@ -526,27 +591,39 @@ def _new_op_state_calls(): return [not_found, started] -def test_processor_item_failures_sets_response_types(): - """item_failures serializes FunctionResponseTypes=REPORT_BATCH_ITEM_FAILURES.""" +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.item_failures("proc", batch_size=25), + 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_FAILURES"] + 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.item_results( + processor=DistributedMapProcessor.report_item_outcomes( "proc", max_retry_attempts=DistributedMapProcessor.UNLIMITED, max_retry_duration=Duration.from_hours(1), @@ -555,7 +632,7 @@ def test_processor_unlimited_retries_maps_to_negative_one(): config=DistributedMapConfig(), ) processor = options["Processor"] - assert processor["FunctionResponseTypes"] == ["REPORT_BATCH_ITEM_RESULTS"] + assert processor["FunctionResponseTypes"] == ["REPORT_BATCH_ITEM_OUTCOMES"] assert processor["MaxRetryAttempts"] == -1 assert processor["MaxRetryDurationSeconds"] == 3600 @@ -565,7 +642,9 @@ def test_processor_explicit_retry_attempts_pass_through(): options = _start_options( _new_op_state_calls(), source=["a"], - processor=DistributedMapProcessor.batch("proc", max_retry_attempts=0), + processor=DistributedMapProcessor.report_batch_outcome( + "proc", max_retry_attempts=0 + ), max_concurrency=1, config=DistributedMapConfig(), ) @@ -581,16 +660,20 @@ def test_s3_source_serializes_config(): source=DistributedMapSource.s3_json_lines( "s3://bucket/data.jsonl", max_items=500 ), - processor=DistributedMapProcessor.batch("proc"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=2, config=DistributedMapConfig(), ) source = options["Source"] assert source["Type"] == "S3" assert source["MaxItemsToRead"] == 500 - assert source["S3SourceConfig"]["Bucket"] == "bucket" - assert source["S3SourceConfig"]["Key"] == "data.jsonl" - assert source["S3SourceConfig"]["Format"] == "JSON_LINES" + 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(): @@ -608,7 +691,7 @@ def test_full_config_serializes_all_blocks(): options = _start_options( _new_op_state_calls(), source=["a"], - processor=DistributedMapProcessor.item_results("proc"), + processor=DistributedMapProcessor.report_item_outcomes("proc"), max_concurrency=8, config=config, ) @@ -639,7 +722,7 @@ def test_result_config_returns_map_run_result_with_items(): name="test_map_run", status=OperationStatus.SUCCEEDED, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + source_exhausted=True, success_count=1, failure_count=1, unprocessed_count=0, @@ -660,7 +743,7 @@ def test_result_config_returns_map_run_result_with_items(): result = distributed_map_handler( source=["a", "b"], - processor=DistributedMapProcessor.item_results("proc"), + processor=DistributedMapProcessor.report_item_outcomes("proc"), max_concurrency=2, state=mock_state, operation_identifier=_identifier("mrr"), @@ -672,8 +755,110 @@ def test_result_config_returns_map_run_result_with_items(): assert result.get_results() == [42] errors = result.get_errors() assert len(errors) == 1 - assert errors[0].error_type == "ItemError" - assert errors[0].error_message == "boom" + 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(): @@ -687,7 +872,7 @@ def test_plain_config_returns_plain_summary(): name="test_map_run", status=OperationStatus.SUCCEEDED, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + source_exhausted=True, success_count=2, ), ) @@ -697,7 +882,7 @@ def test_plain_config_returns_plain_summary(): result = distributed_map_handler( source=["a", "b"], - processor=DistributedMapProcessor.batch("proc"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=2, state=mock_state, operation_identifier=_identifier("mrs"), @@ -714,11 +899,11 @@ def test_csv_source_header_location(): given = _start_options( _new_op_state_calls(), source=DistributedMapSource.s3_csv("s3://b/data.csv", headers=["a", "b"]), - processor=DistributedMapProcessor.batch("proc"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=1, config=DistributedMapConfig(), ) - csv_opts = given["Source"]["S3SourceConfig"]["CsvFormatOptions"] + csv_opts = given["Source"]["S3SourceConfig"]["ObjectConfig"]["CsvFormatOptions"] assert csv_opts["HeaderLocation"] == "GIVEN" assert csv_opts["Headers"] == ["a", "b"] assert csv_opts["Delimiter"] == "COMMA" @@ -727,13 +912,13 @@ def test_csv_source_header_location(): first_row = _start_options( _new_op_state_calls(), source=DistributedMapSource.s3_csv("s3://b/data.csv"), - processor=DistributedMapProcessor.batch("proc"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=1, config=DistributedMapConfig(), ) - s3_cfg = first_row["Source"]["S3SourceConfig"] - assert s3_cfg["CsvFormatOptions"]["HeaderLocation"] == "FIRST_ROW" - assert "Headers" not in s3_cfg["CsvFormatOptions"] + 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 @@ -746,11 +931,16 @@ def test_csv_delimiter_accepts_enum(): source=DistributedMapSource.s3_csv( "s3://b/data.csv", delimiter=DistributedMapCsvDelimiter.PIPE ), - processor=DistributedMapProcessor.batch("proc"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=1, config=DistributedMapConfig(), ) - assert opts["Source"]["S3SourceConfig"]["CsvFormatOptions"]["Delimiter"] == "PIPE" + assert ( + opts["Source"]["S3SourceConfig"]["ObjectConfig"]["CsvFormatOptions"][ + "Delimiter" + ] + == "PIPE" + ) def test_inline_source_over_1mb_rejected(): @@ -759,7 +949,7 @@ def test_inline_source_over_1mb_rejected(): _start_options( _new_op_state_calls(), source=big, - processor=DistributedMapProcessor.batch("proc"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=1, config=DistributedMapConfig(), ) @@ -772,7 +962,7 @@ def test_reader_state_serialized_and_capped(): source=DistributedMapSource.reader_function( "reader", initial_state={"page": 0} ), - processor=DistributedMapProcessor.batch("proc"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=1, config=DistributedMapConfig(), ) @@ -787,7 +977,7 @@ def test_reader_state_serialized_and_capped(): source=DistributedMapSource.reader_function( "reader", initial_state="x" * 40_000 ), - processor=DistributedMapProcessor.batch("proc"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=1, config=DistributedMapConfig(), ) @@ -807,7 +997,7 @@ def deserialize(self, data, _serdes_context): # noqa: ANN001, ANN201 options = _start_options( _new_op_state_calls(), source=DistributedMapSource.inline(["a", "b"], serdes=_UpperSerDes()), - processor=DistributedMapProcessor.batch("proc"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=1, config=DistributedMapConfig(), ) @@ -822,7 +1012,7 @@ def test_map_run_options_from_dict_round_trip(): sent = _start_options( _new_op_state_calls(), source=["a", "b"], - processor=DistributedMapProcessor.batch("proc"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=7, config=DistributedMapConfig(), ) @@ -836,7 +1026,7 @@ def test_destination_only_success(): opts = _start_options( _new_op_state_calls(), source=["a"], - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig( destination=DistributedMapDestination( @@ -853,7 +1043,7 @@ def test_destination_only_failure(): opts = _start_options( _new_op_state_calls(), source=["a"], - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig( destination=DistributedMapDestination( @@ -870,27 +1060,29 @@ def test_s3_objects_source_config(): opts = _start_options( _new_op_state_calls(), source=DistributedMapSource.s3_objects("s3://b/prefix/"), - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig(), ) s3 = opts["Source"]["S3SourceConfig"] - assert s3["Transform"] == "NONE" - assert s3["KeyPrefix"] == "prefix/" - assert "Format" not in s3 + 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.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig(), ) s3 = opts["Source"]["S3SourceConfig"] - assert s3["Transform"] == "LOAD_AND_FLATTEN" - assert s3["Format"] == "JSON_LINES" + assert s3["Type"] == "LISTED_OBJECTS" + assert s3["ListedObjectsConfig"]["KeyPrefix"] == "prefix/" + assert s3["ListedObjectsConfig"]["Format"] == "JSON_LINES" # Validations, config branches, and round-trips @@ -915,15 +1107,22 @@ def _start_executor(source, processor, config, max_concurrency=1): # --- Completion config translation --- -def test_empty_completion_config_omitted(): +def test_completion_config_percentage_and_sample_size_sent(): options = _start_options( _new_op_state_calls(), source=["a"], - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, - config=DistributedMapConfig(completion_config=DistributedMapCompletionConfig()), + config=DistributedMapConfig( + completion_config=DistributedMapCompletionConfig.failure_percentage( + 25, minimum_sample_size=50 + ) + ), ) - assert "CompletionConfig" not in options + assert options["CompletionConfig"] == { + "ToleratedFailurePercentage": 25, + "MinimumSampleSize": 50, + } # --- Source translation variants --- @@ -933,14 +1132,14 @@ def test_objects_whole_bucket_prefix_config(): options = _start_options( _new_op_state_calls(), source=DistributedMapSource.s3_objects("s3://bucket"), - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig(), ) s3 = options["Source"]["S3SourceConfig"] - assert s3["KeyPrefix"] == "" - assert s3["Transform"] == "NONE" - assert "Key" not in s3 + assert s3["Type"] == "LISTED_OBJECT_METADATA" + assert s3["ListedObjectMetadataConfig"]["KeyPrefix"] == "" + assert "ObjectConfig" not in s3 def test_flattened_csv_source_config(): @@ -949,22 +1148,21 @@ def test_flattened_csv_source_config(): source=DistributedMapSource.s3_flattened_csv( "s3://b/prefix", headers=["a", "b"] ), - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig(), ) - s3 = options["Source"]["S3SourceConfig"] - assert s3["Transform"] == "LOAD_AND_FLATTEN" - assert s3["Format"] == "CSV" - assert s3["CsvFormatOptions"]["HeaderLocation"] == "GIVEN" - assert s3["CsvFormatOptions"]["Headers"] == ["a", "b"] + 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.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig(), ) @@ -987,7 +1185,7 @@ def deserialize(self, data, _serdes_context): # noqa: ANN001, ANN201, ARG002 options = _start_options( _new_op_state_calls(), source=DistributedMapSource.inline([1, 2], serdes=_PlainTextSerDes()), - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig(), ) @@ -1001,7 +1199,7 @@ def test_success_destination_include_input_and_owner(): options = _start_options( _new_op_state_calls(), source=["a"], - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig( destination=DistributedMapDestination( @@ -1024,7 +1222,7 @@ def test_failure_destination_error_only_and_owner(): options = _start_options( _new_op_state_calls(), source=["a"], - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig( destination=DistributedMapDestination( @@ -1053,13 +1251,18 @@ def test_operation_to_dict_round_trip_preserves_results(): name="test_map_run", status=OperationStatus.SUCCEEDED, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + 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", - completion_details="done", + error=ErrorObject( + message="tolerated", + type="DistributedMapFailure", + data=None, + stack_trace=None, + ), results=( DistributedMapResultItemApi( item_id="0", status=DistributedMapItemStatus.SUCCEEDED, output=5 @@ -1077,11 +1280,14 @@ def test_operation_to_dict_round_trip_preserves_results(): 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["CompletionDetails"] == "done" + 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" @@ -1090,7 +1296,7 @@ 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.item_results( + processor=DistributedMapProcessor.report_item_outcomes( "proc", batch_size=5, max_retry_attempts=DistributedMapProcessor.UNLIMITED, @@ -1111,7 +1317,7 @@ def test_options_full_round_trip(): assert parsed.max_concurrency == 3 assert parsed.source.source_type is DistributedMapSourceType.S3 assert parsed.processor.function_response_types == [ - DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_RESULTS + DistributedMapFunctionResponseType.REPORT_BATCH_ITEM_OUTCOMES ] assert parsed.processor.max_retry_attempts == -1 assert parsed.processor.durable_execution_name_prefix == "pfx" @@ -1130,11 +1336,11 @@ def test_csv_first_row_no_headers(): options = _start_options( _new_op_state_calls(), source=DistributedMapSource.s3_csv("s3://b/f.csv"), - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig(), ) - csv_opts = options["Source"]["S3SourceConfig"]["CsvFormatOptions"] + csv_opts = options["Source"]["S3SourceConfig"]["ObjectConfig"]["CsvFormatOptions"] assert csv_opts["HeaderLocation"] == "FIRST_ROW" assert "Headers" not in csv_opts assert csv_opts["Delimiter"] == "COMMA" @@ -1146,7 +1352,7 @@ def test_s3_source_expected_bucket_owner(): source=DistributedMapSource.s3_json_lines( "s3://b/k.jsonl", expected_bucket_owner="123456789012" ), - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig(), ) @@ -1157,7 +1363,7 @@ def test_empty_destination_config_omitted(): options = _start_options( _new_op_state_calls(), source=["a"], - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig(destination=DistributedMapDestination()), ) @@ -1168,7 +1374,7 @@ def test_success_destination_input_only(): options = _start_options( _new_op_state_calls(), source=["a"], - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig( destination=DistributedMapDestination( @@ -1185,7 +1391,7 @@ def test_failure_destination_input_only(): options = _start_options( _new_op_state_calls(), source=["a"], - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig( destination=DistributedMapDestination( @@ -1204,7 +1410,7 @@ def test_reader_source_options_round_trip(): source=DistributedMapSource.reader_function( "reader", initial_state={"page": 0} ), - processor=DistributedMapProcessor.batch("p"), + processor=DistributedMapProcessor.report_batch_outcome("p"), max_concurrency=1, config=DistributedMapConfig(), ) @@ -1222,7 +1428,7 @@ def test_operation_to_dict_minimal_details_omits_optionals(): name="test_map_run", status=OperationStatus.SUCCEEDED, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + source_exhausted=True, success_count=1, failure_count=0, unprocessed_count=0, @@ -1282,7 +1488,7 @@ def test_terminal_failure_with_details_resolves_with_summary(status, expected): name="test_map_run", status=status, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.FAILURE_TOLERANCE_EXCEEDED, + source_exhausted=False, success_count=1, failure_count=2, unprocessed_count=0, @@ -1305,8 +1511,8 @@ def test_terminal_failure_with_details_resolves_with_summary(status, expected): summary.throw_if_error() -def test_terminal_without_completion_reason_raises(): - """A terminal operation missing its completion reason is surfaced, not papered over.""" +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( @@ -1316,6 +1522,7 @@ def test_terminal_without_completion_reason_raises(): 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, @@ -1324,7 +1531,7 @@ def test_terminal_without_completion_reason_raises(): mock_state.get_checkpoint_result.return_value = ( CheckpointedResult.create_from_operation(operation) ) - with pytest.raises(ExecutionError, match="carried no CompletionReason"): + with pytest.raises(ExecutionError, match="carried no SourceExhausted"): distributed_map_handler( source=["a"], processor="p", @@ -1358,7 +1565,7 @@ def deserialize(self, data, _serdes_context): # noqa: ANN001, ANN201 name="test_map_run", status=OperationStatus.SUCCEEDED, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + source_exhausted=True, success_count=1, failure_count=0, unprocessed_count=0, @@ -1376,7 +1583,7 @@ def deserialize(self, data, _serdes_context): # noqa: ANN001, ANN201 ) result = distributed_map_handler( source=["a"], - processor=DistributedMapProcessor.item_results("proc"), + processor=DistributedMapProcessor.report_item_outcomes("proc"), max_concurrency=1, state=mock_state, operation_identifier=_identifier("cs"), @@ -1404,7 +1611,7 @@ def test_falsy_output_round_trips(value): name="test_map_run", status=OperationStatus.SUCCEEDED, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + source_exhausted=True, success_count=1, failure_count=0, unprocessed_count=0, @@ -1416,7 +1623,7 @@ def test_falsy_output_round_trips(value): ) result = distributed_map_handler( source=["a"], - processor=DistributedMapProcessor.item_results("proc"), + processor=DistributedMapProcessor.report_item_outcomes("proc"), max_concurrency=1, state=mock_state, operation_identifier=_identifier("fo"), @@ -1430,7 +1637,7 @@ def test_processor_retry_duration_only(): options = _start_options( _new_op_state_calls(), source=["a"], - processor=DistributedMapProcessor.batch( + processor=DistributedMapProcessor.report_batch_outcome( "proc", max_retry_duration=Duration.from_minutes(5), ), @@ -1451,7 +1658,7 @@ def test_resolved_summary_uses_operation_status_not_details(): name="test_map_run", status=OperationStatus.SUCCEEDED, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + source_exhausted=True, success_count=2, failure_count=0, unprocessed_count=0, @@ -1460,7 +1667,7 @@ def test_resolved_summary_uses_operation_status_not_details(): ) executor = DistributedMapOperationExecutor( source=["a", "b"], - processor=DistributedMapProcessor.batch("proc"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=1, state=Mock(spec=ExecutionState), operation_identifier=_identifier("dmap-id"), @@ -1470,7 +1677,7 @@ def test_resolved_summary_uses_operation_status_not_details(): summary = executor._resolve_summary(operation) assert summary.status is DistributedMapStatus.SUCCEEDED - assert summary.completion_reason is DistributedMapCompletionReason.ALL_COMPLETED + assert summary.source_exhausted is True assert summary.success_count == 2 @@ -1485,7 +1692,7 @@ def test_checkpoint_from_a_different_map_is_rejected(): name="a_different_map", status=OperationStatus.SUCCEEDED, distributed_map_details=DistributedMapDetails( - completion_reason=DistributedMapCompletionReason.ALL_COMPLETED, + source_exhausted=True, success_count=999, failure_count=0, unprocessed_count=0, @@ -1499,7 +1706,7 @@ def test_checkpoint_from_a_different_map_is_rejected(): with pytest.raises(NonDeterministicExecutionError, match="name checkpoint="): distributed_map_handler( source=["a", "b"], - processor=DistributedMapProcessor.batch("proc"), + processor=DistributedMapProcessor.report_batch_outcome("proc"), max_concurrency=1, state=mock_state, operation_identifier=_identifier("dmap-id"),