feat(dmap): add the distributed map operation by nvasiu · Pull Request #849 · aws/aws-durable-execution-sdk-python

nvasiu · 2026-10-26T21:20:53Z

Issue #, if available

N/A

Description of changes

Adds the distributed map operation (ctx.distributed_map) to the Python SDK:
A map run processes a bounded dataset in parallel. A customer starts a map run from a durable function, naming a source to read items from, a processor function to invoke per batch, and concurrency, retry, and failure settings. The service reads items from the source, groups them into batches, invokes the processor for each batch, retries failures, tracks progress, routes successful results and failed items to destinations, and reports completion.

config.py

The input types a customer constructs to describe a map run, and the distributed map enums.

DistributedMapConfig: optional settings for distributed map.
DistributedMapResultConfig: subclass that additionally collects item results inline, selecting the return type statically.
DistributedMapSource: factories naming where map run items come from. inline, s3_json_lines, s3_json_array, s3_csv, s3_objects, s3_flattened_json_lines, s3_flattened_json_array, s3_flattened_csv and reader_function. S3Source and ReaderSource are the resolved forms a factory returns.
DistributedMapProcessor: describes the Lambda that processes items. The report_batch_outcome and report_item_outcomes factories pick how the processor reports outcomes back. The retry option takes a DistributedMapProcessorRetryConfig holding max_retry_attempts and max_retry_duration, with DistributedMapProcessorRetryConfig.UNLIMITED for an unbounded attempt count.
DistributedMapCompletionConfig: defines item failure thresholds for marking the overall map run failed.
DistributedMapDestination: the s3_successes and s3_failures factories for routing item records to S3. SuccessDestination and FailureDestination are the resolved forms, and DistributedMapDestinationConfig holds the pair.
Each config validates its own arguments on construction: a destination prefix of at most 900 UTF-8 bytes, a batch size between 1 and 10000, a timeout of at most 90 days, and retry attempts that are non-negative or UNLIMITED.
The enums live here rather than in lambda_service, so config.py has no runtime import from it.

context.py

The customer-facing entry point on the durable execution context.

ctx.distributed_map: the method a customer calls to run a distributed map. Overloaded so passing a DistributedMapResultConfig types the return as DistributedMapResult and anything else as DistributedMapSummary.
Requires max_concurrency to be greater than zero.

dmap/models.py

The result types a customer receives back from a map run.

DistributedMapSummary: what ctx.distributed_map returns, describes the run's overall outcome.
DistributedMapResult: returned when a DistributedMapResultConfig is passed. Contains individual map run item outcomes.
DistributedMapResultItem: represents a single item's result or error.
DistributedMapRunArn: the parsed run ARN, exposing the run id.
The item counts and each collected result's item id are required. A terminal operation that omits a count, or a result that omits its item id, raises an error naming the missing field rather than substituting a value the service never sent.

dmap/handlers.py

Authoring decorators for the processor Lambda, so a customer writes a plain function rather than the item or batch protocol.

distributed_map_item_handler, distributed_map_batch_handler, distributed_map_reader, and the durable variants durable_distributed_map_item_handler and durable_distributed_map_batch_handler.
Named and shaped like durable_execution, so they work bare, with options, or called directly with a function.
The durable variants take boto3_client and plugins and pass them through, so a durable processor can register instrumentation plugins or supply its own client.
Responses are built from envelope dataclasses rather than hand-built dicts.
The wrappers validate the invocation input and reject a record that carries no item id, since every outcome is reported against an item id.

operation/dmap.py

The executor that drives the operation against the durable execution runtime.

Suspends the caller's function while the run executes and resumes it with the finished outcome.
Every terminal state resolves, and throw_if_error() opts into raising.
All translation between the config types and the service shapes happens here.

lambda_service.py

Serialization for carrying the operation and its results to and from the backend service.

One dataclass per API shape, named after the shape rather than with a Wire suffix.
Items, Output, and the record body are opaque strings, matching the API model, so any serdes works rather than JSON only.
The operation parser reads DistributedMapDetails whenever the key is present. The service sends the object on every distributed map operation and leaves it empty until the run starts, so an empty object round trips as an empty object.

state.py

Durable execution state handling, so a run's outcome persists across suspend and resume.

Records that a distributed map operation carries its outcome on the operation itself rather than as an operation-level result or error.

exceptions.py

The error type a customer catches.

DistributedMapError: raised when a run or an item fails.

plugin.py

Operation type registration.

Adds DISTRIBUTED_MAP to OperationType.

__init__.py

The package's public API surface.

Exports the result types, their enums, and the five authoring decorators, matching how durable_execution and durable_step are exported.
The input config and factory types are imported from config.

dmap/__init__.py

Empty package marker for the dmap package.

