diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py index 6592dd79e..afc816dda 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py @@ -51,7 +51,10 @@ CallbackOperationExecutor, wait_for_callback_handler, ) -from aws_durable_execution_sdk_python.operation.child import child_handler +from aws_durable_execution_sdk_python.operation.child import ( + ChildOperationExecutor, + child_handler, +) from aws_durable_execution_sdk_python.operation.invoke import InvokeOperationExecutor from aws_durable_execution_sdk_python.operation.map import map_handler from aws_durable_execution_sdk_python.operation.parallel import parallel_handler @@ -515,12 +518,18 @@ def _next_operation_exists(self) -> bool: @contextmanager def _replay_aware( - self, operation_identifier: OperationIdentifier | None = None + self, + operation_identifier: OperationIdentifier | None = None, + *, + is_virtual: bool = False, ) -> Iterator[None]: """Wrap a single operation with replay-boundary detection. The operation kind is inferred from its own checkpoint (when one - exists), so no per-call-site flag is needed. + exists). The one per-call-site input is `is_virtual`: a virtual child + context never checkpoints itself, so with `is_virtual=True` any + checkpoint at its id is rejected right after the identity check, + before replay and update hooks run. The boundary has these parts: @@ -555,6 +564,10 @@ def _replay_aware( if was_replaying and next_operation is not None: if operation_identifier is not None: operation_identifier.validate_checkpoint(next_operation) + if is_virtual: + ChildOperationExecutor.validate_virtual_checkpoint( + operation_identifier, next_operation + ) # While replaying, an operation that already has a checkpoint was # observed in a prior invocation. If the backend says the operation @@ -601,6 +614,7 @@ def _operation_replay_aware( name: str | None = None, *, operation_type: OperationType, + is_virtual: bool = False, ) -> Iterator[OperationIdentifier]: """Allocate an operation and validate its replay identity before hooks.""" operation_identifier = OperationIdentifier( @@ -610,7 +624,7 @@ def _operation_replay_aware( name=name, operation_type=operation_type, ) - with self._replay_aware(operation_identifier): + with self._replay_aware(operation_identifier, is_virtual=is_virtual): yield operation_identifier # endregion replay status @@ -814,14 +828,15 @@ def run_in_child_context( if config and config.sub_type else OperationSubType.RUN_IN_CHILD_CONTEXT ) + is_virtual: bool = config.is_virtual if config else False with self._operation_replay_aware( sub_type, step_name, operation_type=OperationType.CONTEXT, + is_virtual=is_virtual, ) as operation_identifier: operation_id = operation_identifier.operation_id - is_virtual: bool = config.is_virtual if config else False def callable_with_child_context(): return func( 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 8287be646..f29262673 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 @@ -10,11 +10,13 @@ ChildContextError, ExecutionError, InvocationError, + NonDeterministicExecutionError, SuspendExecution, ) from aws_durable_execution_sdk_python.lambda_service import ( ContextOptions, ErrorObject, + Operation, OperationStatus, OperationSubType, OperationUpdate, @@ -71,6 +73,28 @@ def __init__( self.is_virtual: bool = config.is_virtual self.sub_type = config.sub_type or OperationSubType.RUN_IN_CHILD_CONTEXT + @staticmethod + def validate_virtual_checkpoint( + operation_identifier: OperationIdentifier, + checkpoint: Operation | None, + ) -> None: + """Reject a checkpoint at the operation id of a virtual child context. + + A virtual child never checkpoints itself, so a checkpoint at its id was + written by a non-virtual child in an earlier invocation. Consuming it + would skip the virtual body and leave its inner operations unvalidated. + """ + if not isinstance(checkpoint, Operation): + return + msg = ( + "Non-deterministic child context at " + f"id={operation_identifier.operation_id!r}: checkpoint contains a " + "non-virtual child context but current child context is virtual" + ) + raise NonDeterministicExecutionError( + msg, step_id=operation_identifier.operation_id + ) + def check_result_status(self) -> CheckResult[T]: """Check operation status and create START checkpoint if needed. @@ -82,8 +106,14 @@ def check_result_status(self) -> CheckResult[T]: Raises: ChildContextError: For FAILED operations + NonDeterministicExecutionError: For a virtual child whose operation + id already has a checkpoint """ checkpointed_result = self._get_checkpoint_result() + if self.is_virtual: + self.validate_virtual_checkpoint( + self.operation_identifier, checkpointed_result.operation + ) # Terminal success without replay_children - deserialize and return if ( 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 3a3302e35..e3cc07acd 100644 --- a/packages/aws-durable-execution-sdk-python/tests/context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/context_test.py @@ -3145,6 +3145,67 @@ def on_operation_end(self, info): assert step_body_calls == [] +@pytest.mark.parametrize( + ("checkpoint_status", "updated"), + [ + (OperationStatus.STARTED, False), + (OperationStatus.SUCCEEDED, True), + ], +) +def test_virtual_child_checkpoint_is_rejected_before_replay_hooks( + checkpoint_status: OperationStatus, + updated: bool, +): + """A virtual child fails on a container checkpoint before replay/update hooks.""" + captured: list[str] = [] + + class _CapturingPlugin(DurableInstrumentationPlugin): + def on_operation_start(self, info): + captured.append(f"start:{info.operation_id}") + + def on_operation_end(self, info): + captured.append(f"end:{info.operation_id}") + + plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) + child_body_calls: list[bool] = [] + with plugin_executor.run(): + state = ExecutionState( + durable_execution_arn="arn", + initial_checkpoint_token="token", # noqa: S106 + operations={}, + service_client=Mock(), + plugin_executor=plugin_executor, + updated_operation_ids=[], + ) + ctx = DurableContext( + state=state, + execution_context=ExecutionContext(durable_execution_arn="arn"), + replay_status=ReplayStatus.REPLAY, + ) + next_id = ctx._peek_next_operation_id() # noqa: SLF001 + state._operations[next_id] = Operation( # noqa: SLF001 + operation_id=next_id, + operation_type=OperationType.CONTEXT, + status=checkpoint_status, + sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, + name="child", + ) + if updated: + state._updated_operation_ids.add(next_id) # noqa: SLF001 + + with pytest.raises( + NonDeterministicExecutionError, match="current child context is virtual" + ): + ctx.run_in_child_context( + lambda _child_context: child_body_calls.append(True), + name="child", + config=ChildConfig(is_virtual=True), + ) + + assert captured == [] + assert child_body_calls == [] + + def test_replay_aware_emits_update_hook_for_operation_updated_since_last_invocation(): """Updated terminal operations emit operation_end, not replay start+end.""" captured: list[tuple[str, str, bool, OperationStatus]] = [] diff --git a/packages/aws-durable-execution-sdk-python/tests/execution_test.py b/packages/aws-durable-execution-sdk-python/tests/execution_test.py index 85ea72a97..735ec2e49 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -2957,6 +2957,86 @@ def test_handler(event: Any, context: DurableContext) -> str: assert child_body_calls == [] +@pytest.mark.parametrize( + "child_status", + [OperationStatus.SUCCEEDED, OperationStatus.FAILED, OperationStatus.STARTED], +) +def test_durable_execution_rejects_non_virtual_child_checkpoint_in_virtual_replay( + child_status: OperationStatus, +): + """Replay rejects a child changed from non-virtual to virtual. + + A virtual child never checkpoints itself, so a CONTEXT checkpoint at its id + comes from a non-virtual child. Replay fails before the cached result or + error is used and before the child body runs. + """ + mock_client = Mock(spec=DurableServiceClient) + child_id = OperationIdNamespace().create_id_for_step(1) + execution_operation = Operation( + operation_id="exec1", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + execution_details=ExecutionDetails(input_payload="{}"), + ) + child_operation = Operation( + operation_id=child_id, + operation_type=OperationType.CONTEXT, + status=child_status, + sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, + name="child", + context_details=ContextDetails( + result=( + json.dumps("cached") + if child_status is OperationStatus.SUCCEEDED + else None + ), + error=( + ErrorObject( + message="child failed", + type="ValueError", + data=None, + stack_trace=None, + ) + if child_status is OperationStatus.FAILED + else None + ), + ), + ) + invocation_input = DurableExecutionInvocationInputWithClient( + durable_execution_arn="arn:test:execution/exec1", + checkpoint_token="token123", # noqa: S106 + initial_execution_state=InitialExecutionState( + operations=[execution_operation, child_operation], + next_marker="", + ), + service_client=mock_client, + ) + child_body_calls: list[bool] = [] + + def child_body(_child: DurableContext) -> str: + child_body_calls.append(True) + return "executed" + + @durable_execution + def test_handler(event: Any, context: DurableContext) -> str: + return context.run_in_child_context( + child_body, + name="child", + config=ChildConfig(is_virtual=True), + ) + + result = test_handler(invocation_input, _make_lambda_context()) + + assert result["Status"] == InvocationStatus.FAILED.value + assert ( + result["Error"]["ErrorType"] + == "aws_durable_execution_sdk_python.exceptions.NonDeterministicExecutionError" + ) + assert "current child context is virtual" in result["Error"]["ErrorMessage"] + assert child_body_calls == [] + mock_client.checkpoint.assert_not_called() + + @pytest.mark.parametrize("operation_kind", ["map", "parallel"]) @pytest.mark.parametrize("parent_replay_children", [False, True]) @pytest.mark.parametrize( 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 0c4d44249..a99e7f19e 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 @@ -15,13 +15,16 @@ DurableApiErrorCategory, ExecutionError, InvocationError, + NonDeterministicExecutionError, RetryableSerDesError, SerDesError, StepError, ) from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( + ContextDetails, ErrorObject, + Operation, OperationAction, OperationStatus, OperationSubType, @@ -29,7 +32,7 @@ ) 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.state import CheckpointedResult, ExecutionState from aws_durable_execution_sdk_python.types import SummaryGenerator from tests.serdes_test import ( CustomDictSerDes, @@ -1315,6 +1318,81 @@ def setup_mocks(): assert mock_state2.create_checkpoint.call_count == 0 # No checkpoints +@pytest.mark.parametrize( + ("status", "replay_children"), + [ + (OperationStatus.SUCCEEDED, False), + (OperationStatus.SUCCEEDED, True), + (OperationStatus.FAILED, False), + (OperationStatus.STARTED, False), + ], +) +def test_child_handler_is_virtual_rejects_container_checkpoint( + status: OperationStatus, replay_children: bool +): + """Reject a CONTEXT checkpoint at the id of a virtual child. + + A virtual child never checkpoints itself, so the checkpoint was written by + a non-virtual child in an earlier invocation. Returning its result or + re-raising its error would skip the virtual body, so replay must fail + before the checkpoint is used. + """ + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_from_operation( + Operation( + operation_id="op1", + operation_type=OperationType.CONTEXT, + status=status, + sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, + name="test_name", + context_details=ContextDetails( + replay_children=replay_children, + result=( + json.dumps("cached_result") + if status is OperationStatus.SUCCEEDED + else None + ), + error=( + ErrorObject( + message="Previous failure", + type="ValueError", + data=None, + stack_trace=None, + ) + if status is OperationStatus.FAILED + else None + ), + ), + ) + ) + ) + mock_callable = Mock(return_value="fresh_result") + + with pytest.raises( + NonDeterministicExecutionError, match="current child context is virtual" + ) as exc_info: + child_handler( + mock_callable, + mock_state, + OperationIdentifier( + "op1", + OperationSubType.RUN_IN_CHILD_CONTEXT, + None, + "test_name", + operation_type=OperationType.CONTEXT, + ), + ChildConfig(is_virtual=True), + ) + + assert exc_info.value.step_id == "op1" + mock_callable.assert_not_called() + mock_state.wrap_user_function.assert_not_called() + mock_state.create_checkpoint.assert_not_called() + mock_state.emit_child_context_end_hook.assert_not_called() + + # 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.