From e1293c759547ff38f55e41d723f32c9917db26e5 Mon Sep 17 00:00:00 2001 From: Ayushi Ahjolia Date: Fri, 17 Jul 2026 15:26:03 -0700 Subject: [PATCH] fix: round-trip serdes result on first run across ops --- .../examples-catalog.json | 8 +- .../src/serdes_roundtrip/serdes_roundtrip.py | 128 +++++++++ .../src/step/step_with_custom_serdes.py | 65 ----- .../serdes_roundtrip/test_serdes_roundtrip.py | 37 +++ .../test/step/test_step_with_custom_serdes.py | 44 --- .../lambda_service.py | 6 +- .../operation/child.py | 72 +++-- .../operation/step.py | 3 +- .../operation/wait_for_condition.py | 11 +- .../tests/concurrency_test.py | 7 +- .../tests/context_test.py | 8 +- .../e2e/custom_serdes_roundtrip_int_test.py | 266 +++++++++++++++++- .../tests/operation/child_test.py | 186 +++++++++++- .../tests/operation/map_test.py | 56 +++- .../tests/operation/parallel_test.py | 55 +++- .../operation/wait_for_condition_test.py | 136 +++++++++ 16 files changed, 908 insertions(+), 180 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-examples/src/serdes_roundtrip/serdes_roundtrip.py delete mode 100644 packages/aws-durable-execution-sdk-python-examples/src/step/step_with_custom_serdes.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/test/serdes_roundtrip/test_serdes_roundtrip.py delete mode 100644 packages/aws-durable-execution-sdk-python-examples/test/step/test_step_with_custom_serdes.py diff --git a/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json b/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json index 86878529..0c647339 100644 --- a/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json +++ b/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json @@ -46,15 +46,15 @@ "path": "./src/step/step_with_retry.py" }, { - "name": "Step with Custom SerDes", - "description": "Step with a non-identity custom SerDes; the step returns the canonical round-tripped value, consistent across first run and replay", - "handler": "step_with_custom_serdes.handler", + "name": "Custom SerDes Round-Trip", + "description": "One function proving the first-run/replay result-equality guarantee with a non-identity SerDes across step, wait_for_condition, and run_in_child_context (normal, virtual, large-payload)", + "handler": "serdes_roundtrip.handler", "integration": true, "durableConfig": { "RetentionPeriodInDays": 7, "ExecutionTimeout": 300 }, - "path": "./src/step/step_with_custom_serdes.py" + "path": "./src/serdes_roundtrip/serdes_roundtrip.py" }, { "name": "Wait State", diff --git a/packages/aws-durable-execution-sdk-python-examples/src/serdes_roundtrip/serdes_roundtrip.py b/packages/aws-durable-execution-sdk-python-examples/src/serdes_roundtrip/serdes_roundtrip.py new file mode 100644 index 00000000..685f2fb7 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/serdes_roundtrip/serdes_roundtrip.py @@ -0,0 +1,128 @@ +"""Custom (non-identity) SerDes round-trip across every operation. + +Demonstrates the first-run / replay result-equality guarantee: with a +non-identity ``SerDes``, each durable operation returns the *round-tripped* +value - the value produced by ``serialize`` then ``deserialize`` - on the first +run, which is exactly the value a replay reconstructs from the checkpoint. +Returning the raw in-memory result on the first run would make the first run and +replay disagree whenever the serdes is not a perfect round-trip. + +A single deployed function exercises the guarantee across every operation that +checkpoints a serialized result, so there is no need to deploy a separate +example per operation: + +* ``step`` +* ``wait_for_condition`` +* ``run_in_child_context`` - normal, virtual, and large-payload (ReplayChildren) + +The shared ``MarkerSerDes`` strips a ``round_tripped`` marker on ``serialize`` +and re-adds it on ``deserialize``, so ``deserialize(serialize(x)) != x``. Every +operation's result therefore carries ``round_tripped=True`` only if the SDK +handed back the canonical round-tripped value rather than the raw one. +""" + +import json +from typing import Any + +from aws_durable_execution_sdk_python.config import ChildConfig, StepConfig +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python.serdes import SerDes, SerDesContext +from aws_durable_execution_sdk_python.waits import ( + WaitForConditionConfig, + WaitForConditionDecision, +) + +# Results larger than this are not checkpointed in full; the child context +# switches to ReplayChildren mode (a compact summary is checkpointed and the +# child re-executes on replay). Kept in sync with the SDK constant. +CHECKPOINT_SIZE_LIMIT_BYTES = 256 * 1024 + + +class MarkerSerDes(SerDes[dict[str, Any]]): + """Non-identity serdes: ``deserialize`` re-adds a marker ``serialize`` strips. + + ``deserialize(serialize(value)) == {**value, "round_tripped": True}``, which + differs from ``value`` whenever ``value`` lacks the marker. That makes the + round-trip observable in the value each operation returns. + """ + + def serialize(self, value: dict[str, Any], _: SerDesContext) -> str: + payload = {k: v for k, v in dict(value).items() if k != "round_tripped"} + return json.dumps(payload) + + def deserialize(self, data: str, _: SerDesContext) -> dict[str, Any]: + return {**json.loads(data), "round_tripped": True} + + +def _large_child_summary(_result: dict[str, Any]) -> str: + """Compact summary checkpointed in place of the large child result.""" + return json.dumps({"type": "large-child-result"}) + + +def _stop_immediately( + _state: dict[str, Any], _attempt: int +) -> WaitForConditionDecision: + """Meet the wait_for_condition on the first check.""" + return WaitForConditionDecision.stop_polling() + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, bool]: + """Run every operation with a non-identity serdes and report the round-trip. + + Each returned value carries ``round_tripped=True`` only because the SDK + returned ``deserialize(serialize(result))`` on the first run - the same + value the corresponding replay path produces. + """ + serdes = MarkerSerDes() + + step_result = context.step( + lambda _step_ctx: {"op": "step"}, + name="step", + config=StepConfig(serdes=serdes), + ) + + child_result = context.run_in_child_context( + lambda _child_ctx: {"op": "child"}, + name="child", + config=ChildConfig(serdes=serdes), + ) + + virtual_result = context.run_in_child_context( + lambda _child_ctx: {"op": "virtual-child"}, + name="virtual-child", + config=ChildConfig(serdes=serdes, is_virtual=True), + ) + + # A result larger than the checkpoint limit forces ReplayChildren mode: only + # the summary is checkpointed, yet the returned value is still the full + # round-tripped result. + large_blob = "x" * (CHECKPOINT_SIZE_LIMIT_BYTES + 1) + large_result = context.run_in_child_context( + lambda _child_ctx: {"op": "large-child", "blob": large_blob}, + name="large-child", + config=ChildConfig(serdes=serdes, summary_generator=_large_child_summary), + ) + + wait_for_condition_result = context.wait_for_condition( + check=lambda _state, _ctx: {"op": "wait-for-condition"}, + config=WaitForConditionConfig( + initial_state={}, + wait_strategy=_stop_immediately, + serdes=serdes, + ), + name="wait-for-condition", + ) + + # Only the marker booleans are returned (never the large blob) so the + # handler result stays small and easy to assert on. + return { + "step_round_tripped": step_result.get("round_tripped", False), + "child_round_tripped": child_result.get("round_tripped", False), + "virtual_child_round_tripped": virtual_result.get("round_tripped", False), + "large_child_round_tripped": large_result.get("round_tripped", False), + "wait_for_condition_round_tripped": wait_for_condition_result.get( + "round_tripped", False + ), + } diff --git a/packages/aws-durable-execution-sdk-python-examples/src/step/step_with_custom_serdes.py b/packages/aws-durable-execution-sdk-python-examples/src/step/step_with_custom_serdes.py deleted file mode 100644 index 7657c644..00000000 --- a/packages/aws-durable-execution-sdk-python-examples/src/step/step_with_custom_serdes.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Step with a custom (non-identity) SerDes. - -Demonstrates that a step returns the *canonical, round-tripped* value — the -value produced by ``serialize`` then ``deserialize`` — rather than the raw -in-memory object the step function returned. Because a step's result is -deserialized from the checkpoint on replay, returning the round-tripped value -on the first run keeps the result identical across the first run and every -replay. - -Here the custom SerDes normalizes the order on the way to the checkpoint: it -persists only the canonical fields, sorts the tags, and drops a transient, -non-persisted field. The handler uses the step result immediately, so the value -observed on the first run is already the canonical form (sorted tags, transient -field dropped) — the same value a replay would deserialize from the checkpoint. -""" - -import json -from typing import Any - -from aws_durable_execution_sdk_python.config import StepConfig -from aws_durable_execution_sdk_python.context import DurableContext -from aws_durable_execution_sdk_python.execution import durable_execution -from aws_durable_execution_sdk_python.serdes import SerDes, SerDesContext - - -class OrderSerDes(SerDes[dict[str, Any]]): - """Persist a canonical order: keep stable fields, sort tags, drop transients. - - ``deserialize(serialize(order))`` is intentionally not the identity — the - transient field is dropped and the tags are sorted — so the difference - between the raw result and the round-tripped result is observable. - """ - - def serialize(self, value: dict[str, Any], _: SerDesContext) -> str: - canonical: dict[str, Any] = { - "order_id": value["order_id"], - "tags": sorted(value["tags"]), - } - return json.dumps(canonical) - - def deserialize(self, data: str, _: SerDesContext) -> dict[str, Any]: - return json.loads(data) - - -@durable_execution -def handler(_event: Any, context: DurableContext) -> dict[str, Any]: - """Build an order in a step and use its round-tripped result immediately.""" - order = context.step( - lambda _: { - "order_id": "ORD-42", - "tags": ["priority", "gift", "fragile"], # unsorted on purpose - "transient_score": 0.99, # not persisted by OrderSerDes - }, - name="build_order", - config=StepConfig(serdes=OrderSerDes()), - ) - - # `order` is the canonical, round-tripped value even on the first run: - # tags are sorted and the transient field is gone. A replay would produce - # the exact same value by deserializing the checkpoint. - return { - "order_id": order["order_id"], - "tags": order["tags"], # canonical: sorted - "has_transient": "transient_score" in order, # False after round-trip - } diff --git a/packages/aws-durable-execution-sdk-python-examples/test/serdes_roundtrip/test_serdes_roundtrip.py b/packages/aws-durable-execution-sdk-python-examples/test/serdes_roundtrip/test_serdes_roundtrip.py new file mode 100644 index 00000000..d52f660a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/test/serdes_roundtrip/test_serdes_roundtrip.py @@ -0,0 +1,37 @@ +"""Tests for the combined custom-serdes round-trip example.""" + +import pytest +from aws_durable_execution_sdk_python.execution import InvocationStatus + +from src.serdes_roundtrip import serdes_roundtrip +from test.conftest import deserialize_operation_payload + + +@pytest.mark.example +@pytest.mark.durable_execution( + handler=serdes_roundtrip.handler, + lambda_function_name="Custom SerDes Round-Trip", +) +def test_serdes_roundtrip_all_operations(durable_runner): + """Every operation returns the canonical, round-tripped value on first run. + + With a non-identity serdes, each operation (step, wait_for_condition, and + every run_in_child_context variant - normal, virtual, large-payload) must + return ``deserialize(serialize(result))`` on the first run, so the returned + value matches what a replay reconstructs. The handler reports a boolean per + operation that is True only when the marker added by ``deserialize`` is + present. + """ + with durable_runner: + result = durable_runner.run(input="test", timeout=15) + + assert result.status is InvocationStatus.SUCCEEDED + + result_data = deserialize_operation_payload(result.result) + assert result_data == { + "step_round_tripped": True, + "child_round_tripped": True, + "virtual_child_round_tripped": True, + "large_child_round_tripped": True, + "wait_for_condition_round_tripped": True, + } diff --git a/packages/aws-durable-execution-sdk-python-examples/test/step/test_step_with_custom_serdes.py b/packages/aws-durable-execution-sdk-python-examples/test/step/test_step_with_custom_serdes.py deleted file mode 100644 index 8cc3a66b..00000000 --- a/packages/aws-durable-execution-sdk-python-examples/test/step/test_step_with_custom_serdes.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Tests for the step-with-custom-serdes example.""" - -import pytest -from aws_durable_execution_sdk_python.execution import InvocationStatus - -from src.step import step_with_custom_serdes -from src.step.step_with_custom_serdes import OrderSerDes -from test.conftest import deserialize_operation_payload - - -@pytest.mark.example -@pytest.mark.durable_execution( - handler=step_with_custom_serdes.handler, - lambda_function_name="Step with Custom SerDes", -) -def test_step_with_custom_serdes(durable_runner): - """The handler returns the canonical, round-tripped step result. - - Even on the first run the step returns the value produced by - OrderSerDes.serialize then OrderSerDes.deserialize, so the result reflects - the canonical form (sorted tags, transient field dropped) rather than the - raw value the step function produced. This is the same value a replay would - deserialize from the checkpoint. - """ - with durable_runner: - result = durable_runner.run(input="test", timeout=10) - - assert result.status is InvocationStatus.SUCCEEDED - - result_data = deserialize_operation_payload(result.result) - assert result_data == { - "order_id": "ORD-42", - "tags": ["fragile", "gift", "priority"], # sorted by the serdes - "has_transient": False, # transient field dropped on serialize - } - - # The step checkpoint stores the canonical payload, and deserializing it - # yields exactly what the handler returned for the order fields. - step_result = result.get_step("build_order") - canonical = deserialize_operation_payload(step_result.result, OrderSerDes()) - assert canonical == { - "order_id": "ORD-42", - "tags": ["fragile", "gift", "priority"], - } 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 0843956b..9e1fb516 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 @@ -562,7 +562,7 @@ def create_context_start( def create_context_succeed( cls, identifier: OperationIdentifier, - payload: str, + payload: str | None, sub_type: OperationSubType, context_options: ContextOptions | None = None, ) -> OperationUpdate: @@ -624,7 +624,7 @@ def create_execution_fail(cls, error: ErrorObject) -> OperationUpdate: # region step @classmethod def create_step_succeed( - cls, identifier: OperationIdentifier, payload: str + cls, identifier: OperationIdentifier, payload: str | None ) -> OperationUpdate: """Create an instance of OperationUpdate for type: STEP, action: SUCCEED.""" return cls( @@ -726,7 +726,7 @@ def create_wait_for_condition_start( @classmethod def create_wait_for_condition_succeed( - cls, identifier: OperationIdentifier, payload: str + cls, identifier: OperationIdentifier, payload: str | None ) -> OperationUpdate: """Create an instance of OperationUpdate for type: STEP, action: SUCCEED.""" return cls( diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/child.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/child.py index 2789301d..d4103e29 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/child.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/child.py @@ -133,6 +133,20 @@ def check_result_status(self) -> CheckResult[T]: # Ready to execute (checkpoint exists or was just created) return CheckResult.create_is_ready_to_execute(checkpointed_result) + def _deserialize_payload(self, serialized: str | None) -> T: + """Return the round-tripped value, so the first run matches replay. + + A None payload is returned as-is. + """ + if serialized is None: + return None # type: ignore[return-value] + return deserialize( + serdes=self.config.serdes, + data=serialized, + operation_id=self.operation_identifier.operation_id, + durable_execution_arn=self.state.durable_execution_arn, + ) + def execute(self, checkpointed_result: CheckpointedResult) -> T: """Execute child context function with error handling and large payload support. @@ -162,13 +176,25 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T: ) raw_result: T = wrapped_user_func() + # Serialize once: used as the round-tripped return value in every + # mode, and as the checkpoint payload on the normal path. A custom + # serdes may serialize to None, which is handled below. + serialized_result: str | None = serialize( + serdes=self.config.serdes, + value=raw_result, + operation_id=self.operation_identifier.operation_id, + durable_execution_arn=self.state.durable_execution_arn, + ) + if self.is_virtual: logger.debug( "Virtual context: Exiting child context without creating another checkpoint. id: %s, name: %s", self.operation_identifier.operation_id, self.operation_identifier.name, ) - return raw_result + # Virtual contexts never checkpoint and re-execute on replay; + # round-trip so the first run matches replay. + return self._deserialize_payload(serialized_result) # If in replay_children mode, return without checkpointing if checkpointed_result.is_replay_children(): @@ -177,30 +203,19 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T: self.operation_identifier.operation_id, self.operation_identifier.name, ) - return raw_result - - # Serialize result - serialized_result: str = serialize( - serdes=self.config.serdes, - value=raw_result, - operation_id=self.operation_identifier.operation_id, - durable_execution_arn=self.state.durable_execution_arn, - ) - - # Check payload size and use ReplayChildren mode if needed - # Summary Generator Logic: - # When the serialized result exceeds 256KB, we use ReplayChildren mode to avoid - # checkpointing large payloads. Instead, we checkpoint a compact summary and mark - # the operation for replay. This matches the TypeScript implementation behavior. - # - # See TypeScript reference: - # - aws-durable-execution-sdk-js/src/handlers/run-in-child-context-handler/run-in-child-context-handler.ts (lines ~200-220) - # - # The summary generator creates a JSON summary with metadata (type, counts, status) - # instead of the full BatchResult. During replay, the child context is re-executed - # to reconstruct the full result rather than deserializing from the checkpoint. + # Large payloads re-execute on replay; round-trip so the first + # run matches replay. The checkpoint stays small (summary only). + return self._deserialize_payload(serialized_result) + + # Large results checkpoint a compact summary and use ReplayChildren + # so replay re-executes instead of deserializing. The returned value + # always uses the full serialized_result, never the summary. + payload_to_checkpoint: str | None = serialized_result replay_children: bool = False - if len(serialized_result) > CHECKPOINT_SIZE_LIMIT_BYTES: + if ( + serialized_result is not None + and len(serialized_result) > CHECKPOINT_SIZE_LIMIT_BYTES + ): logger.debug( "Large payload detected, using ReplayChildren mode: id: %s, name: %s, payload_size: %d, limit: %d", self.operation_identifier.operation_id, @@ -209,8 +224,8 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T: CHECKPOINT_SIZE_LIMIT_BYTES, ) replay_children = True - # Use summary generator if provided, otherwise use empty string (matches TypeScript) - serialized_result = ( + # Use the summary generator if provided, otherwise an empty string. + payload_to_checkpoint = ( self.config.summary_generator(raw_result) if self.config.summary_generator else "" @@ -219,7 +234,7 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T: # Checkpoint SUCCEED success_operation: OperationUpdate = OperationUpdate.create_context_succeed( identifier=self.operation_identifier, - payload=serialized_result, + payload=payload_to_checkpoint, sub_type=self.sub_type, context_options=ContextOptions(replay_children=replay_children), ) @@ -234,7 +249,8 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T: self.operation_identifier.operation_id, self.operation_identifier.name, ) - return raw_result # noqa: TRY300 + # Round-trip so the first run matches replay. + return self._deserialize_payload(serialized_result) # noqa: TRY300 except SuspendExecution: # Don't checkpoint SuspendExecution - let it bubble up raise diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/step.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/step.py index 6c425459..f4596212 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/step.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/step.py @@ -228,7 +228,8 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T: ) raw_result: T = wrapped_user_func(step_context) - serialized_result: str = serialize( + # A custom serdes may serialize to None, which is handled below. + serialized_result: str | None = serialize( serdes=self.config.serdes, value=raw_result, operation_id=self.operation_identifier.operation_id, diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/wait_for_condition.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/wait_for_condition.py index 65858c5a..6dd03804 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/wait_for_condition.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/wait_for_condition.py @@ -232,7 +232,16 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T: self.operation_identifier.operation_id, self.operation_identifier.name, ) - return new_state + # Round-trip so the first run matches replay. A None payload is + # returned as-is. + if serialized_state is None: + return None # type: ignore[return-value] + return deserialize( + serdes=self.config.serdes, + data=serialized_state, + operation_id=self.operation_identifier.operation_id, + durable_execution_arn=self.state.durable_execution_arn, + ) # Condition not met - schedule retry # We enforce a minimum delay second of 1, to match model behaviour. diff --git a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py index 887e91c3..298e15c3 100644 --- a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py @@ -2074,9 +2074,14 @@ def success_callable(): execution_state = Mock() execution_state.create_checkpoint = Mock() + # wrap_user_function must return the real function so the branch result is + # the (serializable) "result_N" string that the child round-trips. + child_context = Mock() + child_context.state.wrap_user_function = lambda func, *a, **k: func + executor_context = Mock() executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa SLF001 - executor_context.create_child_context = lambda *args, **kwargs: Mock() + executor_context.create_child_context = lambda *args, **kwargs: child_context result = executor.execute(execution_state, executor_context) 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 d91c0951..6a81b5d0 100644 --- a/packages/aws-durable-execution-sdk-python/tests/context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/context_test.py @@ -1722,7 +1722,9 @@ def test_function(context, item, index, items): with patch( "aws_durable_execution_sdk_python.context.map_handler" ) as mock_map_handler: - mock_map_handler.return_value = Mock() + # The wrapping child context round-trips this result, so it must be + # serializable. + mock_map_handler.return_value = {"result": "value"} with patch.object(context, "run_in_child_context") as mock_run_in_child: # Set up the mock to call the nested function @@ -1763,7 +1765,9 @@ def test_callable_2(context): with patch( "aws_durable_execution_sdk_python.context.parallel_handler" ) as mock_parallel_handler: - mock_parallel_handler.return_value = Mock() + # The wrapping child context round-trips this result, so it must be + # serializable. + mock_parallel_handler.return_value = {"result": "value"} with patch.object(context, "run_in_child_context") as mock_run_in_child: # Set up the mock to call the nested function diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/custom_serdes_roundtrip_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/custom_serdes_roundtrip_int_test.py index 08d74352..83f7c146 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/custom_serdes_roundtrip_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/custom_serdes_roundtrip_int_test.py @@ -1,10 +1,14 @@ """Integration tests for the first-run/replay result equality guarantee. With a non-identity custom ``SerDes`` (serialize/deserialize is not a round-trip -identity), a step must return the *round-tripped* value on the first run so that -it matches the value returned on replay, where the result is deserialized from -the checkpoint. Returning the raw in-memory result on the first run would make -the first-run and replay results diverge. +identity), an operation must return the *round-tripped* value on the first run +so that it matches the value returned on replay, where the result is +deserialized from the checkpoint. Returning the raw in-memory result on the +first run would make the first-run and replay results diverge. + +This is exercised below for ``step``, ``wait_for_condition``, and +``run_in_child_context`` in all three modes: normal, virtual, and large-payload +(ReplayChildren). The serdes used here strips a marker key on ``serialize`` and re-adds it on ``deserialize`` so that ``deserialize(serialize(x)) != x``. That makes the bug @@ -15,7 +19,7 @@ from typing import Any from unittest.mock import Mock, patch -from aws_durable_execution_sdk_python.config import StepConfig +from aws_durable_execution_sdk_python.config import ChildConfig, StepConfig from aws_durable_execution_sdk_python.context import DurableContext from aws_durable_execution_sdk_python.execution import ( InvocationStatus, @@ -27,6 +31,10 @@ OperationAction, ) from aws_durable_execution_sdk_python.serdes import SerDes, SerDesContext +from aws_durable_execution_sdk_python.waits import ( + WaitForConditionConfig, + WaitForConditionDecision, +) class MarkerSerDes(SerDes[Any]): @@ -183,7 +191,251 @@ def handler(event, context: DurableContext) -> dict[str, Any]: assert replay_result["Status"] == InvocationStatus.SUCCEEDED.value replay_data = json.loads(replay_result["Result"]) - # The core invariant: first-run output equals replay output, and both carry - # the marker added by deserialize() (i.e. the round-tripped value, not raw). + # First run equals replay, and both carry the round-trip marker. + assert first_data == replay_data + assert first_data == {"order_id": "ORD-123", "round_tripped": True} + + +def test_child_first_run_matches_replay_with_non_identity_serdes(): + """First-run result equals replay result for run_in_child_context.""" + + @durable_execution + def handler(event, context: DurableContext) -> dict[str, Any]: + return context.run_in_child_context( + lambda _child_ctx: {"order_id": "ORD-123"}, + name="process_order", + config=ChildConfig(serdes=MarkerSerDes()), + ) + + # --- First invocation --- + checkpoint_calls, mock_checkpoint = _patched_client() + 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 + + first_result = handler(_create_initial_event(), _create_lambda_context()) + + assert first_result["Status"] == InvocationStatus.SUCCEEDED.value + first_data = json.loads(first_result["Result"]) + + # The child's SUCCEED payload is the *serialized* value (marker stripped). + all_operations = [op for batch in checkpoint_calls for op in batch] + child_succeed_ops = [ + op + for op in all_operations + if op.action == OperationAction.SUCCEED and op.operation_id != "execution-1" + ] + assert len(child_succeed_ops) == 1 + child_payload = child_succeed_ops[0].payload + assert json.loads(child_payload) == {"order_id": "ORD-123"} # no marker + + # --- Replay: rebuild the execution with the child already SUCCEEDED --- + from tests.test_helpers import operation_id_sequence + + child_id = next(operation_id_sequence()) + + checkpoint_calls_replay, mock_checkpoint_replay = _patched_client() + 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_replay + + replay_event = _create_replay_event( + [ + { + "Id": child_id, + "Type": "CONTEXT", + "Status": "SUCCEEDED", + "ParentId": "execution-1", + "ContextDetails": {"Result": child_payload}, + } + ] + ) + replay_result = handler(replay_event, _create_lambda_context()) + + assert replay_result["Status"] == InvocationStatus.SUCCEEDED.value + replay_data = json.loads(replay_result["Result"]) + + # First run equals replay, and both carry the round-trip marker. + assert first_data == replay_data + assert first_data == {"order_id": "ORD-123", "round_tripped": True} + + +def test_wait_for_condition_first_run_matches_replay_with_non_identity_serdes(): + """First-run result equals replay result for wait_for_condition.""" + + @durable_execution + def handler(event, context: DurableContext) -> dict[str, Any]: + return context.wait_for_condition( + check=lambda _state, _ctx: {"order_id": "ORD-123"}, + config=WaitForConditionConfig( + initial_state={}, + wait_strategy=lambda _s, _a: WaitForConditionDecision.stop_polling(), + serdes=MarkerSerDes(), + ), + name="process_order", + ) + + # --- First invocation --- + checkpoint_calls, mock_checkpoint = _patched_client() + 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 + + first_result = handler(_create_initial_event(), _create_lambda_context()) + + assert first_result["Status"] == InvocationStatus.SUCCEEDED.value + first_data = json.loads(first_result["Result"]) + + # The SUCCEED payload is the *serialized* state (marker stripped). A + # wait_for_condition operation is checkpointed as a STEP-typed operation. + all_operations = [op for batch in checkpoint_calls for op in batch] + succeed_ops = [ + op + for op in all_operations + if op.action == OperationAction.SUCCEED and op.operation_id != "execution-1" + ] + assert len(succeed_ops) == 1 + wfc_payload = succeed_ops[0].payload + assert json.loads(wfc_payload) == {"order_id": "ORD-123"} # no marker + + # --- Replay: rebuild the execution with the operation already SUCCEEDED --- + from tests.test_helpers import operation_id_sequence + + wfc_id = next(operation_id_sequence()) + + checkpoint_calls_replay, mock_checkpoint_replay = _patched_client() + 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_replay + + replay_event = _create_replay_event( + [ + { + "Id": wfc_id, + "Type": "STEP", + "Status": "SUCCEEDED", + "ParentId": "execution-1", + "StepDetails": {"Result": wfc_payload}, + } + ] + ) + replay_result = handler(replay_event, _create_lambda_context()) + + assert replay_result["Status"] == InvocationStatus.SUCCEEDED.value + replay_data = json.loads(replay_result["Result"]) + + # First run equals replay, and both carry the round-trip marker. assert first_data == replay_data assert first_data == {"order_id": "ORD-123", "round_tripped": True} + + +def test_virtual_child_first_run_returns_round_tripped_result(): + """A virtual child returns the round-tripped value and writes no checkpoint.""" + + @durable_execution + def handler(event, context: DurableContext) -> dict[str, Any]: + return context.run_in_child_context( + lambda _child_ctx: {"order_id": "ORD-123"}, + name="process_order", + config=ChildConfig(serdes=MarkerSerDes(), is_virtual=True), + ) + + checkpoint_calls, mock_checkpoint = _patched_client() + 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 + + first_result = handler(_create_initial_event(), _create_lambda_context()) + + assert first_result["Status"] == InvocationStatus.SUCCEEDED.value + first_data = json.loads(first_result["Result"]) + + # Virtual contexts re-execute on replay and never checkpoint, so the + # round-tripped first-run result is what every replay reproduces. + assert first_data == {"order_id": "ORD-123", "round_tripped": True} + all_operations = [op for batch in checkpoint_calls for op in batch] + assert not [op for op in all_operations if op.operation_id != "execution-1"] + + +def test_large_child_first_run_matches_replay_with_non_identity_serdes(): + """First-run result equals replay result for a large-payload child.""" + blob = "x" * (256 * 1024 + 100) + + @durable_execution + def handler(event, context: DurableContext) -> dict[str, Any]: + return context.run_in_child_context( + lambda _child_ctx: {"order_id": "ORD-123", "blob": blob}, + name="process_order", + config=ChildConfig(serdes=MarkerSerDes()), + ) + + # --- First invocation: large result -> ReplayChildren + summary checkpoint --- + checkpoint_calls, mock_checkpoint = _patched_client() + 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 + + first_result = handler(_create_initial_event(), _create_lambda_context()) + + assert first_result["Status"] == InvocationStatus.SUCCEEDED.value + first_data = json.loads(first_result["Result"]) + + all_operations = [op for batch in checkpoint_calls for op in batch] + child_succeed_ops = [ + op + for op in all_operations + if op.action == OperationAction.SUCCEED and op.operation_id != "execution-1" + ] + assert len(child_succeed_ops) == 1 + # The full result is not checkpointed; only a summary with ReplayChildren. + assert child_succeed_ops[0].context_options.replay_children is True + + # --- Replay: child re-executes (ReplayChildren) and round-trips --- + from tests.test_helpers import operation_id_sequence + + child_id = next(operation_id_sequence()) + + checkpoint_calls_replay, mock_checkpoint_replay = _patched_client() + 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_replay + + replay_event = _create_replay_event( + [ + { + "Id": child_id, + "Type": "CONTEXT", + "Status": "SUCCEEDED", + "ParentId": "execution-1", + "ContextDetails": {"Result": "", "ReplayChildren": True}, + } + ] + ) + replay_result = handler(replay_event, _create_lambda_context()) + + assert replay_result["Status"] == InvocationStatus.SUCCEEDED.value + replay_data = json.loads(replay_result["Result"]) + + assert first_data == replay_data + assert first_data == {"order_id": "ORD-123", "blob": blob, "round_tripped": True} diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/child_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/child_test.py index 1f8617ca..76a4ae60 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/child_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/child_test.py @@ -3,7 +3,7 @@ from __future__ import annotations import json -from typing import cast +from typing import Any, cast from unittest.mock import Mock import pytest @@ -21,11 +21,29 @@ OperationType, ) from aws_durable_execution_sdk_python.operation.child import child_handler +from aws_durable_execution_sdk_python.serdes import SerDes, SerDesContext from aws_durable_execution_sdk_python.state import ExecutionState from aws_durable_execution_sdk_python.types import SummaryGenerator from tests.serdes_test import CustomDictSerDes +class _NonIdentitySerDes(SerDes[Any]): + """deserialize() adds a marker that serialize() never removes. + + Makes ``deserialize(serialize(x)) != x`` so first-run vs replay divergence + is observable. + """ + + def serialize(self, value: Any, _: SerDesContext) -> str: + payload = dict(value) + payload.pop("deserialized", None) + return json.dumps(payload) + + def deserialize(self, data: str, _: SerDesContext) -> dict[str, Any]: + parsed = json.loads(data) + return {**parsed, "deserialized": True} + + # region child_handler @pytest.mark.parametrize( ("config", "expected_sub_type"), @@ -1006,3 +1024,169 @@ def setup_mocks(): assert result2 == "test_result" assert mock_state2.create_checkpoint.call_count == 0 # No checkpoints + + +# region first-run round-trip +def test_child_handler_first_run_returns_round_tripped_result(): + """First-run result must equal the replay (deserialized-from-checkpoint) result. + + With a non-identity SerDes, returning the raw child result on the first run + diverges from the value replay reconstructs by deserializing the checkpoint. + The first run must return the serialize-then-deserialize value so both agree. + """ + serdes = _NonIdentitySerDes() + raw_result = {"key": "value"} + config = ChildConfig(serdes=serdes) + op_id = OperationIdentifier( + "child_rt", OperationSubType.RUN_IN_CHILD_CONTEXT, None, "test_name" + ) + + # --- First run --- + first_state = Mock(spec=ExecutionState) + first_state.durable_execution_arn = "test_arn" + first_result = Mock() + first_result.is_succeeded.return_value = False + first_result.is_failed.return_value = False + first_result.is_started.return_value = False + first_result.is_replay_children.return_value = False + first_result.is_existent.return_value = False + first_state.get_checkpoint_result.return_value = first_result + mock_callable = Mock(return_value=raw_result) + first_state.wrap_user_function.return_value = mock_callable + + first_run_result = child_handler(mock_callable, first_state, op_id, config) + + # Grab the payload actually checkpointed (START is call 0, SUCCEED is call 1). + success_call = first_state.create_checkpoint.call_args_list[1] + checkpointed_payload = success_call[1]["operation_update"].payload + + # --- Replay: checkpoint already SUCCEEDED (not replay_children) --- + replay_state = Mock(spec=ExecutionState) + replay_state.durable_execution_arn = "test_arn" + replay_result = Mock() + replay_result.is_succeeded.return_value = True + replay_result.is_replay_children.return_value = False + replay_result.result = checkpointed_payload + replay_state.get_checkpoint_result.return_value = replay_result + + replayed = child_handler( + Mock(return_value="should_not_call"), replay_state, op_id, config + ) + + assert first_run_result == {"key": "value", "deserialized": True} + assert first_run_result == replayed + assert first_run_result != raw_result + + +def test_child_handler_first_run_none_payload_skips_deserialize(): + """A None serialized payload is returned as-is without deserializing. + + Mirrors the replay path, which returns None (without deserializing) when the + checkpointed result is None. + """ + + class NonePayloadSerDes(SerDes[Any]): + """serialize() yields None; deserialize() must never be called for None.""" + + def serialize(self, _value: Any, _ctx: SerDesContext) -> str: + return None # type: ignore[return-value] + + def deserialize(self, _data: str, _ctx: SerDesContext) -> Any: + msg = "deserialize should not be called for a None payload" + raise AssertionError(msg) + + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + mock_result = Mock() + mock_result.is_succeeded.return_value = False + mock_result.is_failed.return_value = False + mock_result.is_started.return_value = False + mock_result.is_replay_children.return_value = False + mock_result.is_existent.return_value = False + mock_state.get_checkpoint_result.return_value = mock_result + mock_callable = Mock(return_value={"key": "value"}) + mock_state.wrap_user_function.return_value = mock_callable + + result = child_handler( + mock_callable, + mock_state, + OperationIdentifier( + "child_none", OperationSubType.RUN_IN_CHILD_CONTEXT, None, "test_name" + ), + ChildConfig(serdes=NonePayloadSerDes()), + ) + + assert result is None + + +def test_child_handler_virtual_returns_round_tripped_result(): + """A virtual child returns the serdes round-tripped result. + + Virtual contexts write no checkpoint and re-execute on replay, so the + round-trip keeps the returned value identical across first run and replay. + """ + serdes = _NonIdentitySerDes() + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + mock_result = Mock() + mock_result.is_succeeded.return_value = False + mock_result.is_failed.return_value = False + mock_result.is_started.return_value = False + mock_result.is_replay_children.return_value = False + mock_result.is_existent.return_value = False + mock_state.get_checkpoint_result.return_value = mock_result + mock_callable = Mock(return_value={"key": "value"}) + mock_state.wrap_user_function.return_value = mock_callable + + result = child_handler( + mock_callable, + mock_state, + OperationIdentifier( + "child_virtual_rt", OperationSubType.RUN_IN_CHILD_CONTEXT, None, "test_name" + ), + ChildConfig(serdes=serdes, is_virtual=True), + ) + + # Round-tripped result (deserialize adds the marker). + assert result == {"key": "value", "deserialized": True} + # Virtual contexts still write no checkpoints. + assert mock_state.create_checkpoint.call_count == 0 + + +def test_child_handler_replay_children_returns_round_tripped_result(): + """A large-payload (ReplayChildren) child round-trips its re-executed result. + + In ReplayChildren mode only a summary is checkpointed and the child + re-executes on replay. The re-executed result is round-tripped so it matches + the value the first run returned; only the checkpoint payload stays small. + """ + serdes = _NonIdentitySerDes() + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + mock_result = Mock() + mock_result.is_succeeded.return_value = True + mock_result.is_failed.return_value = False + mock_result.is_started.return_value = True + mock_result.is_replay_children.return_value = True + mock_result.is_existent.return_value = True + mock_state.get_checkpoint_result.return_value = mock_result + mock_callable = Mock(return_value={"key": "value"}) + mock_state.wrap_user_function.return_value = mock_callable + + result = child_handler( + mock_callable, + mock_state, + OperationIdentifier( + "child_replay_rt", OperationSubType.RUN_IN_CHILD_CONTEXT, None, "test_name" + ), + ChildConfig(serdes=serdes), + ) + + # Round-tripped result (deserialize adds the marker). + assert result == {"key": "value", "deserialized": True} + # ReplayChildren re-executes the function and writes no new checkpoint. + mock_callable.assert_called_once() + mock_state.create_checkpoint.assert_not_called() + + +# endregion first-run round-trip diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py index 5ebae34f..069f307c 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py @@ -437,8 +437,10 @@ def get_checkpoint_result(self, operation_id): operation_identifier, ) - # Verify execute was called - assert result.all[0].result == "RESULT_TEST_ITEM" + # Verify execute was called. The item result is the serdes round-tripped + # value: CustomStrSerDes uppercases on serialize and lowercases on + # deserialize, so "RESULT_TEST_ITEM" round-trips to "result_test_item". + assert result.all[0].result == "result_test_item" def test_map_handler_with_summary_generator(): @@ -455,7 +457,9 @@ def mock_summary_generator(result): executor_context = Mock() executor_context._create_step_id_for_logical_step = Mock(side_effect=["1", "2"]) # noqa SLF001 - executor_context.create_child_context = Mock(return_value=Mock()) + _child_ctx = Mock() + _child_ctx.state.wrap_user_function = lambda func, *a, **k: func + executor_context.create_child_context = Mock(return_value=_child_ctx) class MockExecutionState: def get_checkpoint_result(self, operation_id): @@ -515,7 +519,9 @@ def callable_func(ctx, item, idx, items): executor_context = Mock() executor_context._create_step_id_for_logical_step = Mock(return_value="1") # noqa SLF001 - executor_context.create_child_context = Mock(return_value=Mock()) # SLF001 + _child_ctx = Mock() + _child_ctx.state.wrap_user_function = lambda func, *a, **k: func + executor_context.create_child_context = Mock(return_value=_child_ctx) # SLF001 class MockExecutionState: def get_checkpoint_result(self, operation_id): @@ -595,7 +601,9 @@ def get_checkpoint_result(self, operation_id): executor_context._create_step_id_for_logical_step = Mock( # noqa: SLF001 side_effect=["1", "2", "3"] ) - executor_context.create_child_context = Mock(return_value=Mock()) + _child_ctx = Mock() + _child_ctx.state.wrap_user_function = lambda func, *a, **k: func + executor_context.create_child_context = Mock(return_value=_child_ctx) # Call map_handler map_handler( @@ -962,9 +970,13 @@ def create_id(self, i): call.kwargs["operation_id"]: call.kwargs for call in mock_deserialize.call_args_list } - assert set(deserialize_calls) == {"child-0", "child-1"} + # Children deserialize from their checkpoints with the item serdes; the + # parent also deserializes its BatchResult (the first-run round-trip) with + # the batch serdes. + assert set(deserialize_calls) == {"child-0", "child-1", "parent"} assert deserialize_calls["child-0"]["serdes"] is expected assert deserialize_calls["child-1"]["serdes"] is expected + assert deserialize_calls["parent"]["serdes"] is batch_serdes def test_map_result_serialization_roundtrip(): @@ -1068,7 +1080,14 @@ def create_id(self, i): assert len(mock_serdes_serialize.call_args_list) == 3 parent_call = mock_serdes_serialize.call_args_list[2] - assert parent_call[1]["value"] is result + # The value serialized (checkpointed) at the parent level is the raw + # BatchResult. + assert isinstance(parent_call[1]["value"], BatchResult) + # The first run returns the round-trip of that checkpointed payload, + # matching what replay would deserialize from the checkpoint (with + # serialize mocked to a plain string, the round-trip yields that + # string). + assert result == "serialized" finally: importlib.reload(child) @@ -1130,7 +1149,10 @@ def create_id(self, i): parent_call = mock_serialize.call_args_list[2] assert parent_call[1]["serdes"] is None assert isinstance(parent_call[1]["value"], BatchResult) - assert parent_call[1]["value"] is result + # First run returns the round-tripped BatchResult, which equals the + # value serialized into the checkpoint (default serdes round-trips + # as identity). + assert parent_call[1]["value"] == result finally: importlib.reload(child) @@ -1185,8 +1207,16 @@ def create_id(self, i): else f"child-{i}" ) - with patch.object( - DurableContext, "_create_step_id_for_logical_step", create_id + # serialize is mocked to a plain string, so stub deserialize to a + # canonical BatchResult to exercise the round-trip return. + round_tripped = BatchResult( + all=[], completion_reason=CompletionReason.ALL_COMPLETED + ) + with ( + patch.object(child, "deserialize", return_value=round_tripped), + patch.object( + DurableContext, "_create_step_id_for_logical_step", create_id + ), ): context = create_test_context(state=mock_state) result = context.map( @@ -1195,12 +1225,14 @@ def create_id(self, i): config=MapConfig(serdes=custom_serdes), ) - assert isinstance(result, BatchResult) assert len(mock_serialize.call_args_list) == 3 parent_call = mock_serialize.call_args_list[2] assert parent_call[1]["serdes"] is custom_serdes + # The parent serializes a BatchResult with the custom serdes and + # returns the round-tripped BatchResult. assert isinstance(parent_call[1]["value"], BatchResult) - assert parent_call[1]["value"] is result + assert isinstance(result, BatchResult) + assert result is round_tripped finally: importlib.reload(child) diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py index 0a7da40a..2e84c610 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py @@ -432,7 +432,10 @@ def get_checkpoint_result(self, operation_id): operation_identifier, ) - assert result.all[0].result == "RESULT1" + # The branch result is the serdes round-tripped value: CustomStrSerDes + # uppercases on serialize and lowercases on deserialize, so "RESULT1" + # round-trips to "result1". + assert result.all[0].result == "result1" def test_parallel_handler_with_summary_generator(): @@ -460,7 +463,9 @@ def get_checkpoint_result(self, operation_id): executor_context = Mock() executor_context._create_step_id_for_logical_step = Mock(return_value="1") # noqa SLF001 - executor_context.create_child_context = Mock(return_value=Mock()) + _child_ctx = Mock() + _child_ctx.state.wrap_user_function = lambda func, *a, **k: func + executor_context.create_child_context = Mock(return_value=_child_ctx) # Call parallel_handler parallel_handler( @@ -516,7 +521,9 @@ def get_checkpoint_result(self, operation_id): executor_context = Mock() executor_context._create_step_id_for_logical_step = Mock(side_effect=["1", "2"]) # noqa SLF001 - executor_context.create_child_context = Mock(return_value=Mock()) + _child_ctx = Mock() + _child_ctx.state.wrap_user_function = lambda func, *a, **k: func + executor_context.create_child_context = Mock(return_value=_child_ctx) # Call parallel_handler with None config (should use default) parallel_handler( @@ -564,7 +571,9 @@ def get_checkpoint_result(self, operation_id): executor_context._create_step_id_for_logical_step = Mock( # noqa: SLF001 side_effect=["1", "2", "3"] ) - executor_context.create_child_context = Mock(return_value=Mock()) + _child_ctx = Mock() + _child_ctx.state.wrap_user_function = lambda func, *a, **k: func + executor_context.create_child_context = Mock(return_value=_child_ctx) # Call parallel_handler parallel_handler( @@ -921,9 +930,13 @@ def create_id(self, i): expected = item_serdes or batch_serdes calls_by_operation_id = _mock_call_kwargs_by_operation_id(mock_deserialize) - assert set(calls_by_operation_id) == {"child-0", "child-1"} + # Branches deserialize from their checkpoints with the item serdes; the + # parent also deserializes its BatchResult (the first-run round-trip) with + # the batch serdes. + assert set(calls_by_operation_id) == {"child-0", "child-1", "parent"} assert calls_by_operation_id["child-0"]["serdes"] is expected assert calls_by_operation_id["child-1"]["serdes"] is expected + assert calls_by_operation_id["parent"]["serdes"] is batch_serdes def test_parallel_result_serialization_roundtrip(): @@ -1039,7 +1052,14 @@ def create_id(self, i): assert len(mock_serdes_serialize.call_args_list) == 3 parent_call = mock_serdes_serialize.call_args_list[2] - assert parent_call[1]["value"] is result + # The value serialized (checkpointed) at the parent level is the raw + # BatchResult. + assert isinstance(parent_call[1]["value"], BatchResult) + # The first run returns the round-trip of that checkpointed payload, + # matching what replay would deserialize from the checkpoint (with + # serialize mocked to a plain string, the round-trip yields that + # string). + assert result == "serialized" finally: importlib.reload(child) @@ -1101,7 +1121,10 @@ def create_id(self, i): parent_call = mock_serialize.call_args_list[2] assert parent_call[1]["serdes"] is None assert isinstance(parent_call[1]["value"], BatchResult) - assert parent_call[1]["value"] is result + # First run returns the round-tripped BatchResult, which equals the + # value serialized into the checkpoint (default serdes round-trips + # as identity). + assert parent_call[1]["value"] == result finally: importlib.reload(child) @@ -1156,8 +1179,16 @@ def create_id(self, i): else f"child-{i}" ) - with patch.object( - DurableContext, "_create_step_id_for_logical_step", create_id + # serialize is mocked to a plain string, so stub deserialize to a + # canonical BatchResult to exercise the round-trip return. + round_tripped = BatchResult( + all=[], completion_reason=CompletionReason.ALL_COMPLETED + ) + with ( + patch.object(child, "deserialize", return_value=round_tripped), + patch.object( + DurableContext, "_create_step_id_for_logical_step", create_id + ), ): context = create_test_context(state=mock_state) result = context.parallel( @@ -1165,12 +1196,14 @@ def create_id(self, i): config=ParallelConfig(serdes=custom_serdes), ) - assert isinstance(result, BatchResult) assert len(mock_serialize.call_args_list) == 3 parent_call = mock_serialize.call_args_list[2] assert parent_call[1]["serdes"] is custom_serdes + # The parent serializes a BatchResult with the custom serdes and + # returns the round-tripped BatchResult. assert isinstance(parent_call[1]["value"], BatchResult) - assert parent_call[1]["value"] is result + assert isinstance(result, BatchResult) + assert result is round_tripped finally: importlib.reload(child) diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py index 5166cc5a..8fcb7ec1 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py @@ -2,6 +2,7 @@ import datetime import json +from typing import Any from unittest.mock import Mock import pytest @@ -25,6 +26,7 @@ from aws_durable_execution_sdk_python.operation.wait_for_condition import ( WaitForConditionOperationExecutor, ) +from aws_durable_execution_sdk_python.serdes import SerDes, SerDesContext from aws_durable_execution_sdk_python.state import CheckpointedResult, ExecutionState from aws_durable_execution_sdk_python.types import WaitForConditionCheckContext from aws_durable_execution_sdk_python.waits import ( @@ -1472,3 +1474,137 @@ def mock_wait_strategy(state, attempt): assert result == "final_state" assert mock_state.get_checkpoint_result.call_count == 1 # Single check (async) assert mock_state.create_checkpoint.call_count == 2 # START + SUCCESS checkpoints + + +def test_wait_for_condition_first_run_returns_round_tripped_result(): + """First-run result must match the replay (deserialized-from-checkpoint) result. + + With a non-identity SerDes whose serialize/deserialize is not a round-trip + identity, returning the raw check-function result on the first run diverges + from the value returned on replay. The first run must return the value + obtained by serializing then deserializing, so both runs agree. + """ + + class NonIdentitySerDes(SerDes[Any]): + """deserialize() adds a marker that serialize() never removes.""" + + def serialize(self, value: Any, _: SerDesContext) -> str: + payload = dict(value) + payload.pop("deserialized", None) + return json.dumps(payload) + + def deserialize(self, data: str, _: SerDesContext) -> dict[str, Any]: + parsed = json.loads(data) + return {**parsed, "deserialized": True} + + serdes = NonIdentitySerDes() + raw_new_state = {"key": "value"} + + def check_func(_state, _context): + return raw_new_state + + config = WaitForConditionConfig( + initial_state={}, + wait_strategy=lambda s, a: WaitForConditionDecision.stop_polling(), + serdes=serdes, + ) + + # --- First run --- + first_run_state = Mock(spec=ExecutionState) + first_run_state.durable_execution_arn = "test_arn" + first_run_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_not_found() + ) + first_run_state.wrap_user_function.return_value = check_func + mock_logger = Mock(spec=Logger) + mock_logger.with_log_info.return_value = mock_logger + + op_id = OperationIdentifier( + "wfc_rt", OperationSubType.WAIT_FOR_CONDITION, None, "test_wait" + ) + first_run_result = wait_for_condition_handler( + state=first_run_state, + operation_identifier=op_id, + check=check_func, + config=config, + context_logger=mock_logger, + ) + + # Grab the payload that was actually checkpointed (START is call 0, SUCCEED is call 1). + success_call = first_run_state.create_checkpoint.call_args_list[1] + checkpointed_payload = success_call[1]["operation_update"].payload + + # --- Replay: checkpoint already SUCCEEDED with the serialized payload --- + replay_state = Mock(spec=ExecutionState) + replay_state.durable_execution_arn = "test_arn" + succeeded_op = Operation( + operation_id="wfc_rt", + operation_type=OperationType.STEP, + status=OperationStatus.SUCCEEDED, + step_details=StepDetails(result=checkpointed_payload), + ) + replay_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation(succeeded_op) + ) + replay_result = wait_for_condition_handler( + state=replay_state, + operation_identifier=op_id, + check=Mock(return_value="should_not_call"), + config=config, + context_logger=Mock(spec=Logger), + ) + + # First run must return the round-tripped value, which equals the replay value, + # and must NOT equal the raw check-function result. + assert first_run_result == {"key": "value", "deserialized": True} + assert first_run_result == replay_result + assert first_run_result != raw_new_state + + +def test_wait_for_condition_first_run_none_payload_skips_deserialize(): + """A None serialized payload is returned as-is without deserializing. + + This mirrors the replay path, which returns None without calling deserialize + when the checkpointed result is None. A serdes whose serialize returns None + must therefore see its deserialize skipped on the first run too. + """ + + class NonePayloadSerDes(SerDes[Any]): + """serialize() yields None; deserialize() must never be called for None.""" + + def serialize(self, _value: Any, _ctx: SerDesContext) -> str: + return None # type: ignore[return-value] + + def deserialize(self, _data: str, _ctx: SerDesContext) -> Any: + msg = "deserialize should not be called for a None payload" + raise AssertionError(msg) + + def check_func(_state, _context): + return {"key": "value"} + + config = WaitForConditionConfig( + initial_state={}, + wait_strategy=lambda s, a: WaitForConditionDecision.stop_polling(), + serdes=NonePayloadSerDes(), + ) + + 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.wrap_user_function.return_value = check_func + mock_logger = Mock(spec=Logger) + mock_logger.with_log_info.return_value = mock_logger + + result = wait_for_condition_handler( + state=mock_state, + operation_identifier=OperationIdentifier( + "wfc_none", OperationSubType.WAIT_FOR_CONDITION, None, "test_wait" + ), + check=check_func, + config=config, + context_logger=mock_logger, + ) + + assert result is None