From f6ccf790dd015b87b29f14800fdbecd9be285c5f Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:32:18 +0000 Subject: [PATCH 1/3] fix: enforce cooperative invocation deadlines Use an invocation-local monotonic budget for SDK admission and waits. Reject newly registered descendants of completed contexts before user code or checkpoint dispatch. Preserve acknowledged terminal outcomes without issuing metadata reads after the cutoff. Keep running threads confined and document the unresolved non-cooperative cleanup boundary in #741. --- .../e2e/test_invocation_wait_resume_int.py | 1 + .../README.md | 37 ++ .../concurrency/executor.py | 6 +- .../deadline.py | 49 ++ .../exceptions.py | 4 + .../execution.py | 13 +- .../operation/base.py | 1 + .../operation/step.py | 5 + .../aws_durable_execution_sdk_python/state.py | 110 +++- .../tests/concurrency_test.py | 37 ++ .../tests/context_test.py | 1 + .../tests/deadline_test.py | 103 ++++ .../tests/e2e/checkpoint_response_int_test.py | 11 + .../e2e/custom_serdes_roundtrip_int_test.py | 1 + .../tests/e2e/error_hierarchy_int_test.py | 1 + .../tests/e2e/execution_int_test.py | 6 + .../tests/e2e/filesystem_serdes_int_test.py | 1 + .../tests/e2e/invocation_deadline_int_test.py | 515 ++++++++++++++++++ .../plugin_invocation_operations_int_test.py | 1 + .../e2e/plugin_invocation_payload_int_test.py | 1 + ...plugin_user_function_lifecycle_int_test.py | 1 + .../tests/execution_test.py | 34 ++ .../tests/handler_worker_context_test.py | 1 + .../tests/operation/base_test.py | 4 + .../tests/operation/map_test.py | 79 +++ .../tests/operation/parallel_test.py | 78 +++ 26 files changed, 1089 insertions(+), 12 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/deadline.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/deadline_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/invocation_deadline_int_test.py diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index fa7036920..f0c5bc520 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -52,6 +52,7 @@ def _lambda_context() -> Mock: context = Mock() + context.get_remaining_time_in_millis.return_value = 900000 context.aws_request_id = "test-request-id" context.client_context = None context.identity = None diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index 4b2463e52..19f60ea86 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -145,6 +145,43 @@ core alone with OTel 1.0 isolates worker bindings, but does not add the newer Invocation-view fallback to that older plugin; upgrade both packages for it. Provider API version 1 and the independent registration API are unchanged. +### Invocation deadlines and Lambda Managed Instances + +The SDK captures `LambdaContext.get_remaining_time_in_millis()` at invocation +entry and uses one monotonic budget, reserving 100 ms before the platform +deadline for cooperative cleanup. An already-expired invocation starts no user +handler. Once the budget is exhausted, new user operations and checkpoint/state +requests are rejected with `InvocationDeadlineExceeded`, a retryable +`InvocationError`. Step retry strategies do not turn this interruption into a +persisted step failure. Completed checkpoints retain their stored results and +errors for replay. + +Synchronous checkpoint waits, scheduled refresh waits, checkpoint queue reads, +and map/parallel coordinator waits share that budget. In particular, a losing +branch parked in an SDK refresh wait after early parallel completion unwinds at +the cutoff, allowing invocation cleanup to join it. A newly created descendant +of a completed context is rejected before its body, START hook, or checkpoint +runs, even if it did not exist when its ancestor completed. + +This is cooperative interruption, **not a guarantee of bounded invocation +shutdown or LMI worker reclamation**. Python cannot safely terminate arbitrary +running threads. The SDK still joins invocation-owned workers before returning: +user code or I/O that does not cooperate, an already-running synchronous service +client call (including its retries), or a blocked plugin hook can delay those +joins indefinitely. The shared service client is not reconfigured or closed to +interrupt one invocation. Give application I/O explicit finite timeouts and keep +plugin hooks bounded. A runtime-supported worker-retirement strategy is still +required for non-cooperative work; [issue #741](https://github.com/aws/aws-durable-execution-sdk-python/issues/741) +tracks that remaining work, with real-LMI coverage in [test PR #743](https://github.com/aws/aws-durable-execution-sdk-python/pull/743). + +An expired invocation does not report an unacknowledged result as `SUCCEEDED` or +`PENDING`. If the service has already acknowledged this invocation's terminal +execution success/failure checkpoint, that outcome remains authoritative while +local delivery hooks finish. At the cutoff, further response metadata pages are +skipped for that accepted terminal checkpoint; inline and already-read operations +are retained. Already-running side effects cannot be undone; +interrupted work continues to require the documented retry/idempotency strategy. + ## 🚀 Quick Start Install the execution SDK: diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py index fd049d3bf..c998081c7 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py @@ -292,6 +292,7 @@ def execute( execution_state.register_branch_pool(pool) def submit(branch: Branch[CallableType, ResultType]) -> None: + execution_state.invocation_deadline.check() branch.start() if execution_state._plugin_executor._plugins: # noqa: SLF001 # Every admission/resume gets its own Context. Branch bindings @@ -323,6 +324,7 @@ def submit(branch: Branch[CallableType, ResultType]) -> None: # and the reason, so user code runs at most once per change. decision: CompletionDecision | None = None while True: + execution_state.invocation_deadline.check() if needs_snapshot_rebuild: items_snapshot = ( self._build_items_snapshot() @@ -386,7 +388,9 @@ def submit(branch: Branch[CallableType, ResultType]) -> None: if resume_times: timeout = max(resume_times[0] - time.time(), 0) try: - event: BranchEvent[ResultType] = events.get(timeout=timeout) + event: BranchEvent[ResultType] = events.get( + timeout=execution_state.invocation_deadline.remaining(timeout) + ) except queue.Empty: # A timed resume came due while branches were running. continue diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/deadline.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/deadline.py new file mode 100644 index 000000000..1b6ebd948 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/deadline.py @@ -0,0 +1,49 @@ +"""One monotonic time budget for cooperative invocation interruption.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from aws_durable_execution_sdk_python.exceptions import InvocationDeadlineExceeded + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python.threading import CompletionEvent + from aws_durable_execution_sdk_python.types import LambdaContext + +# Stop admitting work before the platform deadline to leave time for cooperative +# cleanup. This is not a bound on arbitrary user code, clients, or plugin hooks. +_CLEANUP_MARGIN_SECONDS = 0.1 + + +@dataclass(frozen=True) +class InvocationDeadline: + """Internal invocation-local deadline; None keeps standalone state unbounded.""" + + expires_at: float | None = None + + @classmethod + def from_lambda_context(cls, context: LambdaContext) -> InvocationDeadline: + # Capture once: wall-clock adjustments cannot extend this invocation. + started_at = time.monotonic() + remaining_seconds = context.get_remaining_time_in_millis() / 1000 + return cls(started_at + max(0.0, remaining_seconds - _CLEANUP_MARGIN_SECONDS)) + + def remaining(self, timeout: float | None = None) -> float | None: + """Clamp an existing wait to the same budget, without resetting it.""" + if self.expires_at is None: + return timeout + remaining = max(0.0, self.expires_at - time.monotonic()) + return remaining if timeout is None else min(timeout, remaining) + + def check(self) -> None: + if self.remaining() == 0: + raise InvocationDeadlineExceeded("Lambda invocation time budget exhausted") + + def wait(self, event: CompletionEvent, timeout: float | None = None) -> bool: + """Do not turn an acknowledged checkpoint into an error after its wait.""" + completed = event.wait(self.remaining(timeout)) + if not completed: + self.check() + return completed diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py index ea3f163f2..05bf077e4 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py @@ -104,6 +104,10 @@ def is_retryable(self) -> bool: return True +class InvocationDeadlineExceeded(InvocationError): + """The invocation must retry after its cooperative time budget expires.""" + + class DurableApiErrorCategory(Enum): INVOCATION = "INVOCATION" EXECUTION = "EXECUTION" 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 e59820007..5a038ab96 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 @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.deadline import InvocationDeadline from aws_durable_execution_sdk_python.exceptions import ( BackgroundThreadError, BotoClientError, @@ -193,6 +194,8 @@ def durable_execution( plugin_executor = PluginExecutor(configured_plugins) def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: + deadline = InvocationDeadline.from_lambda_context(context) + deadline.check() invocation_input: DurableExecutionInvocationInput service_client: DurableServiceClient @@ -228,6 +231,7 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: service_client=service_client, plugin_executor=plugin_executor, updated_operation_ids=invocation_input.updated_operation_ids, + invocation_deadline=deadline, ) try: @@ -286,6 +290,7 @@ def invoke_handler() -> MutableMapping[str, Any]: ) try: # Background checkpointing errors will propagate through CompletionEvent.wait() as BackgroundThreadError + deadline.check() result = func(input_event, durable_context) # done with userland @@ -445,6 +450,7 @@ def wait_for_checkpoint() -> None: ) execution_operation = execution_state.get_execution_operation() + deadline.check() # execute the plugins plugin_executor.on_invocation_start( execution_arn=invocation_input.durable_execution_arn, @@ -470,6 +476,7 @@ def wait_for_checkpoint() -> None: output = plugin_executor._run_in_invocation_context( invoke_handler ) + execution_state.check_invocation_result() plugin_executor.on_invocation_end( DurableExecutionInvocationOutput.from_dict(output) ) @@ -492,7 +499,11 @@ def wait_for_checkpoint() -> None: else: # Preserve the original fresh-worker context without plugins. invocation_future = executor.submit(invoke_worker) - return invocation_future.result() + result = invocation_future.result() + # Cleanup hooks and all owned workers have returned. A hook can consume + # the remaining budget; do not call End again when this check fails. + execution_state.check_invocation_result() + return result return wrapper diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/base.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/base.py index 126f63d55..871568550 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/base.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/base.py @@ -179,6 +179,7 @@ def process(self) -> T: InvalidStateError: If the check result is in an invalid state May raise operation-specific errors from check_result_status() or execute() """ + self.state.check_operation_active(self.operation_identifier) # Check 1: Entry (handles replay and existing checkpoints) result = self.check_result_status() 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 a9b9151e3..025b25aa8 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 @@ -12,6 +12,7 @@ from aws_durable_execution_sdk_python.exceptions import ( ExecutionError, InvalidStateError, + InvocationDeadlineExceeded, RetryableSerDesError, SerDesError, StepError, @@ -266,6 +267,10 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T: self.operation_identifier.name, ) return return_value + except InvocationDeadlineExceeded: + # Deadline interruption belongs to the invocation, not the step + # retry policy or persisted operation failure. + raise except RetryableSerDesError: # Transient serdes failure: fail the invocation for backend retry, # bypassing the step retry strategy. For AT_MOST_ONCE steps the diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py index de396b11e..ff2b9a1fd 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py @@ -16,12 +16,14 @@ from threading import Lock from typing import TYPE_CHECKING, Callable, NoReturn +from aws_durable_execution_sdk_python.deadline import InvocationDeadline from aws_durable_execution_sdk_python.exceptions import ( BackgroundThreadError, CheckpointError, DurableExecutionsError, DurableOperationError, GetExecutionStateError, + InvocationDeadlineExceeded, OrphanedChildException, ) from aws_durable_execution_sdk_python.identifier import OperationIdentifier @@ -112,9 +114,10 @@ class ScheduledRefresh: earliest_check_time: float completion_event: CompletionEvent cancelled: threading.Event = field(default_factory=threading.Event) + invocation_deadline: InvocationDeadline = field(default_factory=InvocationDeadline) def wait(self, timeout: float | None = None) -> bool: - return self.completion_event.wait(timeout) + return self.invocation_deadline.wait(self.completion_event, timeout) def is_set(self) -> bool: return self.completion_event.is_set() @@ -369,8 +372,11 @@ def __init__( plugin_executor: PluginExecutor, batcher_config: CheckpointBatcherConfig | None = None, updated_operation_ids: list[str] | None = None, + invocation_deadline: InvocationDeadline | None = None, ): self.durable_execution_arn: str = durable_execution_arn + self.invocation_deadline = invocation_deadline or InvocationDeadline() + self._terminal_execution_acknowledged = threading.Event() self._current_checkpoint_token: str = initial_checkpoint_token self._operations: dict[str, Operation] = dict(operations) self._service_client: DurableServiceClient = service_client @@ -453,6 +459,8 @@ def fetch_paginated_operations( initial_operations: list[Operation], checkpoint_token: str, next_marker: str | None, + *, + finish_terminal_checkpoint: bool = False, ) -> list[Operation]: """Add initial operations and fetch all paginated operations from the Durable Functions API. This method is thread_safe. @@ -476,6 +484,15 @@ def fetch_paginated_operations( ) try: while next_marker: + try: + self.invocation_deadline.check() + except InvocationDeadlineExceeded: + if not finish_terminal_checkpoint: + raise + # The terminal execution outcome is authoritative. Preserve + # inline/already-read metadata without starting another I/O + # request after expiry merely to finish its metadata pages. + break output: StateOutput = self._service_client.get_execution_state( durable_execution_arn=self.durable_execution_arn, checkpoint_token=checkpoint_token, @@ -642,12 +659,37 @@ def _reject_if_execution_completed( operation_id=operation_id, ) - def _reject_if_parent_done(self, operation_update: OperationUpdate) -> None: + def check_operation_active(self, identifier: OperationIdentifier) -> None: + """Reject new work under an expired invocation or completed ancestor.""" + self.invocation_deadline.check() + with self._parent_done_lock: + self._reject_if_parent_done(identifier) + if identifier.parent_id: + self._parent_to_children.setdefault(identifier.parent_id, set()).add( + identifier.operation_id + ) + + def check_invocation_result(self) -> None: + # A terminal execution checkpoint is already authoritative. Preserve it + # even if local cleanup took the remaining invocation budget. + if not self._terminal_execution_acknowledged.is_set(): + self.invocation_deadline.check() + + def _reject_if_parent_done( + self, operation_update: OperationUpdate | OperationIdentifier + ) -> None: """Raise OrphanedChildException when the operation is orphaned. Must be called while holding _parent_done_lock so the check can be linearized with parent completion and, at the final call site, enqueue. """ + if ( + operation_update.parent_id in self._parent_done + or operation_update.parent_id in self._terminal_contexts + ): + # Inherit cancellation even when this child did not exist when its + # ancestor completed. Registration and completion share this lock. + self._parent_done.add(operation_update.operation_id) if operation_update.operation_id not in self._parent_done: return logger.debug( @@ -731,6 +773,7 @@ def create_checkpoint( # Empty checkpoint (async for performance) execution_state.create_checkpoint(is_sync=False) """ + self.invocation_deadline.check() # if this is CONTEXT complete, mark incomplete descendants as orphans so the children can't complete after the parent if operation_update is not None: # Use single lock to coordinate completion and checkpoint validation @@ -798,6 +841,7 @@ def create_checkpoint( # _completion_lock and hold both through queue insertion so parent completion # cannot mark the operation orphaned between the final check and enqueue. with self._completion_lock: + self.invocation_deadline.check() if self._checkpointing_failed.is_set(): # Raises the stored BackgroundThreadError. self._checkpointing_failed.wait() @@ -818,7 +862,21 @@ def create_checkpoint( raise DurableExecutionsError(msg) # Wait for completion - will raise BackgroundThreadError if background thread fails - completion_event.wait() + try: + self.invocation_deadline.wait(completion_event) + except InvocationDeadlineExceeded: + if ( + operation_update is not None + and operation_update.operation_type is OperationType.EXECUTION + and operation_update.action + in {OperationAction.SUCCEED, OperationAction.FAIL} + and self._terminal_execution_acknowledged.is_set() + ): + # The service already accepted this terminal outcome. Finish + # its local processing rather than report a retry after it. + completion_event.wait() + else: + raise else: logger.debug("Enqueued checkpoint operation for asynchronous processing") @@ -831,8 +889,13 @@ def schedule_refresh(self, earliest_check_time: float) -> ScheduledRefresh: caller waits on the handle when it needs the refreshed operations, and cancels it if it stops needing them. """ - refresh = ScheduledRefresh(earliest_check_time, CompletionEvent()) + refresh = ScheduledRefresh( + earliest_check_time, + CompletionEvent(), + invocation_deadline=self.invocation_deadline, + ) with self._completion_lock: + self.invocation_deadline.check() if self._checkpointing_failed.is_set(): self._checkpointing_failed.wait() self._reject_if_execution_completed(None) @@ -994,7 +1057,7 @@ def checkpoint_batches_forever(self) -> None: # Collect operations into a batch batch: list[QueuedOperation] = self._collect_checkpoint_batch() - if batch: + if batch or self.invocation_deadline.remaining() == 0: # Extract OperationUpdates, excluding empty checkpoints from API call updates: list[OperationUpdate] = [] empty_count = 0 @@ -1013,6 +1076,9 @@ def checkpoint_batches_forever(self) -> None: ) try: + # Do not dispatch new checkpoint I/O after expiry. An I/O + # call already in progress must still be joined. + self.invocation_deadline.check() # Make API call with batched operations output: CheckpointOutput = self._service_client.checkpoint( durable_execution_arn=self.durable_execution_arn, @@ -1040,6 +1106,14 @@ def checkpoint_batches_forever(self) -> None: "checkpoint." ) + if any( + update.operation_type is OperationType.EXECUTION + and update.action + in {OperationAction.SUCCEED, OperationAction.FAIL} + for update in updates + ): + self._terminal_execution_acknowledged.set() + # Each read of self.operations copies the whole map, and only # plugins use these copies. So take the snapshots, and call # the hook, only when a plugin is registered. @@ -1059,6 +1133,7 @@ def checkpoint_batches_forever(self) -> None: output.new_execution_state.operations, current_checkpoint_token, fetch_marker, + finish_terminal_checkpoint=self._terminal_execution_acknowledged.is_set(), ) if has_plugins: self._plugin_executor.on_operation_update( @@ -1189,9 +1264,10 @@ def stop_checkpointing(self) -> None: """Signal background thread to stop checkpointing. This method sets the checkpointing stopped event, which signals the background - thread to exit. Any remaining async checkpoints in the queue are non-essential - (observability only) and will be abandoned. All critical synchronous checkpoints - will have already completed before this is called. + thread to exit. On ordinary completion, all required synchronous checkpoints + have completed. At deadline expiry their callers instead unwind with an + invocation interruption; an in-flight service call is still joined by the + wrapper. Unsent async checkpoints are abandoned. """ logger.debug("Signaling background thread to stop checkpointing") self._checkpointing_stopped.set() @@ -1205,6 +1281,7 @@ def register_branch_pool(self, pool: ThreadPoolExecutor) -> None: OrphanedChildException at their next checkpoint attempt) before the invocation returns. """ + self.invocation_deadline.check() with self._branch_pools_lock: self._branch_pools.append(pool) @@ -1251,7 +1328,10 @@ def _collect_checkpoint_batch(self) -> list[QueuedOperation]: while not batch.is_full() and not self._checkpointing_stopped.is_set(): now = time.time() self._add_due_refreshes(batch, now) - timeout = self._read_timeout(batch, now, deadline) + timeout = self.invocation_deadline.remaining( + self._read_timeout(batch, now, deadline) + ) + assert timeout is not None if timeout <= 0: break try: @@ -1304,8 +1384,11 @@ def _drain_overflow(self, batch: _BatchAccumulator) -> None: def _wait_for_first_operation(self, batch: _BatchAccumulator) -> bool: """Block until the batch has an operation. Returns False at shutdown.""" while not self._checkpointing_stopped.is_set(): + if self.invocation_deadline.remaining() == 0: + return False now = time.time() - timeout = _STOP_SIGNAL_POLL_SECONDS + timeout = self.invocation_deadline.remaining(_STOP_SIGNAL_POLL_SECONDS) + assert timeout is not None until_refresh = self._seconds_until_next_refresh(now) if until_refresh is not None: timeout = max(0.0, min(timeout, until_refresh)) @@ -1366,6 +1449,11 @@ def close(self): orphaned branch) and unwind. Stopping the batcher first would deadlock that branch until the Lambda timeout. + SDK-owned waits share the invocation deadline and unwind on expiry, so + they cannot pin this join indefinitely. Running user code, service I/O, + and plugin hooks must still finish cooperatively; Python cannot safely + terminate their threads and this method does not detach them. + Drains registrations until no pools remain: a branch joined in one batch can start a nested map or parallel operation, whose coordinator registers a new pool mid-join. The lock is never held @@ -1390,12 +1478,14 @@ def wrap_user_function( ): @functools.wraps(user_function) def wrapper(*args, **kwargs): + self.check_operation_active(operation_identifier) start_info = self._plugin_executor.on_user_function_start( operation_identifier, is_replay_children, attempt ) outcome = UserFunctionOutcome.INCOMPLETE error = None try: + self.check_operation_active(operation_identifier) result = user_function(*args, **kwargs) except Exception as exception: outcome = UserFunctionOutcome.FAILED 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 20adf7ac0..45011cea1 100644 --- a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py @@ -14,6 +14,8 @@ import pytest +from aws_durable_execution_sdk_python.deadline import InvocationDeadline + from aws_durable_execution_sdk_python.concurrency.executor import ( ConcurrentExecutor, ) @@ -697,6 +699,7 @@ def execute_item(self, child_context, executable): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() # Mock ChildConfig from the config module @@ -753,6 +756,7 @@ def failure_callable(): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() executor_context = Mock() @@ -1009,6 +1013,7 @@ def test_flat_replay_rejects_failed_nested_container_checkpoint(replay_path): executor_context = Mock() executor_context._parent_id = "parallel-op" # noqa: SLF001 execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.get_checkpoint_result.return_value = ( CheckpointedResult.create_from_operation( Operation( @@ -1060,6 +1065,7 @@ def test_replay_started_branch_validates_existing_checkpoint( executor_context = Mock() executor_context._parent_id = "parallel-op" # noqa: SLF001 execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.get_checkpoint_result.return_value = ( CheckpointedResult.create_from_operation( Operation( @@ -1102,6 +1108,7 @@ def test_replay_started_nested_branch_allows_missing_checkpoint(): executor_context = Mock() executor_context._parent_id = "parallel-op" # noqa: SLF001 execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.get_checkpoint_result.return_value = ( CheckpointedResult.create_not_found() ) @@ -1218,6 +1225,7 @@ def failure_callable(): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() def mock_run_in_child_context(func, name, config): @@ -1256,6 +1264,7 @@ def execute_item(self, child_context, executable): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() executor_context = Mock() @@ -1305,6 +1314,7 @@ def execute_item(self, child_context, executable): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() executor_context = Mock() @@ -1353,6 +1363,7 @@ def execute_item(self, child_context, executable): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() executor_context = Mock() @@ -1413,6 +1424,7 @@ def execute_item(self, child_context, executable): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() # The coordinator requests the refresh when task 1 suspends and waits on # the returned event when task 1 is due. Fail the wait. The branches' own @@ -1496,6 +1508,7 @@ def execute_item(self, child_context, executable): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() executor_context = Mock() @@ -1548,6 +1561,7 @@ def failure_callable(): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() executor_context = Mock() @@ -1597,6 +1611,7 @@ def make_executor(): # First run: execute() runs the branch and builds the result live. execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() executor_context = Mock() executor_context._create_step_id_for_logical_step = lambda *args: "1" @@ -1655,6 +1670,7 @@ def failure_callable(): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() executor_context = Mock() @@ -1699,6 +1715,7 @@ def success_callable(): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() # wrap_user_function must return the real function so the branch result is @@ -1722,6 +1739,7 @@ def success_callable(): def _serdes_branch_test_context() -> tuple[Mock, Mock]: """Build (execution_state, executor_context) mocks for a single-branch run.""" execution_state: Mock = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() child_context: Mock = Mock() child_context.is_replaying.return_value = False @@ -1834,6 +1852,7 @@ def execute_item(self, child_context, executable): checkpoint.is_succeeded.return_value = False checkpoint.is_failed.return_value = False execution_state: Mock = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.get_checkpoint_result = Mock(return_value=checkpoint) executor_context: Mock = Mock() @@ -1866,6 +1885,7 @@ def test_replay_flat_branch_nondeterminism_escapes_batch(): checkpoint.is_succeeded.return_value = False checkpoint.is_failed.return_value = False execution_state: Mock = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.get_checkpoint_result.return_value = checkpoint executor_context: Mock = Mock() @@ -1905,6 +1925,7 @@ def suspend_callable(): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() executor_context = Mock() @@ -2089,6 +2110,7 @@ def patched_child_handler( return result execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() completion_config = CompletionConfig(min_successful=num_executables) @@ -2404,6 +2426,7 @@ def execute_item(self, child_context, executable): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() executor_context = Mock() executor_context._create_step_id_for_logical_step = lambda *args: "1" @@ -2445,6 +2468,7 @@ def task_func(ctx, item, idx, items): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() executor_context = Mock() executor_context._create_step_id_for_logical_step = lambda *args: "1" @@ -2504,6 +2528,7 @@ def slow_branch(): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() executor_context = Mock() executor_context._create_step_id_for_logical_step = lambda idx: f"step_{idx}" @@ -2579,6 +2604,7 @@ def slow_branch(): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() execution_state.wrap_user_function = lambda func, *args, **kwargs: func executor_context = Mock() @@ -2641,6 +2667,7 @@ def slow_func(): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() execution_state.wrap_user_function = lambda func, *args, **kwargs: func executor_context = Mock() @@ -2836,6 +2863,7 @@ def execute_item(self, child_context, executable): return executable.func(child_context) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() execution_state.wrap_user_function = lambda func, *args, **kwargs: func @@ -2893,6 +2921,7 @@ def execute_item(self, child_context, executable): return executable.func(child_context) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() execution_state.wrap_user_function = lambda func, *args, **kwargs: func @@ -2968,6 +2997,7 @@ def execute_item(self, child_context, executable): def make_run(): execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() mock_checkpoint = Mock() @@ -3342,6 +3372,7 @@ def predicate(s: CompletionStatus) -> CompletionDecision: def _make_executor_mocks(): """Build the execution_state / executor_context mock pair used by executor tests.""" execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() execution_state.wrap_user_function = lambda func, *args, **kwargs: func # Default: no branch has a prior checkpoint (fresh run). @@ -3545,6 +3576,7 @@ def execute_item(self, child_context, executable): operation_id_namespace=_StubNamespace(), ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() executor_context = Mock() executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() @@ -3611,6 +3643,7 @@ def execute_item(self, child_context, executable): operation_id_namespace=_StubNamespace(), ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() executor_context = Mock() executor_context._create_step_id_for_logical_step = lambda *args: "1" child_context = Mock() @@ -3657,6 +3690,7 @@ def execute_item(self, child_context, executable): ) execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() refresh = Mock() def schedule_refresh(_resume_at): @@ -3910,6 +3944,7 @@ def test_summary_index_fields_records_smaller_side(): def _make_replay_mocks(branch_checkpoints: dict): """Mocks for replay: op id 'op_{index}', checkpoints per index.""" execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.durable_execution_arn = ( "arn:aws:durable:us-east-1:123456789012:execution/test" ) @@ -4631,6 +4666,7 @@ def test_operation_id_namespace_derivation_is_stable(): def _fresh_execution_state(): """Mock execution_state for fresh runs (no prior checkpoints).""" execution_state = Mock() + execution_state.invocation_deadline = InvocationDeadline() execution_state.create_checkpoint = Mock() _absent = Mock() _absent.is_succeeded.return_value = False @@ -5308,6 +5344,7 @@ def _execute_item_in_child_context( operation_id_namespace=_StubNamespace(), ) state = Mock(spec=ExecutionState) + state.invocation_deadline = InvocationDeadline() state._plugin_executor = PluginExecutor([DurableInstrumentationPlugin()]) token = binding.set("coordinator") try: 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 e3cc07acd..7acdbc34a 100644 --- a/packages/aws-durable-execution-sdk-python/tests/context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/context_test.py @@ -2345,6 +2345,7 @@ def test_from_lambda_context_creates_execution_context(): mock_state = Mock(spec=ExecutionState) mock_state.durable_execution_arn = expected_arn mock_lambda_context = Mock() + mock_lambda_context.get_remaining_time_in_millis.return_value = 900000 context = DurableContext.from_lambda_context( state=mock_state, lambda_context=mock_lambda_context diff --git a/packages/aws-durable-execution-sdk-python/tests/deadline_test.py b/packages/aws-durable-execution-sdk-python/tests/deadline_test.py new file mode 100644 index 000000000..8d01d7901 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/deadline_test.py @@ -0,0 +1,103 @@ +"""Invocation time budgets preserve acknowledged outcomes and error identity.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from aws_durable_execution_sdk_python import deadline as deadline_module +from aws_durable_execution_sdk_python.deadline import InvocationDeadline +from aws_durable_execution_sdk_python.exceptions import InvocationDeadlineExceeded +from aws_durable_execution_sdk_python.threading import CompletionEvent +from aws_durable_execution_sdk_python.state import ExecutionState +from aws_durable_execution_sdk_python.plugin import PluginExecutor +from aws_durable_execution_sdk_python.lambda_service import StateOutput + + +def test_budget_is_captured_once_and_shared_by_waits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + now = [10.0] + monkeypatch.setattr( + deadline_module, "time", SimpleNamespace(monotonic=lambda: now[0]) + ) + context = Mock() + context.get_remaining_time_in_millis.return_value = 2100 + deadline = InvocationDeadline.from_lambda_context(context) + assert deadline.remaining() == 2.0 + now[0] = 11.5 + assert deadline.remaining(10) == 0.5 + assert deadline.remaining(0.1) == 0.1 + context.get_remaining_time_in_millis.assert_called_once() + now[0] = 12.0 + with pytest.raises(InvocationDeadlineExceeded) as caught: + deadline.check() + assert caught.value.is_retryable() + assert deadline.remaining() == 0 + + +def test_standalone_state_has_no_implicit_invocation_deadline() -> None: + deadline = InvocationDeadline() + assert deadline.remaining() is None + assert deadline.remaining(0.5) == 0.5 + deadline.check() + assert not deadline.wait(CompletionEvent(), 0) + + +def test_acknowledged_event_keeps_success_after_deadline() -> None: + completed = CompletionEvent() + completed.set() + assert InvocationDeadline(0).wait(completed) + + +def test_event_preserves_original_error_after_deadline() -> None: + failure = ValueError("original failure") + completed = CompletionEvent() + completed.set(failure) + with pytest.raises(ValueError) as caught: + InvocationDeadline(0).wait(completed) + assert caught.value is failure + + +def test_unacknowledged_event_expires_without_reporting_success() -> None: + with pytest.raises(InvocationDeadlineExceeded): + InvocationDeadline(0).wait(CompletionEvent()) + + +def test_shorter_caller_timeout_does_not_expire_the_invocation() -> None: + assert not InvocationDeadline(float("inf")).wait(CompletionEvent(), 0) + + +@pytest.mark.parametrize("terminal", [False, True]) +def test_page_admission_samples_deadline_once( + monkeypatch: pytest.MonkeyPatch, terminal: bool +) -> None: + clock = iter([0.9, 1.1]) + monkeypatch.setattr( + deadline_module, "time", SimpleNamespace(monotonic=lambda: next(clock)) + ) + client = Mock() + client.get_execution_state.return_value = StateOutput( + operations=[], next_marker="next" + ) + state = ExecutionState( + "execution", + "token", + {}, + client, + PluginExecutor([]), + invocation_deadline=InvocationDeadline(1.0), + ) + if terminal: + assert ( + state.fetch_paginated_operations( + [], "token", "first", finish_terminal_checkpoint=True + ) + == [] + ) + else: + with pytest.raises(InvocationDeadlineExceeded): + state.fetch_paginated_operations([], "token", "first") + client.get_execution_state.assert_called_once_with( + durable_execution_arn="execution", checkpoint_token="token", next_marker="first" + ) diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/checkpoint_response_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/checkpoint_response_int_test.py index 81c36306b..5192c570b 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/checkpoint_response_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/checkpoint_response_int_test.py @@ -124,6 +124,7 @@ def my_handler(event, context: DurableContext) -> str: } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -187,6 +188,7 @@ def my_handler(event, context: DurableContext) -> list[str]: } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -243,6 +245,7 @@ def my_handler(event, context: DurableContext) -> str: } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -302,6 +305,7 @@ def my_handler(event, context: DurableContext) -> str: } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -415,6 +419,7 @@ def mock_checkpoint( } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -467,6 +472,7 @@ def my_handler(event, context: DurableContext): } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -526,6 +532,7 @@ def my_handler(event, context: DurableContext) -> str: } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -627,6 +634,7 @@ def mock_checkpoint( } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -694,6 +702,7 @@ def my_handler(event, context: DurableContext) -> str: } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -755,6 +764,7 @@ def my_handler(event, context: DurableContext) -> str: } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -833,6 +843,7 @@ def my_handler(event, context: DurableContext) -> str: } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None 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 9768ad851..fb256ebbb 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 @@ -57,6 +57,7 @@ def deserialize(self, data: str, _: SerDesContext) -> dict[str, Any]: def _create_lambda_context(): """Create a mock Lambda context.""" ctx = Mock() + ctx.get_remaining_time_in_millis.return_value = 900000 ctx.aws_request_id = "test-request-id" ctx.client_context = None ctx.identity = None diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/error_hierarchy_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/error_hierarchy_int_test.py index 292a1707e..ec2fafcb7 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/error_hierarchy_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/error_hierarchy_int_test.py @@ -37,6 +37,7 @@ def _lambda_context() -> Mock: lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/execution_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/execution_int_test.py index 40ddcb489..6f8ad2b37 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/execution_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/execution_int_test.py @@ -155,6 +155,7 @@ def mock_checkpoint( # Create mock lambda context lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -241,6 +242,7 @@ def mock_checkpoint( # Create mock lambda context lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -328,6 +330,7 @@ def my_handler(event, context): # Create mock lambda context lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -429,6 +432,7 @@ def mock_checkpoint_failure( # Create mock lambda context lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -483,6 +487,7 @@ def my_handler(event: Any, context: DurableContext): # Create mock lambda context lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None @@ -582,6 +587,7 @@ def mock_checkpoint( } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request-id" lambda_context.client_context = None lambda_context.identity = None diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/filesystem_serdes_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/filesystem_serdes_int_test.py index 12c2efd83..72fddf16d 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/filesystem_serdes_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/filesystem_serdes_int_test.py @@ -39,6 +39,7 @@ def _create_lambda_context(): """Create a mock Lambda context.""" ctx = Mock() + ctx.get_remaining_time_in_millis.return_value = 900000 ctx.aws_request_id = "test-request-id" ctx.client_context = None ctx.identity = None diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/invocation_deadline_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/invocation_deadline_int_test.py new file mode 100644 index 000000000..ebbe15403 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/invocation_deadline_int_test.py @@ -0,0 +1,515 @@ +"""Public lifecycle regressions for #741, including non-cooperative limits.""" + +from __future__ import annotations + +import json +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from typing import Any +from unittest.mock import Mock + +import pytest + +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python import deadline as deadline_module +from aws_durable_execution_sdk_python.config import ( + CompletionConfig, + Duration, + ParallelConfig, + StepConfig, + StepSemantics, +) +from aws_durable_execution_sdk_python.exceptions import ( + InvocationDeadlineExceeded, + OrphanedChildException, + StepError, +) +from aws_durable_execution_sdk_python.execution import ( + DurableExecutionInvocationInputWithClient, + InitialExecutionState, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + ContextDetails, + ExecutionDetails, + Operation, + OperationAction, + OperationStatus, + OperationType, + OperationUpdate, + StepDetails, + StateOutput, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + OperationChangeInfo, +) +from aws_durable_execution_sdk_python.retries import RetryPresets + + +class RecordingService: + """A serialized checkpoint service retaining the real SDK's operation IDs.""" + + def __init__(self) -> None: + self.operations: dict[str, Operation] = {} + self.calls: list[list[OperationUpdate]] = [] + self.block_transport = threading.Event() + self.transport_entered = threading.Event() + self.release_transport = threading.Event() + + def checkpoint( + self, *, updates: list[OperationUpdate], **kwargs: Any + ) -> CheckpointOutput: + if self.block_transport.is_set(): + self.transport_entered.set() + assert self.release_transport.wait(5) + self.calls.append(updates) + changed = [] + for update in updates: + status = { + OperationAction.START: OperationStatus.STARTED, + OperationAction.SUCCEED: OperationStatus.SUCCEEDED, + OperationAction.FAIL: OperationStatus.FAILED, + }[update.action] + operation = Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=status, + parent_id=update.parent_id, + name=update.name, + sub_type=update.sub_type, + step_details=StepDetails(result=update.payload, error=update.error) + if update.operation_type is OperationType.STEP + else None, + context_details=ContextDetails( + result=update.payload, error=update.error + ) + if update.operation_type is OperationType.CONTEXT + else None, + ) + self.operations[operation.operation_id] = operation + changed.append(operation) + return CheckpointOutput( + checkpoint_token="next-token", + new_execution_state=CheckpointUpdatedExecutionState(operations=changed), + ) + + def event( + self, name: str = "execution" + ) -> DurableExecutionInvocationInputWithClient: + initial = Operation( + operation_id=name, + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + execution_details=ExecutionDetails( + input_payload=json.dumps({"name": name}) + ), + ) + return DurableExecutionInvocationInputWithClient( + durable_execution_arn=f"arn:test:execution/{name}", + checkpoint_token="token", + initial_execution_state=InitialExecutionState( + operations=[initial, *self.operations.values()], next_marker="" + ), + service_client=self, # type: ignore[arg-type] + ) + + +def context(remaining_ms: int = 900000) -> Mock: + value = Mock() + value.get_remaining_time_in_millis.return_value = remaining_ms + value.aws_request_id = "request" + value.client_context = None + value.identity = None + value.invoked_function_arn = "test" + value.tenant_id = None + return value + + +def test_expired_invocation_never_enters_handler_or_dispatches_checkpoint() -> None: + service = RecordingService() + body = Mock() + wrapped = durable_execution(body, plugins=[]) + with pytest.raises(InvocationDeadlineExceeded): + wrapped(service.event(), context(0)) + body.assert_not_called() + assert not service.calls + + +def test_deadline_retry_replays_saved_success_and_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + now = [100.0] + monkeypatch.setattr( + deadline_module, "time", SimpleNamespace(monotonic=lambda: now[0]) + ) + service = RecordingService() + effects: list[str] = [] + attempts = 0 + observed_errors: list[tuple[str, str | None]] = [] + + @durable_execution(plugins=[]) # type: ignore[call-arg] + def handler(_event: Any, ctx: DurableContext) -> Any: + nonlocal attempts + attempts += 1 + + def save(_: Any) -> str: + effects.append("saved") + return "saved" + + saved = ctx.step(save, name="saved") + + def fail(_: Any) -> None: + effects.append("failed") + raise ValueError("persisted failure") + + try: + ctx.step( + fail, + name="failed", + config=StepConfig(retry_strategy=RetryPresets.none()), + ) + except StepError as error: + observed_errors.append((error.error_type, error.message)) + if attempts == 1: + now[0] += 1000 + + def last(_: Any) -> str: + effects.append("last") + return saved + + return ctx.step(last, name="last") + + with pytest.raises(InvocationDeadlineExceeded): + handler(service.event(), context()) + assert effects == ["saved", "failed"] + outcome = handler(service.event(), context()) + assert outcome["Status"] == "SUCCEEDED" + assert json.loads(outcome["Result"]) == "saved" + assert effects == ["saved", "failed", "last"] + assert observed_errors == [("ValueError", "persisted failure")] * 2 + + +def test_expiry_inside_step_does_not_call_user_retry_strategy( + monkeypatch: pytest.MonkeyPatch, +) -> None: + now = [100.0] + monkeypatch.setattr( + deadline_module, "time", SimpleNamespace(monotonic=lambda: now[0]) + ) + service = RecordingService() + retry = Mock() + + @durable_execution(plugins=[]) # type: ignore[call-arg] + def handler(_event: Any, ctx: DurableContext) -> Any: + def body(_: Any) -> str: + now[0] += 1000 + return "late" + + return ctx.step(body, name="work", config=StepConfig(retry_strategy=retry)) + + with pytest.raises(InvocationDeadlineExceeded): + handler(service.event(), context()) + retry.assert_not_called() + assert all( + update.action is OperationAction.START + for batch in service.calls + for update in batch + ) + + +def test_completed_ancestor_rejects_new_step_before_body_or_checkpoint() -> None: + service = RecordingService() + loser_started = threading.Event() + parent_completed = threading.Event() + rejected = threading.Event() + late_body = Mock() + + def loser(child: DurableContext) -> Any: + def held_io(_: Any) -> None: + loser_started.set() + assert parent_completed.wait(5) + + try: + child.step(held_io, name="held") + finally: + try: + child.step(late_body, name="late") + except OrphanedChildException: + rejected.set() + raise + + def winner(child: DurableContext) -> Any: + assert loser_started.wait(5) + return child.step(lambda _: "winner", name="win") + + @durable_execution(plugins=[]) # type: ignore[call-arg] + def handler(_event: Any, ctx: DurableContext) -> Any: + result = ctx.parallel( + [loser, winner], + name="race", + config=ParallelConfig( + max_concurrency=2, completion_config=CompletionConfig.first_successful() + ), + ) + parent_completed.set() + return result.get_results() + + outcome = handler(service.event(), context()) + assert outcome["Status"] == "SUCCEEDED" + assert rejected.is_set() + late_body.assert_not_called() + assert all(update.name != "late" for batch in service.calls for update in batch) + + +def test_early_completion_joins_orphan_waiter_after_budget_expires() -> None: + service = RecordingService() + waiter_started = threading.Event() + result_computed = threading.Event() + waiter_finished = threading.Event() + + def loser(child: DurableContext) -> Any: + # Park in the same SDK refresh wait used by the coordinator. Its refresh + # is deliberately due after the invocation budget, with no blocked I/O. + refresh = child.state.schedule_refresh(time.time() + 60) + waiter_started.set() + try: + refresh.wait() + finally: + waiter_finished.set() + + def winner(child: DurableContext) -> Any: + assert waiter_started.wait(5) + return child.step(lambda _: "winner", name="win") + + @durable_execution(plugins=[]) # type: ignore[call-arg] + def handler(_event: Any, ctx: DurableContext) -> Any: + ctx.parallel( + [loser, winner], + name="race", + config=ParallelConfig( + max_concurrency=2, completion_config=CompletionConfig.first_successful() + ), + ) + result_computed.set() + return "winner" + + with ThreadPoolExecutor(max_workers=1) as callers: + future = callers.submit(handler, service.event(), context(2000)) + assert result_computed.wait(2) + with pytest.raises(InvocationDeadlineExceeded): + future.result(timeout=5) + assert waiter_finished.is_set() + assert not any( + thread.name.startswith("dex-handler") for thread in threading.enumerate() + ) + + +@pytest.mark.parametrize("blocked_part", ["transport", "plugin", "user"]) +def test_noncooperative_work_is_joined_and_other_invocations_stay_healthy( + blocked_part: str, +) -> None: + service = RecordingService() + entered = threading.Event() + release = threading.Event() + + class HeldUpdate(DurableInstrumentationPlugin): + def on_operation_change(self, _info: OperationChangeInfo) -> None: + entered.set() + assert release.wait(5) + + plugins: list[DurableInstrumentationPlugin] = ( + [HeldUpdate()] if blocked_part == "plugin" else [] + ) + if blocked_part == "transport": + service.block_transport.set() + entered = service.transport_entered + release = service.release_transport + + @durable_execution(plugins=plugins) # type: ignore[call-arg] + def handler(_event: Any, ctx: DurableContext) -> Any: + def body(_: Any) -> str: + if blocked_part == "user": + entered.set() + assert release.wait(5) + return "done" + + return ctx.step( + body, + name="work", + config=StepConfig(step_semantics=StepSemantics.AT_MOST_ONCE_PER_RETRY), + ) + + healthy_service = RecordingService() + healthy = durable_execution( + lambda _, ctx: ctx.step(lambda _: "healthy", name="healthy"), plugins=[] + ) + with ThreadPoolExecutor(max_workers=2) as callers: + expired = callers.submit(handler, service.event(), context(2000)) + try: + assert entered.wait(2) + assert not expired.done() + healthy_result = callers.submit( + healthy, healthy_service.event(), context() + ).result(timeout=2) + assert healthy_result["Status"] == "SUCCEEDED" + # A blocked transport/hook/user body is explicitly not detached. + time.sleep(2.0) + assert not expired.done() + finally: + release.set() + with pytest.raises(InvocationDeadlineExceeded): + expired.result(timeout=5) + assert not any( + update.action is OperationAction.SUCCEED + for batch in service.calls + for update in batch + ) + + +@pytest.mark.parametrize("fail", [False, True]) +def test_accepted_terminal_result_survives_slow_delivery_hook( + monkeypatch: pytest.MonkeyPatch, fail: bool +) -> None: + monkeypatch.setattr( + "aws_durable_execution_sdk_python.execution.LAMBDA_RESPONSE_SIZE_LIMIT", 20 + ) + service = RecordingService() + entered = threading.Event() + release = threading.Event() + ends: list[InvocationEndInfo] = [] + + class HeldUpdate(DurableInstrumentationPlugin): + def on_operation_change(self, _info: OperationChangeInfo) -> None: + entered.set() + assert release.wait(5) + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + ends.append(info) + + @durable_execution(plugins=[HeldUpdate()]) # type: ignore[call-arg] + def handler(_event: Any, _ctx: DurableContext) -> str: + if fail: + raise ValueError("a large persisted failure" * 10) + return "a large persisted result" * 10 + + with ThreadPoolExecutor(max_workers=1) as callers: + future = callers.submit(handler, service.event(), context(2000)) + try: + assert entered.wait(2) + time.sleep(2.0) + assert not future.done() + finally: + release.set() + result = future.result(timeout=5) + assert result["Status"] == ("FAILED" if fail else "SUCCEEDED") + assert len(ends) == 1 + assert ends[0].status.value == result["Status"] + updates = [update for batch in service.calls for update in batch] + assert len(updates) == 1 + assert updates[0].operation_type is OperationType.EXECUTION + assert updates[0].action is ( + OperationAction.FAIL if fail else OperationAction.SUCCEED + ) + + +@pytest.mark.parametrize("fail", [False, True]) +@pytest.mark.parametrize("expire_after_first_page", [False, True]) +def test_accepted_terminal_checkpoint_stops_pagination_at_cutoff( + monkeypatch: pytest.MonkeyPatch, fail: bool, expire_after_first_page: bool +) -> None: + now = [100.0] + monkeypatch.setattr( + deadline_module, "time", SimpleNamespace(monotonic=lambda: now[0]) + ) + monkeypatch.setattr( + "aws_durable_execution_sdk_python.execution.LAMBDA_RESPONSE_SIZE_LIMIT", 20 + ) + pages: list[str] = [] + + class TerminalPagesService(RecordingService): + def checkpoint( + self, *, updates: list[OperationUpdate], **kwargs: Any + ) -> CheckpointOutput: + response = super().checkpoint(updates=updates, **kwargs) + if not expire_after_first_page: + now[0] += 1000 + return CheckpointOutput( + checkpoint_token=response.checkpoint_token, + new_execution_state=CheckpointUpdatedExecutionState( + operations=response.new_execution_state.operations, + next_marker="terminal-page", + ), + ) + + def get_execution_state( + self, *, next_marker: str, **kwargs: Any + ) -> StateOutput: + pages.append(next_marker) + now[0] += 1000 + return StateOutput(operations=[], next_marker="must-not-fetch") + + def handler(_event: Any, _ctx: DurableContext) -> str: + if fail: + raise ValueError("a large persisted failure" * 10) + return "a large persisted result" * 10 + + service = TerminalPagesService() + wrapped = durable_execution(handler, plugins=[]) + outcome = wrapped(service.event(), context()) + assert outcome["Status"] == ("FAILED" if fail else "SUCCEEDED") + assert pages == (["terminal-page"] if expire_after_first_page else []) + + +def test_public_parallel_wait_coordinator_unwinds_at_deadline() -> None: + service = RecordingService() + held = threading.Event() + release = threading.Event() + root_unwound = threading.Event() + + def waiting(child: DurableContext) -> None: + child.wait(Duration.from_seconds(60), name="durable-wait") + + def busy(child: DurableContext) -> str: + def body(_: Any) -> str: + held.set() + assert release.wait(8) + return "late" + + return child.step(body, name="held-step") + + def handler(_event: Any, ctx: DurableContext) -> None: + try: + ctx.parallel( + [waiting, busy], + name="parallel", + config=ParallelConfig(max_concurrency=2), + ) + finally: + root_unwound.set() + + wrapped = durable_execution(handler, plugins=[]) + with ThreadPoolExecutor(max_workers=1) as callers: + future = callers.submit(wrapped, service.event(), context(2000)) + try: + assert held.wait(3) + # The actual public wait enters a timed resume wave while its peer + # runs. The coordinator must leave its queue wait at the deadline. + assert root_unwound.wait(4) + assert not future.done() # the uncooperative peer remains confined + finally: + release.set() + with pytest.raises(InvocationDeadlineExceeded): + future.result(timeout=5) + assert any( + update.name == "durable-wait" for batch in service.calls for update in batch + ) + assert all( + update.action is OperationAction.START + for batch in service.calls + for update in batch + ) diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py index 3cad34121..f7f1019f0 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py @@ -46,6 +46,7 @@ def on_invocation_end(self, info) -> None: def _lambda_context() -> Mock: ctx = Mock() + ctx.get_remaining_time_in_millis.return_value = 900000 ctx.aws_request_id = "test-request-id" ctx.client_context = None ctx.identity = None diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py index 6b9344238..9784545f9 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py @@ -48,6 +48,7 @@ def on_invocation_end(self, info) -> None: def _lambda_context() -> Mock: ctx = Mock() + ctx.get_remaining_time_in_millis.return_value = 900000 ctx.aws_request_id = "test-request-id" ctx.client_context = None ctx.identity = None diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py index 7b30828d0..069a7753f 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py @@ -68,6 +68,7 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: def _lambda_context() -> Mock: context = Mock() + context.get_remaining_time_in_millis.return_value = 900000 context.aws_request_id = "test-request-id" context.client_context = None context.identity = None 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 735ec2e49..914fd6ca8 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -400,6 +400,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -452,6 +453,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -500,6 +502,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -547,6 +550,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -602,6 +606,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -648,6 +653,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -695,6 +701,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -733,6 +740,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -791,6 +799,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -830,6 +839,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -869,6 +879,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -908,6 +919,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -951,6 +963,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -991,6 +1004,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1041,6 +1055,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1093,6 +1108,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1144,6 +1160,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1194,6 +1211,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1241,6 +1259,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1285,6 +1304,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1334,6 +1354,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1380,6 +1401,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1430,6 +1452,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1496,6 +1519,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1586,6 +1610,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1669,6 +1694,7 @@ def test_handler(event: Any, context: DurableContext) -> str: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1739,6 +1765,7 @@ def test_handler(event: Any, context: DurableContext) -> str: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1794,6 +1821,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1857,6 +1885,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1914,6 +1943,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: ) lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1971,6 +2001,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: } lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -1998,6 +2029,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: regular_event = {"key": "value", "data": "test"} lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -2029,6 +2061,7 @@ def test_handler(event: Any, context: DurableContext) -> dict: non_dict_event = "not a dict" lambda_context = Mock() + lambda_context.get_remaining_time_in_millis.return_value = 900000 lambda_context.aws_request_id = "test-request" lambda_context.client_context = None lambda_context.identity = None @@ -2731,6 +2764,7 @@ def _make_invocation_input(mock_client, next_marker="", input_payload="{}"): def _make_lambda_context(): """Helper to create a standard mock Lambda context.""" ctx = Mock() + ctx.get_remaining_time_in_millis.return_value = 900000 ctx.aws_request_id = "test-request" ctx.client_context = None ctx.identity = None diff --git a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py index 8cb081029..3f0e48fee 100644 --- a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py @@ -63,6 +63,7 @@ def invocation(client: Any = None) -> tuple[Any, Any]: service_client=client, ) context = Mock() + context.get_remaining_time_in_millis.return_value = 900000 context.aws_request_id = "worker-request" context.client_context = context.identity = context.invoked_function_arn = None context._epoch_deadline_time_in_ms = 0 diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/base_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/base_test.py index 686f2c2c6..c9153629a 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/base_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/base_test.py @@ -30,6 +30,10 @@ class ConcreteOperationExecutor(OperationExecutor[str]): """Concrete implementation for testing the abstract base class.""" def __init__(self): + self.state = Mock(spec=ExecutionState) + self.operation_identifier = OperationIdentifier( + "test_op", OperationSubType.STEP, operation_type=OperationType.STEP + ) self.check_result_status_called = 0 self.execute_called = 0 self.check_result_to_return = None 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 0ee83fdf3..126bf2675 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 @@ -6,6 +6,8 @@ import pytest +from aws_durable_execution_sdk_python.deadline import InvocationDeadline + # Mock the executor.execute method from aws_durable_execution_sdk_python.plugin import PluginExecutor from aws_durable_execution_sdk_python.concurrency.models import ( @@ -54,6 +56,7 @@ def create_test_context( """Helper to create DurableContext for tests with required execution_context.""" if state is None: state = Mock(spec=ExecutionState) + state.invocation_deadline = InvocationDeadline() state.durable_execution_arn = ( "arn:aws:durable:us-east-1:123456789012:execution/test" ) @@ -188,6 +191,11 @@ def mock_run_in_child_context(func, name, config): # Create a minimal ExecutionState mock class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -232,6 +240,11 @@ def mock_run_in_child_context(func, name, config): return func("mock_context") class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -370,6 +383,11 @@ def callable_func(ctx, item, idx, items): ) as mock_execute: class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -429,6 +447,11 @@ def callable_func(ctx, item, idx, items): executor_context.create_child_context = lambda *args, **kwargs: Mock() class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -489,6 +512,11 @@ def callable_func(ctx, item, idx, items): executor_context.create_child_context = lambda *args, **kwargs: child_context class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -544,6 +572,11 @@ def mock_summary_generator(result): executor_context.create_child_context = Mock(return_value=_child_ctx) class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -590,6 +623,11 @@ def callable_func(ctx, item, idx, items): executor_context.create_child_context = Mock(return_value=Mock()) # SLF001 class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -635,6 +673,11 @@ def func(ctx, item, index, array): config = MapConfig(summary_generator=None) class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -686,6 +729,11 @@ def callable_func(ctx, item, idx, items): # Mock execution state that indicates operation already succeeded class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -761,6 +809,11 @@ def callable_func(ctx, item, idx, items): # Mock execution state that indicates operation succeeded but children need replay class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -870,6 +923,11 @@ def test_func(ctx, item, idx, items): execution_count = 0 class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -974,6 +1032,7 @@ def get_checkpoint(op_id): return child_checkpoint if op_id.startswith("child-") else parent_checkpoint mock_state = Mock() + mock_state.invocation_deadline = InvocationDeadline() mock_state.durable_execution_arn = "arn:test" mock_state.get_checkpoint_result = Mock(side_effect=get_checkpoint) mock_state.create_checkpoint = Mock() @@ -1046,6 +1105,7 @@ def get_checkpoint(op_id): return child_checkpoint if op_id.startswith("child-") else parent_checkpoint mock_state = Mock() + mock_state.invocation_deadline = InvocationDeadline() mock_state.durable_execution_arn = "arn:test" mock_state.get_checkpoint_result = Mock(side_effect=get_checkpoint) mock_state.create_checkpoint = Mock() @@ -1102,6 +1162,11 @@ def func(ctx, item, idx, items): return {"item": item.upper(), "index": idx} class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -1183,6 +1248,7 @@ def get_checkpoint(op_id): ) mock_state = Mock() + mock_state.invocation_deadline = InvocationDeadline() mock_state.durable_execution_arn = "arn:test" mock_state.get_checkpoint_result = Mock(side_effect=get_checkpoint) mock_state.create_checkpoint = Mock() @@ -1256,6 +1322,7 @@ def get_checkpoint(op_id): ) mock_state = Mock() + mock_state.invocation_deadline = InvocationDeadline() mock_state.durable_execution_arn = "arn:test" mock_state.get_checkpoint_result = Mock(side_effect=get_checkpoint) mock_state.create_checkpoint = Mock() @@ -1342,6 +1409,7 @@ def get_checkpoint(op_id): ) mock_state = Mock() + mock_state.invocation_deadline = InvocationDeadline() mock_state.durable_execution_arn = "arn:test" mock_state.get_checkpoint_result = Mock(side_effect=get_checkpoint) mock_state.create_checkpoint = Mock() @@ -1397,6 +1465,7 @@ def map_func(ctx, item, idx, items): return f"processed_{item}" mock_state = Mock() + mock_state.invocation_deadline = InvocationDeadline() mock_state.durable_execution_arn = "arn:test" parent_checkpoint = Mock() @@ -1557,6 +1626,11 @@ def test_map_handler_defaults_summary_generator_for_user_config(): """ class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -1638,6 +1712,11 @@ def predicate(s: CompletionStatus) -> CompletionDecision: ) as mock_execute: class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): 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 930197c23..56332aee4 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 @@ -8,6 +8,8 @@ import pytest +from aws_durable_execution_sdk_python.deadline import InvocationDeadline + from aws_durable_execution_sdk_python.plugin import PluginExecutor from aws_durable_execution_sdk_python.concurrency.executor import ConcurrentExecutor from aws_durable_execution_sdk_python.identifier import OperationIdNamespace @@ -59,6 +61,7 @@ def create_test_context( """Helper to create DurableContext for tests with required execution_context.""" if state is None: state = Mock(spec=ExecutionState) + state.invocation_deadline = InvocationDeadline() state.durable_execution_arn = ( "arn:aws:durable:us-east-1:123456789012:execution/test" ) @@ -215,6 +218,11 @@ def func2(ctx): config = ParallelConfig(max_concurrency=2) class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -265,6 +273,11 @@ def func1(ctx): callables = [func1] class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -315,6 +328,11 @@ def func1(ctx): config = ParallelConfig(max_concurrency=5) class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -373,6 +391,11 @@ def func1(ctx): callables = [func1] class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -494,6 +517,11 @@ def func1(ctx): callables = [func1] class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -547,6 +575,11 @@ def mock_summary_generator(result): config = ParallelConfig(summary_generator=mock_summary_generator) class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -596,6 +629,11 @@ def func2(ctx): callables = [func1, func2] class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -652,6 +690,11 @@ def func3(ctx): config = ParallelConfig(summary_generator=None) class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -706,6 +749,11 @@ def func2(ctx): # Mock execution state that indicates operation already succeeded class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -781,6 +829,11 @@ def func1(ctx): # Mock execution state that indicates operation succeeded but children need replay class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -892,6 +945,11 @@ def task2(ctx): execution_count = 0 class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -994,6 +1052,7 @@ def get_checkpoint(op_id): return child_checkpoint if op_id.startswith("child-") else parent_checkpoint mock_state = Mock() + mock_state.invocation_deadline = InvocationDeadline() mock_state.durable_execution_arn = "arn:test" mock_state.get_checkpoint_result = Mock(side_effect=get_checkpoint) mock_state.create_checkpoint = Mock() @@ -1063,6 +1122,7 @@ def get_checkpoint(op_id): return child_checkpoint if op_id.startswith("child-") else parent_checkpoint mock_state = Mock() + mock_state.invocation_deadline = InvocationDeadline() mock_state.durable_execution_arn = "arn:test" mock_state.get_checkpoint_result = Mock(side_effect=get_checkpoint) mock_state.create_checkpoint = Mock() @@ -1122,6 +1182,11 @@ def func3(ctx): callables = [func1, func2, func3] class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -1204,6 +1269,7 @@ def get_checkpoint(op_id): ) mock_state = Mock() + mock_state.invocation_deadline = InvocationDeadline() mock_state.durable_execution_arn = "arn:test" mock_state.get_checkpoint_result = Mock(side_effect=get_checkpoint) mock_state.create_checkpoint = Mock() @@ -1277,6 +1343,7 @@ def get_checkpoint(op_id): ) mock_state = Mock() + mock_state.invocation_deadline = InvocationDeadline() mock_state.durable_execution_arn = "arn:test" mock_state.get_checkpoint_result = Mock(side_effect=get_checkpoint) mock_state.create_checkpoint = Mock() @@ -1363,6 +1430,7 @@ def get_checkpoint(op_id): ) mock_state = Mock() + mock_state.invocation_deadline = InvocationDeadline() mock_state.durable_execution_arn = "arn:test" mock_state.get_checkpoint_result = Mock(side_effect=get_checkpoint) mock_state.create_checkpoint = Mock() @@ -1550,6 +1618,11 @@ def test_parallel_handler_defaults_summary_generator_for_user_config(): """A user config without a summary generator gets the default (JS parity).""" class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): @@ -1633,6 +1706,11 @@ def predicate(s: CompletionStatus) -> CompletionDecision: ) as mock_execute: class MockExecutionState: + invocation_deadline = InvocationDeadline() + + def check_operation_active(self, _identifier): + pass + _plugin_executor = PluginExecutor([]) def register_branch_pool(self, pool): From 5141c9cbe8cb6bb7615ab287bc48ca38484ddacb Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:43:54 +0000 Subject: [PATCH 2/3] test: supply deadline in wheel context fixture --- .github/tests/otel_lifecycle_compatibility_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/tests/otel_lifecycle_compatibility_test.py b/.github/tests/otel_lifecycle_compatibility_test.py index 8e9ce857b..c4febcd65 100644 --- a/.github/tests/otel_lifecycle_compatibility_test.py +++ b/.github/tests/otel_lifecycle_compatibility_test.py @@ -115,6 +115,7 @@ def handler(event, durable_context): }, } lambda_context = SimpleNamespace( + get_remaining_time_in_millis=lambda: 900000, aws_request_id="installed", client_context=None, identity=None, From 9c39c781c03d0630c32bf34c80cc43105dd0708f Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Thu, 22 Oct 2026 10:01:42 -0700 Subject: [PATCH 3/3] fix: preserve replay and late checkpoint outcomes Allow saved checkpoint reads and scoped ReplayChildren reconstruction without admitting new orphaned work, including virtual child contexts. Settle terminal checkpoint requests across the invocation cutoff before classifying their outcomes. Cover same-invocation map/parallel resumes, saved successes and errors, reconstruction scope cleanup, late terminal transport outcomes, and expired terminal requests that were never dispatched. --- .../README.md | 13 +- .../operation/base.py | 5 +- .../operation/child.py | 1 + .../aws_durable_execution_sdk_python/state.py | 50 ++++- .../tests/e2e/invocation_deadline_int_test.py | 194 +++++++++++++++++- .../operation/invocation_admission_test.py | 173 ++++++++++++++++ 6 files changed, 424 insertions(+), 12 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/invocation_admission_test.py diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index 106d6723c..e29724249 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -163,6 +163,11 @@ the cutoff, allowing invocation cleanup to join it. A newly created descendant of a completed context is rejected before its body, START hook, or checkpoint runs, even if it did not exist when its ancestor completed. +Saved checkpoint reads and `ReplayChildren` reconstruction, including virtual +children, remain available when a concurrent branch resumes in the same +invocation. Reconstructing a saved result does not permit new durable work +under that completed context. + This is cooperative interruption, **not a guarantee of bounded invocation shutdown or LMI worker reclamation**. Python cannot safely terminate arbitrary running threads. The SDK still joins invocation-owned workers before returning: @@ -177,9 +182,11 @@ tracks that remaining work, with real-LMI coverage in [test PR #743](https://git An expired invocation does not report an unacknowledged result as `SUCCEEDED` or `PENDING`. If the service has already acknowledged this invocation's terminal execution success/failure checkpoint, that outcome remains authoritative while -local delivery hooks finish. At the cutoff, further response metadata pages are -skipped for that accepted terminal checkpoint; inline and already-read operations -are retained. Already-running side effects cannot be undone; +local delivery hooks finish. A terminal checkpoint already in flight at the +cutoff is settled before reporting the invocation outcome, including an +acknowledgment received during cleanup. At the cutoff, further response metadata +pages are skipped for that accepted terminal checkpoint; inline and already-read +operations are retained. Already-running side effects cannot be undone; interrupted work continues to require the documented retry/idempotency strategy. ### Chained-invoke propagation diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/base.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/base.py index 871568550..da75adfcd 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/base.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/base.py @@ -109,6 +109,7 @@ class OperationExecutor(ABC, Generic[T]): state: ExecutionState operation_identifier: OperationIdentifier + is_virtual: bool = False def _get_checkpoint_result(self) -> CheckpointedResult: """Return this operation's checkpoint after validating replay identity.""" @@ -179,7 +180,9 @@ def process(self) -> T: InvalidStateError: If the check result is in an invalid state May raise operation-specific errors from check_result_status() or execute() """ - self.state.check_operation_active(self.operation_identifier) + self.state.check_operation_active( + self.operation_identifier, allow_replay=True, is_virtual=self.is_virtual + ) # Check 1: Entry (handles replay and existing checkpoints) result = self.check_result_status() 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 f29262673..840e3bfa9 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 @@ -204,6 +204,7 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T: self.operation_identifier, checkpointed_result.is_replay_children(), attempt=None if checkpointed_result.is_existent() else 1, + is_virtual=self.is_virtual, ) raw_result: T = wrapped_user_func() diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py index b737ae3cb..16c6f7793 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py @@ -2,6 +2,7 @@ from __future__ import annotations +import contextvars import functools import heapq import itertools @@ -61,6 +62,12 @@ # blocked caller stays fast. _BLOCKED_CALLER_WAIT_SECONDS = 0.001 +# Only saved ReplayChildren reconstruction may enter checkpointless virtual +# children. Keep the permission local to that body and that invocation's state. +_RECONSTRUCTING_COMPLETED_CONTEXT: contextvars.ContextVar[ExecutionState | None] = ( + contextvars.ContextVar("reconstructing_completed_context", default=None) +) + @dataclass(frozen=True) class CheckpointBatcherConfig: @@ -678,11 +685,24 @@ def _reject_if_execution_completed( operation_id=operation_id, ) - def check_operation_active(self, identifier: OperationIdentifier) -> None: + def check_operation_active( + self, + identifier: OperationIdentifier, + *, + allow_replay: bool = False, + is_virtual: bool = False, + ) -> None: """Reject new work under an expired invocation or completed ancestor.""" self.invocation_deadline.check() + # A terminal checkpoint is immutable. Reading it (or reconstructing its + # ReplayChildren result) does not admit new work under a completed parent. + replay = allow_replay and ( + self.get_checkpoint_result(identifier.operation_id).is_terminal() + or (is_virtual and _RECONSTRUCTING_COMPLETED_CONTEXT.get() is self) + ) with self._parent_done_lock: - self._reject_if_parent_done(identifier) + if not replay: + self._reject_if_parent_done(identifier) if identifier.parent_id: self._parent_to_children.setdefault(identifier.parent_id, set()).add( identifier.operation_id @@ -909,10 +929,11 @@ def create_checkpoint( and operation_update.operation_type is OperationType.EXECUTION and operation_update.action in {OperationAction.SUCCEED, OperationAction.FAIL} - and self._terminal_execution_acknowledged.is_set() ): - # The service already accepted this terminal outcome. Finish - # its local processing rather than report a retry after it. + # Settle this terminal request before reporting a retry. + # The batcher rejects unsent work at expiry; a request + # already in flight must be joined regardless, and may + # still succeed or return its original transport error. completion_event.wait() else: raise @@ -1549,17 +1570,30 @@ def wrap_user_function( operation_identifier: OperationIdentifier, is_replay_children: bool = False, attempt: int | None = None, + *, + is_virtual: bool = False, ): @functools.wraps(user_function) def wrapper(*args, **kwargs): - self.check_operation_active(operation_identifier) + self.check_operation_active( + operation_identifier, + allow_replay=is_replay_children or is_virtual, + is_virtual=is_virtual, + ) start_info = self._plugin_executor.on_user_function_start( operation_identifier, is_replay_children, attempt ) outcome = UserFunctionOutcome.INCOMPLETE error = None + replay_token = None try: - self.check_operation_active(operation_identifier) + self.check_operation_active( + operation_identifier, + allow_replay=is_replay_children or is_virtual, + is_virtual=is_virtual, + ) + if is_replay_children: + replay_token = _RECONSTRUCTING_COMPLETED_CONTEXT.set(self) result = user_function(*args, **kwargs) except Exception as exception: outcome = UserFunctionOutcome.FAILED @@ -1569,6 +1603,8 @@ def wrapper(*args, **kwargs): outcome = UserFunctionOutcome.SUCCEEDED return result finally: + if replay_token is not None: + _RECONSTRUCTING_COMPLETED_CONTEXT.reset(replay_token) # Runs on the thread that executed the user function, which is # the only thread that can release state bound to it. self._plugin_executor.on_user_function_end( diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/invocation_deadline_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/invocation_deadline_int_test.py index ca922bffc..50a2c4c24 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/invocation_deadline_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/invocation_deadline_int_test.py @@ -15,8 +15,10 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python import deadline as deadline_module from aws_durable_execution_sdk_python.config import ( + ChildConfig, CompletionConfig, Duration, + MapConfig, ParallelConfig, StepConfig, StepSemantics, @@ -49,6 +51,9 @@ OperationChangeInfo, ) from aws_durable_execution_sdk_python.retries import RetryPresets +from aws_durable_execution_sdk_python.deadline import InvocationDeadline +from aws_durable_execution_sdk_python.state import ExecutionState +from aws_durable_execution_sdk_python.threading import CompletionEvent class RecordingService: @@ -86,7 +91,12 @@ def checkpoint( if update.operation_type is OperationType.STEP else None, context_details=ContextDetails( - result=update.payload, error=update.error + result=update.payload, + error=update.error, + replay_children=bool( + update.context_options + and update.context_options.replay_children + ), ) if update.operation_type is OperationType.CONTEXT else None, @@ -140,6 +150,93 @@ def test_expired_invocation_never_enters_handler_or_dispatches_checkpoint() -> N assert not service.calls +@pytest.mark.parametrize("kind", ["map", "parallel"]) +@pytest.mark.parametrize("virtual", [False, True]) +def test_completed_children_replay_during_same_invocation_resume( + monkeypatch: pytest.MonkeyPatch, kind: str, virtual: bool +) -> None: + monkeypatch.setattr( + "aws_durable_execution_sdk_python.operation.child.CHECKPOINT_SIZE_LIMIT_BYTES", + 100, + ) + resumed = threading.Event() + visits: list[str] = [] + effects: list[str] = [] + + class WaitService(RecordingService): + def checkpoint( + self, *, updates: list[OperationUpdate], **kwargs: Any + ) -> CheckpointOutput: + response = super().checkpoint(updates=updates, **kwargs) + if updates: + return response + completed = [] + for operation in list(self.operations.values()): + if operation.operation_type is OperationType.WAIT: + operation = Operation.from_dict( + {**operation.to_dict(), "Status": "SUCCEEDED"} + ) + self.operations[operation.operation_id] = operation + completed.append(operation) + return CheckpointOutput( + "next-token", CheckpointUpdatedExecutionState(operations=completed) + ) + + service = WaitService() + + def saved(child: DurableContext) -> str: + visits.append("saved") + + def step(_: Any) -> str: + effects.append("once") + return "large" * 100 + + return child.step(step, name="value") + + def reconstruct(child: DurableContext) -> str: + visits.append("reconstruct") + return child.run_in_child_context( + saved, name="inner", config=ChildConfig(is_virtual=virtual) + ) + + def first(child: DurableContext) -> str: + result = child.run_in_child_context(reconstruct, name="saved-child") + child.wait(Duration.from_seconds(1), name="resume") + resumed.set() + return result + + def sibling(_child: DurableContext) -> str: + assert resumed.wait(5) + return "sibling" + + def handler(_event: Any, ctx: DurableContext) -> Any: + if kind == "parallel": + return ctx.parallel( + [first, sibling], config=ParallelConfig(max_concurrency=2) + ).get_results() + return ctx.map( + [0, 1], + lambda child, item, _index, _items: [first, sibling][item](child), + config=MapConfig(max_concurrency=2), + ).get_results() + + wrapped = durable_execution(handler, plugins=[]) + outcome = wrapped(service.event(), context()) + assert outcome["Status"] == "SUCCEEDED" + assert json.loads(outcome["Result"]) == ["large" * 100, "sibling"] + assert visits == ["reconstruct", "saved", "reconstruct", "saved"] + assert effects == ["once"] + saved_checkpoints = [ + update + for batch in service.calls + for update in batch + if update.name == "saved-child" and update.action is OperationAction.SUCCEED + ] + assert len(saved_checkpoints) == 1 + assert saved_checkpoints[0].context_options is not None + assert saved_checkpoints[0].context_options.replay_children + + def test_deadline_retry_replays_saved_success_and_failure( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -429,6 +526,101 @@ def handler(_event: Any, _ctx: DurableContext) -> str: ) +@pytest.mark.parametrize("fail", [False, True]) +@pytest.mark.parametrize("transport_error", [False, True]) +@pytest.mark.parametrize("checkpoint_token", ["next-token", None]) +def test_terminal_checkpoint_in_flight_at_cutoff_keeps_its_outcome( + monkeypatch: pytest.MonkeyPatch, + fail: bool, + transport_error: bool, + checkpoint_token: str | None, +) -> None: + now = [100.0] + monkeypatch.setattr( + deadline_module, "time", SimpleNamespace(monotonic=lambda: now[0]) + ) + monkeypatch.setattr( + "aws_durable_execution_sdk_python.execution.LAMBDA_RESPONSE_SIZE_LIMIT", 20 + ) + entered = threading.Event() + release = threading.Event() + settling = threading.Event() + failure = OSError("terminal transport failed") + ends: list[InvocationEndInfo] = [] + + class TerminalService(RecordingService): + def checkpoint( + self, *, updates: list[OperationUpdate], **kwargs: Any + ) -> CheckpointOutput: + entered.set() + assert release.wait(5) + if transport_error: + raise failure + response = super().checkpoint(updates=updates, **kwargs) + return CheckpointOutput(checkpoint_token, response.new_execution_state) + + class Recorder(DurableInstrumentationPlugin): + def on_invocation_end(self, info: InvocationEndInfo) -> None: + ends.append(info) + + original_deadline_wait = InvocationDeadline.wait + original_completion_wait = CompletionEvent.wait + original_close = ExecutionState.close + + def expire_wait( + deadline: InvocationDeadline, + event: CompletionEvent, + timeout: float | None = None, + ) -> bool: + assert entered.wait(5) + now[0] += 1000 + return original_deadline_wait(deadline, event, timeout) + + def observe_settlement( + event: CompletionEvent, timeout: float | None = None + ) -> bool: + if timeout is None: + settling.set() + return original_completion_wait(event, timeout) + + def observe_cleanup(state: ExecutionState) -> None: + settling.set() + original_close(state) + + monkeypatch.setattr(InvocationDeadline, "wait", expire_wait) + monkeypatch.setattr(CompletionEvent, "wait", observe_settlement) + monkeypatch.setattr(ExecutionState, "close", observe_cleanup) + + def handler(_event: Any, _ctx: DurableContext) -> str: + if fail: + raise ValueError("a large persisted failure" * 10) + return "a large persisted result" * 10 + + service = TerminalService() + wrapped = durable_execution(handler, plugins=[Recorder()]) + with ThreadPoolExecutor(max_workers=1) as callers: + future = callers.submit(wrapped, service.event(), context()) + try: + # Release only once the timed wait has expired and the wrapper is + # either settling that request or joining it during cleanup. + assert settling.wait(5) + assert not future.done() + assert not ends + finally: + release.set() + if transport_error: + with pytest.raises(OSError) as caught: + future.result(timeout=5) + assert caught.value is failure + expected_status = "RETRY" + else: + outcome = future.result(timeout=5) + expected_status = "FAILED" if fail else "SUCCEEDED" + assert outcome["Status"] == expected_status + assert len(ends) == 1 + assert ends[0].status.value == expected_status + + @pytest.mark.parametrize("fail", [False, True]) @pytest.mark.parametrize("expire_after_first_page", [False, True]) def test_accepted_terminal_checkpoint_stops_pagination_at_cutoff( diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/invocation_admission_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/invocation_admission_test.py new file mode 100644 index 000000000..560d764b5 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/invocation_admission_test.py @@ -0,0 +1,173 @@ +"""Saved replay is allowed without reopening completed contexts to new work.""" + +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from aws_durable_execution_sdk_python import deadline as deadline_module +from aws_durable_execution_sdk_python.config import StepConfig +from aws_durable_execution_sdk_python.deadline import InvocationDeadline +from aws_durable_execution_sdk_python.exceptions import ( + InvocationDeadlineExceeded, + OrphanedChildException, + StepError, +) +from aws_durable_execution_sdk_python.identifier import OperationIdentifier +from aws_durable_execution_sdk_python.lambda_service import ( + ContextDetails, + ErrorObject, + Operation, + OperationStatus, + OperationSubType, + OperationType, + OperationUpdate, + StepDetails, +) +from aws_durable_execution_sdk_python.operation.step import StepOperationExecutor +from aws_durable_execution_sdk_python.plugin import PluginExecutor +from aws_durable_execution_sdk_python.state import ExecutionState + + +def complete_context(state: ExecutionState, identifier: OperationIdentifier) -> None: + state.create_checkpoint( + OperationUpdate.create_context_succeed( + identifier, "null", OperationSubType.RUN_IN_CHILD_CONTEXT + ), + is_sync=False, + ) + + +@pytest.mark.parametrize("fail", [False, True]) +def test_completed_parent_allows_saved_step_result_without_user_work( + fail: bool, +) -> None: + body = Mock() + client = Mock() + identifier = OperationIdentifier( + "saved", + OperationSubType.STEP, + parent_id="parent", + operation_type=OperationType.STEP, + ) + operation = Operation( + operation_id="saved", + parent_id="parent", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + status=OperationStatus.FAILED if fail else OperationStatus.SUCCEEDED, + step_details=StepDetails( + result=None if fail else '"saved"', + error=ErrorObject.from_exception(ValueError("saved failure")) + if fail + else None, + ), + ) + state = ExecutionState( + "execution", "token", {"saved": operation}, client, PluginExecutor([]) + ) + complete_context( + state, + OperationIdentifier( + "parent", + OperationSubType.RUN_IN_CHILD_CONTEXT, + operation_type=OperationType.CONTEXT, + ), + ) + executor = StepOperationExecutor(body, StepConfig(), state, identifier, Mock()) + if fail: + with pytest.raises(StepError, match="saved failure"): + executor.process() + else: + assert executor.process() == "saved" + body.assert_not_called() + client.checkpoint.assert_not_called() + + +@pytest.mark.parametrize("fail", [False, True]) +def test_virtual_reconstruction_scope_never_admits_new_work_or_leaks( + fail: bool, +) -> None: + replay_id = OperationIdentifier( + "replay", + OperationSubType.RUN_IN_CHILD_CONTEXT, + operation_type=OperationType.CONTEXT, + ) + virtual_id = OperationIdentifier( + "virtual", + OperationSubType.RUN_IN_CHILD_CONTEXT, + parent_id="replay", + operation_type=OperationType.CONTEXT, + ) + operation = Operation( + operation_id="replay", + operation_type=OperationType.CONTEXT, + status=OperationStatus.SUCCEEDED, + context_details=ContextDetails(replay_children=True), + ) + state = ExecutionState( + "execution", "token", {"replay": operation}, Mock(), PluginExecutor([]) + ) + complete_context(state, replay_id) + new_body = Mock() + virtual_body = Mock(return_value="reconstructed") + virtual = state.wrap_user_function(virtual_body, virtual_id, is_virtual=True) + + def reconstruct() -> str: + assert virtual() == "reconstructed" + identifier = OperationIdentifier( + "new", + OperationSubType.STEP, + parent_id="replay", + operation_type=OperationType.STEP, + ) + with pytest.raises(OrphanedChildException): + StepOperationExecutor( + new_body, StepConfig(), state, identifier, Mock() + ).process() + if fail: + raise ValueError("reconstruction failed") + return "saved" + + replay = state.wrap_user_function(reconstruct, replay_id, is_replay_children=True) + if fail: + with pytest.raises(ValueError, match="reconstruction failed"): + replay() + else: + assert replay() == "saved" + with pytest.raises(OrphanedChildException): + virtual() + virtual_body.assert_called_once() + new_body.assert_not_called() + + +def test_unsent_terminal_checkpoint_expires_without_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + now = [0.0] + monkeypatch.setattr( + deadline_module, "time", SimpleNamespace(monotonic=lambda: now[0]) + ) + client = Mock() + state = ExecutionState( + "execution", + "token", + {}, + client, + PluginExecutor([]), + invocation_deadline=InvocationDeadline(1.0), + ) + with ThreadPoolExecutor(max_workers=1) as callers: + future = callers.submit( + state.create_checkpoint_sync, + OperationUpdate.create_execution_succeed('"saved"'), + ) + queued = state._checkpoint_queue.get(timeout=5) # noqa: SLF001 + state._checkpoint_queue.task_done() # noqa: SLF001 + now[0] = 2.0 + state._checkpoint_queue.put(queued) # noqa: SLF001 + state.checkpoint_batches_forever() + with pytest.raises(InvocationDeadlineExceeded): + future.result(timeout=5) + client.checkpoint.assert_not_called()