diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py index c0ada9a6..8ca5ef90 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py @@ -299,6 +299,12 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: ), is_first_invocation=not has_prior_operations, execution_input=input_event, + # Read the map through a callable rather than snapshotting it + # here: the invocation-end hook needs the state as of the end of + # the invocation, and neither hook pays for the conversion until + # a plugin actually reads it. + operations_provider=lambda: execution_state.operations, + updated_operation_ids=invocation_input.updated_operation_ids, ) # Thread 1: Run background checkpoint processing executor.submit(execution_state.checkpoint_batches_forever) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py index 96bfb373..33d36a53 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py @@ -9,7 +9,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum -from typing import Any, Callable, MutableMapping, cast +from typing import Any, Callable, Iterator, MutableMapping, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -107,6 +107,102 @@ def from_operation( ) +def _to_operation_info_map( + operations: Mapping[str, Operation], +) -> dict[str, OperationInfo]: + """Convert a map of checkpointed operations to the plugin ``OperationInfo`` view. + + ``is_replayed`` is left at its default ``False``: these entries describe the + stored state of an operation, not a replay event for it. Replay is signalled + through the dedicated operation hooks. + """ + return { + operation_id: OperationInfo.from_operation(operation) + for operation_id, operation in operations.items() + } + + +class _LazyOperationInfoMap(Mapping[str, OperationInfo]): + """Read-only ``operation id -> OperationInfo`` map built on first access. + + Invocation hook infos carry the execution's whole operation map. Converting + it eagerly would charge every invocation of an operation-heavy execution for + a view most plugins never read, so the conversion is deferred to the first + mapping operation and then cached. The underlying operations are snapshotted + before this map is handed out, so deferring the conversion does not move the + point in time the map describes. Behaves like a plain read-only ``dict`` + (iteration, ``len``, ``in``, ``get``, ``==`` against any mapping), and + materializes into an ordinary ``dict`` when deep-copied so the enclosing + info stays copyable and ``dataclasses.asdict()``-able. + """ + + __slots__ = ("_provider", "_resolved") + + def __init__(self, provider: Callable[[], dict[str, OperationInfo]] | None) -> None: + self._provider = provider + self._resolved: dict[str, OperationInfo] | None = None + + def _resolve(self) -> dict[str, OperationInfo]: + # Deliberately lock-free. Building the map is pure and idempotent, so a + # race can only duplicate the work, never corrupt the result, and the + # assignment below is atomic. A lock here would make the enclosing + # frozen info undeepcopyable, which breaks the plugins this view exists + # to serve: dataclasses.asdict() deep-copies non-dict fields, and a + # thread lock cannot be copied. + if self._resolved is None: + if self._provider is None: + self._resolved = {} + else: + try: + self._resolved = self._provider() + except Exception: + # A plugin-facing view must never break the execution. + logger.exception( + "Failed to build plugin operations map; using empty map" + ) + self._resolved = {} + return self._resolved + + def __getitem__(self, key: str) -> OperationInfo: + return self._resolve()[key] + + def __iter__(self) -> Iterator[str]: + return iter(self._resolve()) + + def __len__(self) -> int: + return len(self._resolve()) + + def __eq__(self, other: object) -> bool: + if isinstance(other, Mapping): + return self._resolve() == dict(other) + return NotImplemented + + def __ne__(self, other: object) -> bool: + result = self.__eq__(other) + if result is NotImplemented: + return result + return not result + + __hash__ = None # type: ignore[assignment] # mutable-by-materialization view + + def __deepcopy__(self, memo: dict[int, Any]) -> dict[str, OperationInfo]: + """Materialize into a plain ``dict`` when copied. + + Copying this view has no reason to preserve its laziness, and a plain + dict is what callers actually want: it makes + ``dataclasses.asdict(info)`` yield an ordinary mapping instead of an + opaque object, and keeps the provider closure out of the copy. + """ + resolved = { + key: copy.deepcopy(value, memo) for key, value in self._resolve().items() + } + memo[id(self)] = resolved + return resolved + + def __repr__(self) -> str: + return f"{type(self).__name__}({self._resolve()!r})" + + @dataclass(frozen=True) class OperationStartInfo(OperationInfo): pass @@ -190,11 +286,29 @@ class InvocationInfo: without it); ``durable_execution()`` always populates it with the deserialized input payload, which is ``{}`` when the payload is empty. """ + operations: Mapping[str, OperationInfo] = field(default_factory=dict, kw_only=True) + """Checkpointed operations for this execution, keyed by operation id. + + A point-in-time view of the execution's operation map: as observed at the + start of the invocation on ``on_invocation_start``, and as observed at the + end of the invocation on ``on_invocation_end``. Empty on the very first + invocation-start, before any operation has been checkpointed. + """ @dataclass(frozen=True) class InvocationStartInfo(InvocationInfo): - pass + updated_operations: Mapping[str, OperationInfo] = field( + default_factory=dict, kw_only=True + ) + """Operations updated externally while this execution was suspended. + + A wait timer that expired, a callback that was delivered, or a chained + invoke that completed between the previous invocation and this one. This is + the subset of :attr:`InvocationInfo.operations` named by the durable + invocation input's ``UpdatedOperationIds``, so it is empty on the first + invocation. + """ @dataclass(frozen=True) @@ -222,6 +336,7 @@ def from_durable_execution_invocation_output( cls, invocation_start_info: InvocationStartInfo, output: "DurableExecutionInvocationOutput", + operations: Mapping[str, OperationInfo] | None = None, ): return InvocationEndInfo( request_id=invocation_start_info.request_id, @@ -229,6 +344,13 @@ def from_durable_execution_invocation_output( is_first_invocation=invocation_start_info.is_first_invocation, execution_start_time=invocation_start_info.execution_start_time, execution_input=invocation_start_info.execution_input, + # Default to the start-of-invocation view when the caller has no + # fresher snapshot to offer. + operations=( + operations + if operations is not None + else invocation_start_info.operations + ), status=output.status, error=output.error, execution_result=output.result, @@ -319,6 +441,7 @@ def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._plugins = plugins or [] self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None + self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None @contextlib.contextmanager def run(self): @@ -331,6 +454,7 @@ def run(self): yield finally: self._invocation_status = None + self._operations_provider = None # Shut down the thread pool, waiting for pending tasks to complete. if self._executor: self._executor.shutdown(wait=True) @@ -371,6 +495,28 @@ def execute_plugins(self, info, sync): # this is called asynchronously, so plugins cannot manipulate thread local objects self._executor.submit(self._dispatch_plugin, plugin, info) + @staticmethod + def _snapshot_operation_infos( + operations_provider: Callable[[], Mapping[str, Operation]] | None, + ) -> Mapping[str, OperationInfo]: + """Capture the operation map now; defer the ``OperationInfo`` conversion. + + Copying the map is cheap, building an ``OperationInfo`` per operation is + not. Snapshotting eagerly pins the point in time the hook reports, so a + plugin that stashes the info and reads it later still sees the state as + of its hook; deferring the conversion keeps operation-heavy executions + from paying for a view no plugin reads. + """ + if operations_provider is None: + return _LazyOperationInfoMap(None) + try: + snapshot = dict(operations_provider()) + except Exception: + # A plugin-facing view must never break the execution. + logger.exception("Failed to snapshot operations for plugin hook") + return _LazyOperationInfoMap(None) + return _LazyOperationInfoMap(lambda: _to_operation_info_map(snapshot)) + def on_invocation_start( self, execution_arn: str, @@ -378,14 +524,40 @@ def on_invocation_start( execution_start_time: datetime.datetime | None, lambda_context: LambdaContext | None, execution_input: Any = None, + operations_provider: Callable[[], Mapping[str, Operation]] | None = None, + updated_operation_ids: Sequence[str] | None = None, ) -> None: + """Fire the invocation-start hook. + + Args: + execution_arn: ARN of the durable execution. + is_first_invocation: False when prior operations exist (a replay). + execution_start_time: Start timestamp of the execution operation. + lambda_context: Lambda context, for the request id. + execution_input: The deserialized execution input event. + operations_provider: Returns the current checkpointed operation map. + Called once here to snapshot it; the conversion to the plugin's + ``OperationInfo`` view is deferred to first access. + updated_operation_ids: Operation ids from the invocation input's + ``UpdatedOperationIds`` -- those updated while suspended. + """ aws_request_id = lambda_context.aws_request_id if lambda_context else None + self._operations_provider = operations_provider + operations = self._snapshot_operation_infos(operations_provider) self._invocation_status = InvocationStartInfo( execution_arn=execution_arn, request_id=aws_request_id, is_first_invocation=is_first_invocation, execution_start_time=execution_start_time, execution_input=self._snapshot_execution_input(execution_input), + operations=operations, + updated_operations=_LazyOperationInfoMap( + lambda: { + operation_id: operations[operation_id] + for operation_id in (updated_operation_ids or []) + if operation_id in operations + } + ), ) self.execute_plugins(self._invocation_status, sync=True) @@ -426,9 +598,13 @@ def on_invocation_end( # on_invocation_start not called, skip return + # Re-read the operation map so the end hook sees the state as of the end + # of this invocation, not the snapshot taken at its start. invocation_end_info = ( InvocationEndInfo.from_durable_execution_invocation_output( - self._invocation_status, output + self._invocation_status, + output, + operations=self._snapshot_operation_infos(self._operations_provider), ) ) self.execute_plugins(invocation_end_info, sync=True) @@ -607,10 +783,7 @@ def on_operation_update( operation.operation_id: OperationInfo.from_operation(operation) for operation in changed_operations }, - operations={ - operation_id: OperationInfo.from_operation(operation) - for operation_id, operation in operations.items() - }, + operations=_to_operation_info_map(operations), ), sync=True, ) diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 556dac0a..b64e191d 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -1,8 +1,9 @@ import datetime import logging import unittest -from dataclasses import fields -from unittest.mock import MagicMock +from copy import deepcopy +from dataclasses import asdict, fields +from unittest.mock import MagicMock, patch from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -203,6 +204,69 @@ def test_invocation_end_info_from_invocation_output_carries_input_and_result(sel self.assertEqual(end_info.status, InvocationStatus.SUCCEEDED) self.assertIsNone(end_info.error) + def test_invocation_info_operation_maps_default_to_empty(self): + """Both maps default to empty so existing constructor calls keep working.""" + self.assertEqual({}, INVOCATION_START_INFO.operations) + self.assertEqual({}, INVOCATION_START_INFO.updated_operations) + self.assertEqual({}, INVOCATION_END_INFO.operations) + + def test_invocation_start_info_carries_operation_maps(self): + info = InvocationStartInfo( + request_id="req-1", + execution_arn="arn:test", + is_first_invocation=False, + operations={"op-1": OPERATION_END_INFO, "op-2": OPERATION_START_INFO}, + updated_operations={"op-1": OPERATION_END_INFO}, + ) + + self.assertEqual( + {"op-1": OPERATION_END_INFO, "op-2": OPERATION_START_INFO}, + info.operations, + ) + self.assertEqual({"op-1": OPERATION_END_INFO}, info.updated_operations) + + def test_invocation_end_info_inherits_start_operations(self): + """The end factory reuses the start snapshot when given no override.""" + start = InvocationStartInfo( + request_id="req-1", + execution_arn="arn:test", + is_first_invocation=True, + operations={"op-1": OPERATION_END_INFO}, + updated_operations={"op-1": OPERATION_END_INFO}, + ) + output = DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + + end = InvocationEndInfo.from_durable_execution_invocation_output(start, output) + + self.assertEqual({"op-1": OPERATION_END_INFO}, end.operations) + # updated_operations is a start-hook-only surface. + self.assertFalse(hasattr(end, "updated_operations")) + + def test_invocation_end_info_accepts_fresher_operations(self): + """An explicit map wins, letting the end hook see end-of-invocation state.""" + start = InvocationStartInfo( + request_id="req-1", + execution_arn="arn:test", + is_first_invocation=True, + operations={"op-1": OPERATION_START_INFO}, + ) + output = DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + + end = InvocationEndInfo.from_durable_execution_invocation_output( + start, + output, + operations={"op-1": OPERATION_END_INFO, "op-2": OPERATION_START_INFO}, + ) + + self.assertEqual( + {"op-1": OPERATION_END_INFO, "op-2": OPERATION_START_INFO}, + end.operations, + ) + def test_user_function_start_info(self): self.assertEqual(USER_FUNCTION_START_INFO.operation_id, "op-1") self.assertEqual(USER_FUNCTION_START_INFO.operation_type, OperationType.STEP) @@ -579,6 +643,293 @@ def test_pending_fires_invocation_end(self): self.assertIn("invocation_end:req-1", self.plugin.calls) +class TestInvocationHookOperationMaps(unittest.TestCase): + """Tests for the operations / updated_operations maps on invocation hooks.""" + + def setUp(self): + self.captured: list[object] = [] + + class _CapturingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(_self, info): # noqa: N805 + self.captured.append(info) + + def on_invocation_end(_self, info): # noqa: N805 + self.captured.append(info) + + self.executor = PluginExecutor(plugins=[_CapturingPlugin()]) + + @staticmethod + def _operation(operation_id, status=OperationStatus.SUCCEEDED): + return Operation( + operation_id=operation_id, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=f"name-{operation_id}", + parent_id="root", + status=status, + start_timestamp=START_TS, + end_timestamp=END_TS, + step_details=StepDetails(attempt=1), + ) + + def _start(self, *, operations=None, updated_operation_ids=None): + self.executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=False, + operations_provider=(lambda: operations) + if operations is not None + else None, + updated_operation_ids=updated_operation_ids, + ) + + def test_start_info_converts_operations_to_operation_info(self): + operations = {"op-1": self._operation("op-1")} + + with self.executor.run(): + self._start(operations=operations) + + (start_info,) = self.captured + self.assertEqual(["op-1"], list(start_info.operations)) + converted = start_info.operations["op-1"] + self.assertIsInstance(converted, OperationInfo) + self.assertEqual("op-1", converted.operation_id) + self.assertEqual(OperationStatus.SUCCEEDED, converted.status) + self.assertEqual(START_TS, converted.start_time) + self.assertEqual(END_TS, converted.end_time) + self.assertEqual(1, converted.attempt) + # Stored state, not a replay event for the operation. + self.assertFalse(converted.is_replayed) + + def test_updated_operations_is_the_subset_named_by_updated_ids(self): + operations = { + "op-1": self._operation("op-1"), + "op-2": self._operation("op-2"), + } + + with self.executor.run(): + self._start(operations=operations, updated_operation_ids=["op-2"]) + + (start_info,) = self.captured + self.assertEqual(2, len(start_info.operations)) + self.assertEqual(["op-2"], list(start_info.updated_operations)) + # Same converted instance as in the full map, not a second conversion. + self.assertIs( + start_info.operations["op-2"], start_info.updated_operations["op-2"] + ) + + def test_updated_ids_absent_from_the_map_are_skipped(self): + operations = {"op-1": self._operation("op-1")} + + with self.executor.run(): + self._start( + operations=operations, updated_operation_ids=["op-1", "op-missing"] + ) + + (start_info,) = self.captured + self.assertEqual(["op-1"], list(start_info.updated_operations)) + + def test_first_invocation_has_empty_maps(self): + with self.executor.run(): + self._start(operations={}, updated_operation_ids=[]) + + (start_info,) = self.captured + self.assertEqual({}, start_info.operations) + self.assertEqual({}, start_info.updated_operations) + + def test_missing_provider_yields_empty_maps(self): + """Callers that supply no provider still get a usable, empty map.""" + with self.executor.run(): + self._start() + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + + start_info, end_info = self.captured + self.assertEqual({}, start_info.operations) + self.assertEqual({}, start_info.updated_operations) + self.assertEqual({}, end_info.operations) + + def test_operation_info_conversion_is_deferred_and_cached(self): + conversions: list[str] = [] + real_from_operation = OperationInfo.from_operation + + def counting(operation, **kwargs): + conversions.append(operation.operation_id) + return real_from_operation(operation, **kwargs) + + with ( + patch.object(OperationInfo, "from_operation", staticmethod(counting)), + self.executor.run(), + ): + self._start(operations={"op-1": self._operation("op-1")}) + # Nothing read the map inside the hook, so nothing was converted. + self.assertEqual([], conversions) + + (start_info,) = self.captured + self.assertEqual(1, len(start_info.operations)) + self.assertEqual(["op-1"], conversions) + # Repeated reads reuse the cached conversion. + self.assertEqual(["op-1"], list(start_info.operations)) + self.assertIn("op-1", start_info.operations) + self.assertEqual(["op-1"], conversions) + + def test_end_info_reflects_operations_added_during_the_invocation(self): + operations = {"op-1": self._operation("op-1")} + + with self.executor.run(): + self.executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + operations_provider=lambda: operations, + ) + operations["op-2"] = self._operation("op-2") + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + + start_info, end_info = self.captured + self.assertEqual(["op-1", "op-2"], sorted(end_info.operations)) + # The start map is snapshotted at its hook, so a later checkpoint does + # not retroactively change what the start hook reported. + self.assertEqual(["op-1"], sorted(start_info.operations)) + + def test_failing_provider_degrades_to_an_empty_map(self): + def provider(): + raise RuntimeError("boom") + + with ( + self.executor.run(), + self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level="ERROR" + ) as logs, + ): + self.executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + operations_provider=provider, + updated_operation_ids=["op-1"], + ) + + self.assertTrue(any("Failed to snapshot operations" in m for m in logs.output)) + (start_info,) = self.captured + self.assertEqual(0, len(start_info.operations)) + self.assertEqual(0, len(start_info.updated_operations)) + + def test_provider_is_released_when_the_run_scope_exits(self): + with self.executor.run(): + self._start(operations={"op-1": self._operation("op-1")}) + self.assertIsNotNone(self.executor._operations_provider) # noqa: SLF001 + + self.assertIsNone(self.executor._operations_provider) # noqa: SLF001 + + +class TestInvocationInfoCopySafety(unittest.TestCase): + """Invocation infos carrying the lazy operation map must stay copyable. + + ``dataclasses.asdict()`` deep-copies any field that is not a dict/list/tuple + or dataclass, so anything uncopyable embedded in the lazy map would make + every plugin that serializes an info raise -- and because plugin exceptions + are swallowed, the hook would be silently dropped. + """ + + def setUp(self): + self.operation = Operation( + operation_id="op-1", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name="name-op-1", + parent_id="root", + status=OperationStatus.SUCCEEDED, + start_timestamp=START_TS, + end_timestamp=END_TS, + step_details=StepDetails(attempt=1), + ) + self.captured: list[object] = [] + + class _CapturingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(_self, info): # noqa: N805 + self.captured.append(info) + + def on_invocation_end(_self, info): # noqa: N805 + self.captured.append(info) + + self.executor = PluginExecutor(plugins=[_CapturingPlugin()]) + + def _fire_hooks(self): + with self.executor.run(): + self.executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=False, + operations_provider=lambda: {"op-1": self.operation}, + updated_operation_ids=["op-1"], + ) + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + return self.captured + + def test_asdict_on_invocation_start_info(self): + start_info, _ = self._fire_hooks() + + as_dict = asdict(start_info) + + # Both lazy maps materialize as ordinary dicts. + self.assertIsInstance(as_dict["operations"], dict) + self.assertIsInstance(as_dict["updated_operations"], dict) + self.assertEqual(["op-1"], list(as_dict["operations"])) + self.assertEqual(["op-1"], list(as_dict["updated_operations"])) + + def test_asdict_on_invocation_end_info(self): + _, end_info = self._fire_hooks() + + as_dict = asdict(end_info) + + self.assertIsInstance(as_dict["operations"], dict) + self.assertEqual(["op-1"], list(as_dict["operations"])) + + def test_deepcopy_on_invocation_infos(self): + start_info, end_info = self._fire_hooks() + + for info in (start_info, end_info): + copied = deepcopy(info) + self.assertIsInstance(copied.operations, dict) + self.assertEqual(["op-1"], list(copied.operations)) + + def test_deepcopy_is_independent_of_the_original(self): + start_info, _ = self._fire_hooks() + + copied = deepcopy(start_info) + copied.operations["op-2"] = OPERATION_END_INFO + + self.assertEqual(["op-1"], list(start_info.operations)) + + def test_unresolved_map_is_still_copyable(self): + """Copying must not require the map to have been read first.""" + start_info, _ = self._fire_hooks() + fresh = InvocationStartInfo( + request_id=start_info.request_id, + execution_arn=start_info.execution_arn, + is_first_invocation=start_info.is_first_invocation, + operations=start_info.operations, + ) + + self.assertIsInstance(asdict(fresh)["operations"], dict) + + class TestPluginExecutorOnOperationAction(unittest.TestCase): """Tests for PluginExecutor.on_operation_action."""