feat: Implement efficient replay mechanism

Parallel and Map now replays without using threading.

Changes:
 - context.map and context.parallel now call child_handler directly and
   create their own operation identifier
 - `map_handler` and `parallel_handler` now checks the state of the
   operation and when we've already succeeded, we use `executor.replay`

We test this by mocking the state, and ensuring we first run through
execute and then through replay and get back the correct results.

Author: FullyTyped

src/aws_durable_execution_sdk_python/concurrency.py

@@ -22,6 +22,7 @@
 from aws_durable_execution_sdk_python.identifier import OperationIdentifier
 from aws_durable_execution_sdk_python.lambda_service import ErrorObject
 from aws_durable_execution_sdk_python.operation.child import child_handler
+from aws_durable_execution_sdk_python.serdes import deserialize
 from aws_durable_execution_sdk_python.types import BatchResult as BatchResultProtocol
 
 if TYPE_CHECKING:
@@ -824,5 +825,46 @@ def run_in_child_handler():
             ),
         )
 
+    def replay(self, execution_state: ExecutionState, executor_context: DurableContext):
+        # if we are here, then we are in replay_children
+        # we will generate again all operation_ids for the children, and collect them
+        items: list[BatchItem[ResultType]] = []
+        for executable in self.executables:
+            operation_id = executor_context._create_step_id_for_logical_step(  # noqa: SLF001
+                executable.index
+            )
+            checkpoint = execution_state.get_checkpoint_result(operation_id)
+
+            result: ResultType | None = None
+            error = None
+            status: BatchItemStatus
+            if checkpoint.is_succeeded():
+                status = BatchItemStatus.SUCCEEDED
+                if checkpoint.is_replay_children():
+                    result = self._execute_item_in_child_context(
+                        executor_context, executable
+                    )
+                else:
+                    serialized_result = checkpoint.result
+                    if serialized_result is None:
+                        result = None
+                    else:
+                        result = deserialize(  # type: ignore[assignment]
+                            serdes=self.item_serdes or self.serdes,
+                            data=serialized_result,
+                            operation_id=operation_id,
+                            durable_execution_arn=execution_state.durable_execution_arn,
+                        )
+
+            elif checkpoint.is_failed():
+                error = checkpoint.error
+                status = BatchItemStatus.FAILED
+            else:
+                status = BatchItemStatus.STARTED
+
+            batch_item = BatchItem(executable.index, status, result=result, error=error)
+            items.append(batch_item)
+        return BatchResult.from_items(items, self.completion_config)
+
 
 # endregion concurrency logic

src/aws_durable_execution_sdk_python/context.py

@@ -317,7 +317,13 @@ def map(
         """Execute a callable for each item in parallel."""
         map_name: str | None = self._resolve_step_name(name, func)
 
-        def map_in_child_context(map_context) -> BatchResult[R]:
+        operation_id = self._create_step_id()
+        operation_identifier = OperationIdentifier(
+            operation_id=operation_id, parent_id=self._parent_id, name=map_name
+        )
+        map_context = self.create_child_context(parent_id=operation_id)
+
+        def map_in_child_context() -> BatchResult[R]:
             # map_context is a child_context of the context upon which `.map`
             # was called. We are calling it `map_context` to make it explicit
             # that any operations happening from hereon are done on the context
@@ -328,11 +334,13 @@ def map_in_child_context(map_context) -> BatchResult[R]:
                 config=config,
                 execution_state=self.state,
                 map_context=map_context,
+                operation_identifier=operation_identifier,
             )
 
