Skip to content

feat(dmap): add the distributed map operation - #849

Draft
nvasiu wants to merge 1 commit into
mainfrom
feat/distributed-map
Draft

nvasiu wants to merge 1 commit into
mainfrom
feat/distributed-map

Conversation

@nvasiu

@nvasiu nvasiu commented Oct 26, 2026 •

Copy link
Copy Markdown
Contributor

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.

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
nvasiu force-pushed the feat/distributed-map branch from c95ee1f to b12de90 Compare October 26, 2026 21:40
@nvasiu

nvasiu commented Oct 26, 2026

Copy link
Copy Markdown
Contributor Author

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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant