From 34b66535d63e689c1ff6b7619efebaaf5adf4a05 Mon Sep 17 00:00:00 2001 From: Ayushi Ahjolia Date: Wed, 22 Jul 2026 14:28:53 -0700 Subject: [PATCH] fix: handle missing checkpoint token from service --- .../concurrency/executor.py | 7 + .../lambda_service.py | 18 +- .../aws_durable_execution_sdk_python/state.py | 144 +++++++++++++--- .../tests/concurrency_test.py | 73 +++++++++ .../tests/lambda_service_test.py | 39 ++++- .../tests/state_test.py | 155 ++++++++++++++++++ 6 files changed, 404 insertions(+), 32 deletions(-) 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 256589ae..8539c252 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 @@ -252,6 +252,13 @@ def resubmitter(ready: list[ExecutableWithState]) -> None: """ try: execution_state.create_checkpoint() + except OrphanedChildException: + # The execution has already completed, so there is nothing left + # to resume. Stop cleanly instead of recording a resume error. + # OrphanedChildException is a BaseException and so is not caught + # by the except below. + self._completion_event.set() + return except Exception as exc: # noqa: BLE001 # resubmitter runs only on the single timer thread, so this # check-then-set needs no lock. First error wins: keep the diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py index 1db42ec3..339a0e9f 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py @@ -1086,7 +1086,8 @@ def from_dict( class CheckpointOutput: """Representation of the CheckpointDurableExecutionOutput structure of the DEX CheckpointDurableExecution API.""" - checkpoint_token: str + # None on the terminal checkpoint that ends the execution. + checkpoint_token: str | None new_execution_state: CheckpointUpdatedExecutionState @classmethod @@ -1109,8 +1110,7 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> CheckpointOutput: new_execution_state = CheckpointUpdatedExecutionState() return cls( - # TODO: maybe should throw if empty? - checkpoint_token=data.get("CheckpointToken", ""), + checkpoint_token=data.get("CheckpointToken"), new_execution_state=new_execution_state, ) @@ -1201,6 +1201,11 @@ def checkpoint( updates: list[OperationUpdate], client_token: str | None, ) -> CheckpointOutput: + # A checkpoint token is required. Raise a clear, retryable error (so the + # invocation re-drives) rather than letting the client reject an empty + # value with an opaque validation error. + if not checkpoint_token: + raise CheckpointError("Cannot checkpoint without a checkpoint token.") try: optional_params: dict[str, str] = {} if client_token is not None: @@ -1230,6 +1235,13 @@ def get_execution_state( next_marker: str, max_items: int = 1000, ) -> StateOutput: + # A checkpoint token is required. Raise a clear, retryable error (so the + # invocation re-drives) rather than letting the client reject an empty + # value with an opaque validation error. + if not checkpoint_token: + raise GetExecutionStateError( + "Cannot get execution state without a checkpoint token." + ) try: result: GetDurableExecutionStateResponseTypeDef = ( self.client.get_durable_execution_state( 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 baec099b..6b9609dd 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 @@ -15,6 +15,7 @@ from aws_durable_execution_sdk_python.exceptions import ( BackgroundThreadError, + CheckpointError, DurableExecutionsError, DurableOperationError, GetExecutionStateError, @@ -294,6 +295,12 @@ def __init__( self._overflow_queue: queue.Queue[QueuedOperation] = queue.Queue() self._checkpointing_stopped: threading.Event = threading.Event() self._checkpointing_failed: CompletionEvent = CompletionEvent() + # Set once the service confirms the execution has completed (a checkpoint + # response with no token). No further checkpoint can succeed afterward. + self._execution_completed: threading.Event = threading.Event() + # Serializes the execution-completed check with enqueueing so a checkpoint + # is never enqueued after _settle_after_execution_completed drains the queue. + self._completion_lock: Lock = Lock() # Concurrency management for parallel operations: parent_id -> {child_operation_ids} self._parent_to_children: dict[str, set[str]] = {} @@ -608,8 +615,27 @@ def create_checkpoint( # Create wrapper object for queue queued_op = QueuedOperation(operation_update, completion_event) - # Enqueue the wrapper object (operation_update can be None for empty checkpoints) - self._checkpoint_queue.put(queued_op) + # Enqueue under the same lock the background loop holds while it stops and + # drains the queue - on completion via _settle_after_execution_completed, or + # on failure in the exception handler. Re-check both terminal conditions + # inside the lock so a checkpoint is never enqueued after a drain and left + # with a waiter that blocks forever. + with self._completion_lock: + if self._checkpointing_failed.is_set(): + # Raises the stored BackgroundThreadError. + self._checkpointing_failed.wait() + if self._execution_completed.is_set(): + operation_id = ( + operation_update.operation_id + if operation_update is not None + else "" + ) + raise OrphanedChildException( + "Execution already completed; checkpoint will not be processed.", + operation_id=operation_id, + ) + # Enqueue the wrapper object (operation_update can be None for empty checkpoints) + self._checkpoint_queue.put(queued_op) # Conditionally wait for completion based on is_sync parameter if is_sync: @@ -756,15 +782,34 @@ def checkpoint_batches_forever(self) -> None: logger.debug("Checkpoint batch processed successfully") - # Update local token for next iteration - current_checkpoint_token = output.checkpoint_token + # The service omits the token only when it completes the + # execution. Advance on a returned token; otherwise treat a + # terminal batch as completion and any other omission as an + # invalid response. + execution_completed: bool = False + if output.checkpoint_token: + current_checkpoint_token = output.checkpoint_token + elif any( + update.operation_type is OperationType.EXECUTION + and ( + update.action is OperationAction.SUCCEED + or update.action is OperationAction.FAIL + ) + for update in updates + ): + execution_completed = True + else: + raise CheckpointError( + "Checkpoint response omitted the token outside of " + "execution completion." + ) previous_operations = self.operations # Fetch new operations from the API before unblocking sync waiters updated_operations = self.fetch_paginated_operations( output.new_execution_state.operations, - output.checkpoint_token, + current_checkpoint_token, output.new_execution_state.next_marker, ) self._plugin_executor.on_operation_update( @@ -777,6 +822,15 @@ def checkpoint_batches_forever(self) -> None: for queued_op in batch: if queued_op.completion_event is not None: queued_op.completion_event.set() + + if execution_completed: + # The execution is complete; no further checkpoint can + # succeed. Stop the loop and settle any still-queued + # operations (e.g. from orphaned concurrent branches) so + # their waiters do not block and no further checkpoint + # request is issued with the consumed token. + self._settle_after_execution_completed() + break except Exception as e: # Checkpoint failed - wake all blocked threads so they can raise error # Drain both queues and signal all completion events @@ -785,38 +839,74 @@ def checkpoint_batches_forever(self) -> None: "Checkpoint creation failed", e ) - # FIFO: although at this point order not really import any anymore - # Signal completion events for the failed batch + # Signal completion events for the failed batch (already dequeued, + # so no producer can race these). for queued_op in batch: if queued_op.completion_event is not None: queued_op.completion_event.set(bg_error) - # overflow 1st: although at this point order not really import any anymore - while not self._overflow_queue.empty(): - try: - item = self._overflow_queue.get_nowait() - if item.completion_event: - item.completion_event.set(bg_error) - except queue.Empty: - break - - # finally Wake all blocked threads in main queue - while not self._checkpoint_queue.empty(): - try: - item = self._checkpoint_queue.get_nowait() - if item.completion_event: - item.completion_event.set(bg_error) - except queue.Empty: - break - - # Set the failure event so future checkpoint attempts fail immediately - self._checkpointing_failed.set(bg_error) + # Drain the queues and set the failure flag under the completion + # lock so a concurrent create_checkpoint either observes the + # failure and raises, or has its operation drained here - never + # enqueued after the drain and left blocked. + with self._completion_lock: + while not self._overflow_queue.empty(): + try: + item = self._overflow_queue.get_nowait() + if item.completion_event: + item.completion_event.set(bg_error) + except queue.Empty: + break + + while not self._checkpoint_queue.empty(): + try: + item = self._checkpoint_queue.get_nowait() + if item.completion_event: + item.completion_event.set(bg_error) + except queue.Empty: + break + + # Future checkpoint attempts fail immediately. + self._checkpointing_failed.set(bg_error) # Exit the loop - error has been signaled to main thread via completion events break logger.debug("Background checkpoint processing stopped") + def _settle_after_execution_completed(self) -> None: + """Stop checkpointing and settle queued operations after the execution ends. + + Called when a checkpoint response omits the token on a terminal batch, + which the service does only when it completes the execution. Any operation + still queued belongs to work that can no longer be checkpointed (typically + an orphaned concurrent branch), so its waiter is settled with + OrphanedChildException rather than left blocking or sent with the consumed + token. + """ + with self._completion_lock: + self._execution_completed.set() + self._checkpointing_stopped.set() + + for pending_queue in (self._overflow_queue, self._checkpoint_queue): + while not pending_queue.empty(): + try: + queued_op: QueuedOperation = pending_queue.get_nowait() + except queue.Empty: + break + if queued_op.completion_event is not None: + operation_id: str = ( + queued_op.operation_update.operation_id + if queued_op.operation_update is not None + else "" + ) + queued_op.completion_event.set( + OrphanedChildException( + "Execution already completed; checkpoint will not be processed.", + operation_id=operation_id, + ) + ) + def stop_checkpointing(self) -> None: """Signal background thread to stop checkpointing. 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 d532d98d..82f5e7bf 100644 --- a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py @@ -1,5 +1,6 @@ """Tests for the concurrency module.""" +import contextlib import json import random import threading @@ -38,6 +39,7 @@ from aws_durable_execution_sdk_python.exceptions import ( ChildContextError, InvalidStateError, + OrphanedChildException, SuspendExecution, TimedSuspendExecution, ) @@ -1574,6 +1576,77 @@ def checkpoint(*args, **kwargs): executor.long_runner_release.set() +def test_concurrent_executor_resume_orphaned_stops_cleanly(): + """An orphaned resume refresh must be handled, not propagated. + + When the execution completes while a concurrent operation is orphaned, the + timer resubmit's checkpoint refresh raises OrphanedChildException. That is a + BaseException, so the resubmitter's generic `except Exception` does not catch + it; it must be handled explicitly so it neither crashes the timer thread nor + escapes execute(). Contrast with the RuntimeError case above, which does + propagate. + """ + + class TestExecutor(ConcurrentExecutor): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.long_runner_release = threading.Event() + + def execute_item(self, child_context, executable): + if executable.index == 0: + # Long-runner keeps the map alive so task 1 resumes in-process. + self.long_runner_release.wait(timeout=5) + return "result_A" + # Task 1 suspends with a past timestamp -> immediate in-process resume. + msg = "resume-me" + raise TimedSuspendExecution(msg, time.time() - 1) + + executables = [Executable(0, lambda: "task_A"), Executable(1, lambda: "task_B")] + completion_config = CompletionConfig( + min_successful=2, + tolerated_failure_count=None, + tolerated_failure_percentage=None, + ) + + executor = TestExecutor( + executables=executables, + max_concurrency=2, + completion_config=completion_config, + sub_type_top="TOP", + sub_type_iteration="ITER", + name_prefix="test_", + serdes=None, + ) + + execution_state = Mock() + + def checkpoint(*args, **kwargs): + # The resume refresh calls create_checkpoint() with no arguments; the + # execution has completed, so it is orphaned. + if not args and not kwargs: + raise OrphanedChildException( + "execution already completed", operation_id="op" + ) + + execution_state.create_checkpoint = Mock(side_effect=checkpoint) + + executor_context = Mock() + executor_context._create_step_id_for_logical_step = lambda *args: "1" # noqa: SLF001 + child_context = Mock() + child_context.state.wrap_user_function = lambda func, *args, **kwargs: func + executor_context.create_child_context = lambda *args, **kwargs: child_context + + try: + # The orphaned resume must not surface as an error; the map may suspend + # or complete, but OrphanedChildException must never escape. + with contextlib.suppress(SuspendExecution): + executor.execute(execution_state, executor_context) + except OrphanedChildException: + pytest.fail("OrphanedChildException from the resume refresh must be handled") + finally: + executor.long_runner_release.set() + + def test_concurrent_executor_with_timed_resubmit_while_other_task_running(): """Test timed resubmission while other tasks are still running.""" diff --git a/packages/aws-durable-execution-sdk-python/tests/lambda_service_test.py b/packages/aws-durable-execution-sdk-python/tests/lambda_service_test.py index 2d0ff829..55f1ebce 100644 --- a/packages/aws-durable-execution-sdk-python/tests/lambda_service_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/lambda_service_test.py @@ -1641,10 +1641,14 @@ def test_checkpoint_output_from_dict(): def test_checkpoint_output_from_dict_empty(): - """Test CheckpointOutput.from_dict with empty data.""" + """Test CheckpointOutput.from_dict with empty data. + + The service omits CheckpointToken on the terminal checkpoint; from_dict + surfaces that as None rather than an empty string. + """ data = {} output = CheckpointOutput.from_dict(data) - assert not output.checkpoint_token + assert output.checkpoint_token is None assert len(output.new_execution_state.operations) == 0 assert output.new_execution_state.next_marker is None @@ -1825,6 +1829,37 @@ def test_lambda_client_checkpoint_with_explicit_none_client_token(): assert result.checkpoint_token == "new_token" # noqa: S105 +@pytest.mark.parametrize("token", ["", None]) +def test_lambda_client_checkpoint_empty_token_raises(token): + """An empty/None checkpoint token raises a retryable CheckpointError, not a client error.""" + mock_client = Mock() + lambda_client = LambdaClient(mock_client) + update = OperationUpdate( + operation_id="op1", + operation_type=OperationType.STEP, + action=OperationAction.START, + ) + + with pytest.raises(CheckpointError) as exc_info: + lambda_client.checkpoint("arn123", token, [update], None) + + assert exc_info.value.is_retryable() + mock_client.checkpoint_durable_execution.assert_not_called() + + +@pytest.mark.parametrize("token", ["", None]) +def test_lambda_client_get_execution_state_empty_token_raises(token): + """An empty/None checkpoint token raises a retryable GetExecutionStateError, not a client error.""" + mock_client = Mock() + lambda_client = LambdaClient(mock_client) + + with pytest.raises(GetExecutionStateError) as exc_info: + lambda_client.get_execution_state("arn123", token, "marker") + + assert exc_info.value.is_retryable() + mock_client.get_durable_execution_state.assert_not_called() + + def test_lambda_client_checkpoint_with_empty_string_client_token(): """Test LambdaClient.checkpoint method with empty string client_token - should pass empty string.""" mock_client = Mock() diff --git a/packages/aws-durable-execution-sdk-python/tests/state_test.py b/packages/aws-durable-execution-sdk-python/tests/state_test.py index 4fb277ff..008ce91a 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -15,6 +15,7 @@ from aws_durable_execution_sdk_python.exceptions import ( BackgroundThreadError, + CheckpointError, DurableApiErrorCategory, GetExecutionStateError, OrphanedChildException, @@ -1783,6 +1784,160 @@ def run_batching(): pass # Expected +def test_checkpoint_missing_token_on_non_terminal_batch_fails(): + """A non-terminal checkpoint response without a token fails the checkpoint. + + The service omits the token only when it completes the execution, so a + missing token on any other checkpoint is an invalid response. + """ + mock_lambda_client = Mock(spec=LambdaClient) + mock_lambda_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token=None, + new_execution_state=CheckpointUpdatedExecutionState( + operations=[], next_marker=None + ), + ) + + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + + completion_event = CompletionEvent() + state._checkpoint_queue.put( + QueuedOperation( + OperationUpdate( + operation_id="op1", + operation_type=OperationType.STEP, + action=OperationAction.START, + ), + completion_event, + ) + ) + + thread = threading.Thread(daemon=True, target=state.checkpoint_batches_forever) + thread.start() + thread.join(timeout=2.0) + + assert completion_event.is_set() + try: + completion_event.wait() + pytest.fail("Should have raised BackgroundThreadError") + except BackgroundThreadError as bg_error: + assert isinstance(bg_error.source_exception, CheckpointError) + assert bg_error.source_exception.is_retryable() + + +@pytest.mark.parametrize( + "terminal_update", + [ + OperationUpdate.create_execution_succeed(payload="{}"), + OperationUpdate.create_execution_fail( + error=ErrorObject(message="boom", type="Error", data=None, stack_trace=None) + ), + ], +) +def test_checkpoint_missing_token_on_terminal_batch_succeeds(terminal_update): + """A terminal checkpoint response without a token completes normally. + + The service omits the token when it ends the execution (SUCCEED or FAIL); + that is expected and must not fail the checkpoint. + """ + mock_lambda_client = Mock(spec=LambdaClient) + mock_lambda_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token=None, + new_execution_state=CheckpointUpdatedExecutionState( + operations=[], next_marker=None + ), + ) + + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + + completion_event = CompletionEvent() + state._checkpoint_queue.put(QueuedOperation(terminal_update, completion_event)) + + thread = threading.Thread(daemon=True, target=state.checkpoint_batches_forever) + thread.start() + try: + # Returns without raising once the terminal checkpoint is processed. + completion_event.wait() + finally: + state.stop_checkpointing() + thread.join(timeout=2.0) + + assert completion_event.is_set() + assert state._execution_completed.is_set() + + +def test_settle_after_execution_completed_orphans_pending_operations(): + """When the execution completes, still-queued operations are settled as orphaned. + + Their waiters must be released with OrphanedChildException rather than left + blocking or sent with the consumed token, and checkpointing must stop. + """ + mock_lambda_client = Mock(spec=LambdaClient) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + + completion_event = CompletionEvent() + state._checkpoint_queue.put( + QueuedOperation( + OperationUpdate( + operation_id="orphan_op", + operation_type=OperationType.STEP, + action=OperationAction.START, + ), + completion_event, + ) + ) + + state._settle_after_execution_completed() + + assert state._execution_completed.is_set() + assert state._checkpointing_stopped.is_set() + with pytest.raises(OrphanedChildException): + completion_event.wait() + + +def test_create_checkpoint_after_execution_completed_is_orphaned(): + """A checkpoint attempted after the execution completed is rejected, not enqueued.""" + mock_lambda_client = Mock(spec=LambdaClient) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + state._execution_completed.set() + + with pytest.raises(OrphanedChildException): + state.create_checkpoint( + OperationUpdate( + operation_id="late_op", + operation_type=OperationType.STEP, + action=OperationAction.START, + ), + is_sync=True, + ) + + assert state._checkpoint_queue.empty() + + def test_collect_checkpoint_batch_shutdown_path(): """Test _collect_checkpoint_batch during shutdown with operations in queue.