Tests

Tests are split by layer so each sits alongside its module.

tests/config_test.py: config and argument validation.
tests/lambda_service_test.py and tests/dmap/models_test.py: serialization round trips and the result types.
tests/operation/dmap_test.py and tests/context_test.py: executor behaviour and the ctx.distributed_map surface.
tests/dmap/handlers_test.py: authoring wrapper tests, checking that they process items, report failures, and reject bad inputs.
tests/e2e/dmap_int_test.py and tests/e2e/dmap_helpers_int_test.py: end to end tests mocking backend responses, covering suspend and resume, result collection, and throwing on failure.

Future tasks

Add distributed map to the local emulator (in the testing package). When this is done, we can add full end to end tests using the emulator.
Add distributed map examples to the examples package.

Design

See Rules for every change.

Design note (rule 3), for a new feature, a new option or a public API change: Feature planned internally.
Simpler options considered (rule 2), if the change adds logic: N/A
Net change in source lines and in public API members:
+2916 lines in src and +3952 in tests.
New public members: ctx.distributed_map, 11 names added to the package __all__, 19 types and 16 factory methods in config.py, 4 result types in dmap/models.py, 5 authoring decorators, and 1 exception type. 24 dataclasses in lambda_service.py carry the API shapes.

Assumptions

Each assumption about the service, Lambda, the runtime or another library, with its evidence, or marked "unverified" (rule 6):

A distributed map operation ends in one of SUCCEEDED, FAILED, STOPPED or TIMED_OUT. Unverified.
Every distributed map operation carries DistributedMapDetails, empty until the run starts. Unverified.
A terminal operation's DistributedMapDetails carries SuccessCount, FailureCount and UnprocessedCount, carries TotalCount once the run succeeded, carries SourceExhausted once DistributedMapRunArn is present, and reports a run failure as Error there rather than as an operation-level error. Unverified.
A run ARN is the parent durable execution ARN, then /distributed-map-run/, then a run id containing no further /. Unverified.
Every collected result in Results carries ItemId and a Status of SUCCEEDED or FAILED, with the item's value in Output as a string the customer's serdes decodes. Unverified.
A processor invocation payload carries a records list, each record holding a non-empty string itemId and a body the customer's serdes decodes. Unverified.
The service matches a reported outcome back to its item by the itemId the processor echoes, and accepts the batch-failure and item-outcome response shapes the wrappers return. Unverified.
A reader invocation payload carries an integer maxItemsPerPage, and state as the reader last returned it. Unverified.
The service accepts the source, processor, completion and destination configuration the checkpoint sends, with MaxRetryAttempts of -1 meaning unlimited attempts and MaxRetryDurationSeconds a whole number of seconds. Unverified.
The service's limits are the ones the configs validate against: a destination key prefix of 900 characters and 900 UTF-8 bytes, a batch size of 10000 items, and a map run timeout of 90 days. The ceiling on concurrent processor invocations is the service's alone, since the SDK enforces none. Unverified.

Decisions not covered by the issue or the design note

One line each, with the reason (rule 7):

A missing count or item id raises ExecutionError, because the other checks in the same parse path raise it.
DistributedMapProcessorRetryConfig.UNLIMITED is the string "unlimited" rather than the wire sentinel, so a caller never writes a magic number.
The source factories are flat on DistributedMapSource, named s3_json_lines and so on, so the factory name states the source kind in one call.
DistributedMapDestination carries the factories and DistributedMapDestinationConfig holds on_success and on_failure, because a destination factory returns one half of that pair rather than the pair itself.
The destination prefix check is 900 UTF-8 bytes, which is the lower of the service's character and byte limits, so a multi-byte prefix is rejected at the real limit.
max_concurrency is required and checked only as greater than zero. No upper bound is enforced in the SDK, so a service-side limit change needs no SDK release.
The validation guards raise rather than assert, because python -O removes asserts and the check would then be absent from the build that most wants it.
A processor record with no item id fails the invocation before any item handler runs, rather than failing that record alone, because an outcome cannot be reported for an item that was not named.
The durable wrappers take boto3_client and plugins rather than one config object, because durable_execution itself takes those two keyword arguments.
Items, Output and the record body are opaque strings, so a customer serdes other than JSON works without an SDK change.

Checklist

Every new branch has a test that fails without it (rule 4)
For a bug fix: the first commit adds a test that fails on the current code (rule 5)
Every edited test expectation is listed in the description, with the reason (rule 5)

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

nvasiu · 2026-10-26T22:01:11Z

Summary of what changed since the last update on the previous PR.

Changes suggested by @yaythomas