-        return self.run_in_child_context(
+        return child_handler(
             func=map_in_child_context,
-            name=map_name,
+            state=self.state,
+            operation_identifier=operation_identifier,
             config=ChildConfig(
                 sub_type=OperationSubType.MAP,
                 serdes=config.serdes if config is not None else None,
@@ -346,8 +354,12 @@ def parallel(
         config: ParallelConfig | None = None,
     ) -> BatchResult[T]:
         """Execute multiple callables in parallel."""
+        # _create_step_id() is thread-safe. rest of method is safe, since using local copy of parent id
+        operation_id = self._create_step_id()
+        parallel_context = self.create_child_context(parent_id=operation_id)
+        operation_identifier: OperationIdentifier
 
-        def parallel_in_child_context(parallel_context) -> BatchResult[T]:
+        def parallel_in_child_context() -> BatchResult[T]:
             # parallel_context is a child_context of the context upon which `.map`
             # was called. We are calling it `parallel_context` to make it explicit
             # that any operations happening from hereon are done on the context
@@ -357,11 +369,18 @@ def parallel_in_child_context(parallel_context) -> BatchResult[T]:
                 config=config,
                 execution_state=self.state,
                 parallel_context=parallel_context,
+                operation_identifier=operation_identifier,
             )
 
-        return self.run_in_child_context(
+        step_name: str | None = self._resolve_step_name(name, parallel_in_child_context)
+        operation_identifier = OperationIdentifier(
+            operation_id=operation_id, parent_id=self._parent_id, name=step_name
+        )
+
+        return child_handler(
             func=parallel_in_child_context,
-            name=name,
+            state=self.state,
+            operation_identifier=operation_identifier,
             config=ChildConfig(
                 sub_type=OperationSubType.PARALLEL,
                 serdes=config.serdes if config is not None else None,

src/aws_durable_execution_sdk_python/operation/map.py

@@ -17,6 +17,7 @@
 
 if TYPE_CHECKING:
     from aws_durable_execution_sdk_python.context import DurableContext
+    from aws_durable_execution_sdk_python.identifier import OperationIdentifier
     from aws_durable_execution_sdk_python.serdes import SerDes
     from aws_durable_execution_sdk_python.state import ExecutionState
     from aws_durable_execution_sdk_python.types import SummaryGenerator
@@ -94,6 +95,7 @@ def map_handler(
     config: MapConfig | None,
     execution_state: ExecutionState,
     map_context: DurableContext,
+    operation_identifier: OperationIdentifier,
 ) -> BatchResult[R]:
     """Execute a callable for each item in parallel."""
     # Summary Generator Construction (matches TypeScript implementation):
@@ -107,6 +109,14 @@ def map_handler(
         func=func,
         config=config or MapConfig(summary_generator=MapSummaryGenerator()),
     )
+
+    checkpoint = execution_state.get_checkpoint_result(
+        operation_identifier.operation_id
+    )
+    if checkpoint.is_succeeded():
+        # if we've reached this point, then not only is the step succeeded, but it is also `replay_children`.
+        # `map_handler` only executes on replay children and gives back
+        return executor.replay(execution_state, map_context)
     # we are making it explicit that we are now executing within the map_context
     return executor.execute(execution_state, executor_context=map_context)
 

src/aws_durable_execution_sdk_python/operation/parallel.py

@@ -14,6 +14,7 @@
 if TYPE_CHECKING:
     from aws_durable_execution_sdk_python.concurrency import BatchResult
     from aws_durable_execution_sdk_python.context import DurableContext
+    from aws_durable_execution_sdk_python.identifier import OperationIdentifier
     from aws_durable_execution_sdk_python.serdes import SerDes
     from aws_durable_execution_sdk_python.state import ExecutionState
     from aws_durable_execution_sdk_python.types import SummaryGenerator
@@ -82,6 +83,7 @@ def parallel_handler(
     config: ParallelConfig | None,
     execution_state: ExecutionState,
     parallel_context: DurableContext,
+    operation_identifier: OperationIdentifier,
 ) -> BatchResult[R]:
     """Execute multiple operations in parallel."""
     # Summary Generator Construction (matches TypeScript implementation):
@@ -94,6 +96,12 @@ def parallel_handler(
         callables,
         config or ParallelConfig(summary_generator=ParallelSummaryGenerator()),
     )
+
+    checkpoint = execution_state.get_checkpoint_result(
+        operation_identifier.operation_id
+    )
+    if checkpoint.is_succeeded():
+        return executor.replay(execution_state, parallel_context)
     return executor.execute(execution_state, executor_context=parallel_context)