[nvasiu]

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]

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]

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

[nvasiu]

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

[yaythomas]

yes, sorry, I didn't phrase that well. Wondered if there's some sort of preview check, we could usefully do while in pre-release.

[nvasiu]

Ah yeah, I couldn't find anything that we could use in the SDK to know whether an API is unreleased/preview or not.

[yaythomas]

Could we add an end-to-end test that begins with an item-handler error carrying errorData and stackTrace, then verifies that both fields are present in the parent's DistributedMapResult? The current tests verify the processor response, but not the complete round trip.

[nvasiu]

Added the test, but it shows the opposite behaviour. We send all 4 members from the processor, but when we get the error back from the service (in DistributedMapResult), data and stack_trace come back as None.

[yaythomas]

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

[nvasiu]

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.

[yaythomas]

FileSystemSerDes keys a file by the serialization context's execution ARN and operation ID. This loop reuses that key for every inline item. The processor and reader wrappers also reuse an empty context for every output, item, and state. I reproduced the result: different values returned the same file pointer, which then read the last value written. Separate handler invocations also used the same path.

Can each serialized value keep an immutable, deterministic reference, with matching context at each producer and consumer boundary? Can tests cover inline items, processor outputs, reader items and state, replay, and concurrent calls with FileSystemSerDes?

[nvasiu]

Added a new DistributedMapRunArn.serdes_context(role, *parts) which hashes [run id / role / parts] into the context's operation ID (existing convention in the SDK with create_id_for_step), and fills in the parent execution ARN. Used across each producer/consumer, so they always derive the same context. And added all the related test cases.

[nvasiu]

Need to make a clarification here.

Only the processor output path derives the same context on both sides. The new serdes_context() still fixes the above defect, every value gets its own address with no key collisions. But there are 3 input paths that read under a different context than the one they were written with:

• inline items are serialized before the run exists, so there is no run id or item id to derive an address from yet
• reader items are addressed by the reader with its cursor and the item's index, and by the processor with the service's item id
• the reader's next_state is written as one call's output and read as the next call's input

This doesn't affect any serdes in the SDK, none of them use the context to locate a value on read. E.g. FileSystemSerDes puts the file path in the payload. A custom serdes that does use the read context will not find the value, so it raises instead of silently returning the wrong one.

ctx.invoke has a similar issue: one Lambda function serializes the payload and a different Lambda function deserializes the payload. The SDK runs only the serializing side.

This wasn't documented before, so I will add the constraint into the serdes.py docstring: a custom SerDes used with ctx.distributed_map or ctx.invoke must write a payload it can read back without the context it was given.

[nvasiu]

Summary of what changed since the last push.

Addressing review comments

• operation/dmap.py:118: added DistributedMapRunArn.serdes_context(role, *parts) and used it for every value the parent, processor and reader serialize, so a pointer-style serdes no longer writes several values to one key.
  • Documented the serdes context requirement in the serdes.py module docstring. Only the processor output path derives the same context on both sides, so a custom serdes must write a payload it can read back without the context it was given. That constraint was not written down anywhere.
• lambda_service.py:505: the two distributed map parsers test for error presence instead of truthiness, so an empty error object survives a parse and serialize round trip.
• dmap/models.py:186: throw_if_error raises from the rebuilt operation error, so the run-level branch keeps the reconstructed cause chain that the item branch already had, and attaches no cause when the error object is empty.
• exceptions.py:333: added DistributedMapError to _DURABLE_OPERATION_ERROR_REGISTRY, so it reconstructs as itself rather than as the DurableOperationError base.
• plugin.py:50: _extract_error reads distributed_map_details, so a failed map run reaches the end and change hooks instead of arriving with error=None.
• lambda_service.py:1644: the start factory copies the subtype from the identifier and the context passes a plain string, so every hook sees a str rather than an enum on the first run and a str on replay. The string sent is unchanged.
• config.py:779: added _validate_not_bool to every numeric distributed map field, so True and False are rejected instead of being sent as JSON booleans where the API requires a number.
• context.py:732: source is back to DistributedMapSource | Sequence[Any], consistent with .inline().
• operation/dmap.py:171: the reader path uses a narrowing cast and the unreachable error branch is deleted.
• dmap/handlers.py:271: both item wrappers' docstrings state that func, item_serdes and result_serdes are called from several threads when concurrency is above 1.
• dmap/handlers.py:512: added tests/e2e/dmap_durable_item_replay_test.py, covering a wait inside an item, a resume with a custom result_serdes, and a completed sibling not running again.
• config.py:890 and config.py:996: added the five S3Source branch tests plus one parameterized test over every named S3 factory, so those invariants and mappings cannot regress silently.
• dmap/handlers.py:138: added the end-to-end test from an item handler error through to the parent's DistributedMapResult. It pins current behaviour, which is the opposite of what the test was expected to show: the processor sends the error type, message, errorData and stackTrace, and only the type and message reach the parent.

Run arn and item id reach the processor and reader functions

Updated the helper wrappers to pass the map run ARN and the item IDs to the customer's functions, so customers can use both for their own logging. Before this change the wrappers passed only the arguments a customer's function needs to do the work. I'll make the same update in the other SDKs.

• types.py: added DistributedMapItemContext, DistributedMapBatchContext and DistributedMapReaderContext, all exported from the package root.
• Each of the five wrappers passes the matching context as the last argument to the customer's function, with a logger that stamps the run arn and, for an item, the item id.
• ProcessorEvent and ReaderEvent require and parse distributedMapRunArn, raising ValidationError when it is absent or invalid.

Other

• DistributedMapResult.from_operation: new override that raises ValidationError naming from_operation_and_items. The inherited factory returned an empty all while the operation retained every item result.
• distributed_map_reader: comment on the repeated-state check, which cannot detect a repeat when the serdes returns a pointer.

Tests

• tests/dmap/handlers_test.py, tests/operation/dmap_test.py, tests/e2e/dmap_helpers_int_test.py: per-value keys and contexts across the parent, processor and reader paths, including concurrency, retries and a second run.
• tests/e2e/dmap_int_test.py: FileSystemSerDes carrying every item from parent to processor and back, and the subtype arriving as a plain str on first run and replay.
• tests/config_test.py, tests/context_test.py: the S3 factories and their rejection messages, bool on every numeric field, and any Sequence as a source.
• Mechanical edits elsewhere: the trailing context parameter, distributedMapRunArn in every envelope, a run arn in the parent fixtures, and distributed_map_details on the PluginExecutor mock.