DistributedMapS3SourceConfig.from_source removed all fallbacks because they invented a default value to satisfy mypy (and 3/5 were unreachable). Replaced with checks that raise and name what is missing. These use if ... raise instead of assert because python -O removes asserts.
report_batch_outcome() / report_item_outcomes(). Names unchanged. Their retry option now nested: one retry argument of type DistributedMapProcessorRetryConfig, holding max_retry_attempts and max_retry_duration, with UNLIMITED moved onto it.
ProcessorResponseMode export. Dropped. Only report_batch_outcome and report_item_outcomes set it, and nothing hands it back to a customer, so a customer never needs to name it. The other distributed map config enums were already unexported. It stays importable from config.

Processor retry settings

DistributedMapProcessorRetryConfig is a new frozen dataclass holding max_retry_attempts and max_retry_duration. DistributedMapProcessor and both its factories take one retry argument of that type in place of the two flat options. The UNLIMITED sentinel moved onto it.
UNLIMITED is declared ClassVar[Literal["unlimited"]] rather than ClassVar[str]. As a str it failed type checking when passed to max_retry_attempts, so a caller using the SDK's own constant got a mypy error. The test that passes the constant now carries a return annotation, since mypy skips an unannotated function body and that is why the suite never caught it.

Destination factories

DistributedMapDestination.s3_successes and DistributedMapDestination.s3_failures replace SuccessDestination.from_uri and FailureDestination.from_uri, which are deleted.
The class holding on_success and on_failure is renamed DistributedMapDestinationConfig, which frees the name DistributedMapDestination for the factories. DistributedMapConfig.destination takes the renamed type.
A destination key prefix over 900 UTF-8 bytes is rejected when the destination is constructed.

Fields the service is required to send

SuccessCount, FailureCount and UnprocessedCount are now int | None, defaulting to None. A terminal operation that omits one raises an error naming that count. They previously defaulted to 0, which cannot be told apart from a count the service sent as 0. TotalCount is required once the status is SUCCEEDED, since a run that failed, timed out or was stopped may never have counted its source.
A collected result with no ItemId raises, naming the missing field. An empty item id gives a result the caller cannot match to an item.
A processor record with no itemId is rejected. The wrappers copy that id onto every outcome they report, so an empty id reports the outcome against no item. The check runs over the whole batch before any item handler is called, so one unusable record fails the invocation rather than the batch running partly.

Durable processor wrappers

durable_distributed_map_item_handler and durable_distributed_map_batch_handler take boto3_client and plugins and pass them to durable_execution, so a durable processor can supply its own client or register instrumentation plugins.

Operation parsing

The operation parser reads DistributedMapDetails when the key is present rather than when its value is truthy. The service sends the object on every distributed map operation and leaves it empty until the run starts, and an empty object was previously discarded, so it did not survive a parse and serialize round trip.

And resolved some conflicts for the rebase.

yaythomas · 2026-10-28T17:34:54Z

+                f"{self.tolerated_failure_count}"
+            )
+            raise ValidationError(msg)
+        if self.tolerated_failure_percentage is not None and not isinstance(


Shouldn't this support fractions also? Can the public and serialization types accept float, while rejecting NaN and values outside 0 through 100?

Updated tolerated_failure_percentage and the failure_percentage factory to take int | float. I included both because existing code does that, CompletionConfig declares tolerated_failure_percentage: int | float | None. A bare float would accept an int argument too, but it would report the attribute as float | None when an int was stored.

The current check already rejects NaN and values outside 0-100, I added tests for those cases.

yaythomas · 2026-10-28T17:34:54Z

+    return a ``DistributedMapResult``, which carries the per-item outcomes.
+    """
+
+    result_serdes: SerDes | None = None  # None = DEFAULT_JSON_SERDES


The result type from result_serdes is lost here. A caller supplying a SerDes[Order] gets list[Any] from DistributedMapResult.get_results().

Can DistributedMapResultConfig, DistributedMapResult, and DistributedMapResultItem carry the result type so get_results() returns list[R], like BatchResult[R]?

DistributedMapResultConfig, DistributedMapResult and DistributedMapResultItem carry the output type now, so a SerDes[Order] gives DistributedMapResult[Order] and get_results() gives list[Order | None] instead of list[Any]. The None is because a succeeded item can have no output.

This change followed ctx.invoke's example. ctx.map does not actually pass its result type through: its type variable is in no parameter, so ctx.map([1], f).get_results() is list[Never].

yaythomas · 2026-10-28T17:34:54Z

+    @overload
+    def distributed_map(
+        self,
+        source: DistributedMapSource | Sequence[Any],


The signature accepts any Sequence, but the runtime check accepts only DistributedMapSource, list, or tuple. Mypy therefore accepts values such as range(3) and str, which this method then rejects.

Can the annotation name the accepted types, or can the implementation honor the Sequence contract?

Narrowed source to DistributedMapSource | list[Any] | tuple[Any, ...], in both overloads, the implementation and DistributedMapOperationExecutor.

yaythomas · 2026-10-28T17:34:54Z

+        *,
+        config: DistributedMapConfig | None = None,
+    ) -> DistributedMapSummary:
+        """Start a distributed map run and resolve with its summary.


Since this API is prelaunch, shouldn't DurableContext.distributed_map include the preview check here?

Did you mean there is an existing preview check I could use here? Couldn't find any example in the SDK.

yaythomas · 2026-10-28T17:34:55Z

            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:


Note: pending model updates in botocore for DistributedMapOptions, DistributedMapDetails, OperationType.DISTRIBUTED_MAP, and the related shapes. Nothing to change in this PR.

yaythomas · 2026-10-28T17:34:55Z

+    ) -> 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"))


Should this calculation use ensure_ascii=False? With a custom serializer that returns non-ASCII text, the default json.dumps behavior counts each character as a \uXXXX escape.

For example, 180,000 é characters count as 1,080,004 bytes here, while their UTF-8 JSON representation is 360,004 bytes. Could we verify which representation the limit applies to and match that representation exactly?

Updated to use ensure_ascii=False.

The limit is 1 MB on the serialized item collection, measured as the UTF-8 bytes of its compact JSON array, with non-ASCII characters unescaped and JSON's own quote and backslash escaping counted. Confirmed, and the check now matches it exactly.

yaythomas · 2026-10-28T17:34:55Z

+                self.state.create_checkpoint(
+                    operation_update=start_operation, is_sync=True
+                )
+            except CheckpointError as e:


ExecutionState.create_checkpoint(is_sync=True) reports checkpoint-worker failures as BackgroundThreadError, with the original CheckpointError as its source. It does not raise CheckpointError directly here.

The current tests pass because their mock raises CheckpointError directly at this call site. Can this remove the unreachable translation, or test the intended behavior through a real ExecutionState checkpoint worker?

Removed, with _rejected_config_error. ExecutionState.create_checkpoint reports the failure as BackgroundThreadError, which inherits BaseException, so the except CheckpointError never matched. A rejected START now fails the execution like any other rejected checkpoint, as ctx.invoke does. Tested through a real ExecutionState and checkpoint worker.

Add ctx.distributed_map, which runs a processor function over a large item source and resolves with a summary when the run reaches a terminal state. Covers the configuration types, the handler decorators, the result types and the checkpoint translation.

nvasiu · 2026-10-29T00:33:13Z

Another update, addressing the above comments plus a few changes to match repo conventions.
Details in replies to comments above, other changes below:

Shared constants

constants.py: added READER_STATE_LIMIT_BYTES. The 32 KB limit was written out in both dmap/handlers.py and operation/dmap.py, so a change could miss one.
dmap/handlers.py, operation/dmap.py: import it, and the two _READER_STATE_LIMIT declarations are deleted. _INLINE_SIZE_LIMIT stays in operation/dmap.py, its only user.

Removed unused code

dmap/handlers.py: deleted ProcessorRecord.to_dict, ProcessorEvent.to_dict, ItemResult.from_dict, ItemHandlerResponse.from_dict, ReaderEvent.to_dict and ReaderResponse.from_dict. Each wrapper reads the event and writes the response, so the other direction is the service's side and had no caller. ItemFailure.from_dict stays, test_item_error_data_and_stack_trace_reach_the_parent_result uses it.

Docstrings

dmap/handlers.py: the module docstring said these wrappers are not imported from the package. __init__.py exports all five.

Tests

tests/e2e/dmap_helpers_int_test.py: durable_distributed_map_item_handler with a result type only result_serdes can write. Fails without the per_item change.
tests/dmap/typing_test.py: new, checks the five decorator signatures and DistributedMapResult's output type.
tests/config_test.py: NaN and infinity cases for tolerated_failure_percentage.
tests/dmap/handlers_test.py: handler calls pass a context.

nvasiu force-pushed the feat/distributed-map branch from c95ee1f to b12de90 Compare October 26, 2026 21:40

yaythomas requested changes Oct 28, 2026

View reviewed changes

nvasiu force-pushed the feat/distributed-map branch from b12de90 to 5cbea0b Compare October 29, 2026 00:32

Conversation

nvasiu commented Oct 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Issue #, if available

Description of changes

Tests

Future tasks

Design

Assumptions

Decisions not covered by the issue or the design note

Checklist

Uh oh!

nvasiu commented Oct 26, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

nvasiu commented Oct 29, 2026

Shared constants

Removed unused code

Docstrings

Tests

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

nvasiu commented Oct 26, 2026 •

edited

Loading