Skip to content

Commit 47977fb

Browse files
authored
feat(testing): invoker state and checkpoint stamping
Align the local test runner's data-plane responses with the deployed service, validated against the JS example suite as a cross-language conformance check (failures reduced from 27 to 7; the remaining 7 are the unimplemented CHAINED_INVOKE operation and an upstream JS SDK replay bug). Invoker state and checkpoint stamping: - Track invocation state (PRE_INVOKE/INVOKING/COMPLETED) and gate checkpoints so at most one handler invocation is in flight per execution - Record per-update timestamps and stamp checkpoint operations with a monotonic sequence so delta computation and pagination share one snapshot - Offload handler invocation off the scheduler event loop so concurrent map/parallel branches no longer serialize on a single thread - Return a locked snapshot from in-process state reads History generation (GetDurableExecutionHistory): - Emit one history event per checkpoint update so step retries report the correct StepStarted/StepFailed counts and attempt numbers, instead of one pair per operation - Reconstruct events from recorded update transitions using the real operation for details, with a fallback for operations completed outside the checkpoint path Fidelity fixes: - Serialize Error and ReplayChildren in ContextDetails so child-context failures reconstruct the correct error class on replay - Store the RETRY payload as the step result so waitForCondition state replays, and map a RETRY without error to StepSucceeded - Set callback timeout errorType to Callback.Timeout/Callback.Heartbeat with the expected message strings Invocation timeout: - Add execution_timeout and invocation_timeout to the runner constructor and CLI to simulate the Lambda function Timeout; a timed-out invocation leaves the step STARTED and replays Closes #464 Closes #460 Closes #462 Closes #438 Closes #439
1 parent c575f46 commit 47977fb

36 files changed

Lines changed: 5043 additions & 1455 deletions

packages/aws-durable-execution-sdk-python-examples/test/conftest.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -106,14 +106,23 @@ def run(
106106
input: str | None = None, # noqa: A002
107107
timeout: int = 60,
108108
) -> DurableFunctionTestResult:
109-
"""Execute the durable function and return results."""
109+
"""Execute the durable function and return results.
110+
111+
The caller's ``timeout`` is the execution deadline. Local runs
112+
pass it as ``execution_timeout``; cloud runs pass it as the
113+
runner's ``timeout``.
114+
"""
115+
if isinstance(self._runner, DurableFunctionTestRunner):
116+
return self._runner.run(input=input, execution_timeout=timeout)
110117
return self._runner.run(input=input, timeout=timeout)
111118

112119
def run_async(
113120
self,
114121
input: str | None = None, # noqa: A002
115122
timeout: int = 60,
116123
) -> str:
124+
if isinstance(self._runner, DurableFunctionTestRunner):
125+
return self._runner.run_async(input=input, execution_timeout=timeout)
117126
return self._runner.run_async(input=input, timeout=timeout)
118127

119128
def send_callback_success(
@@ -222,7 +231,7 @@ def test_hello_world(durable_runner):
222231
if not handler:
223232
pytest.fail("handler is required for local mode tests")
224233
# Create local runner (needs cleanup via context manager)
225-
runner = DurableFunctionTestRunner(handler=handler)
234+
runner = DurableFunctionTestRunner(handler=handler, execution_timeout=60)
226235

227236
# Wrap in adapter and use context manager for proper cleanup
228237
with TestRunnerAdapter(runner, runner_mode) as adapter:

packages/aws-durable-execution-sdk-python-examples/test/wait_for_callback/test_wait_for_callback_timeout.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,4 +29,4 @@ def test_handle_wait_for_callback_timeout_scenarios(durable_runner):
2929
assert result_data["success"] is False
3030
assert isinstance(result_data["error"], str)
3131
assert len(result_data["error"]) > 0
32-
assert "Callback timed out: Callback.Timeout" == result_data["error"]
32+
assert "Callback timed out" == result_data["error"]
Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
"""The shared checkpoint write-transaction core.
2+
3+
Both checkpoint entry points need to detect an idempotent replay before
4+
doing real work. This class holds that shared logic so neither caller
5+
duplicates it.
6+
"""
7+
8+
from __future__ import annotations
9+
10+
from typing import TYPE_CHECKING
11+
12+
if TYPE_CHECKING:
13+
from aws_durable_execution_sdk_python_testing.execution import (
14+
CheckpointIdempotencyRecord,
15+
Execution,
16+
)
17+
18+
19+
class CheckpointCore:
20+
"""Shared checkpoint logic used by both the web and in-process paths."""
21+
22+
@staticmethod
23+
def match_cached(
24+
execution: Execution,
25+
checkpoint_token: str,
26+
client_token: str | None,
27+
) -> CheckpointIdempotencyRecord | None:
28+
"""Return the cached idempotency record when the incoming
29+
``(client_token, checkpoint_token)`` matches the last checkpoint.
30+
31+
Returns ``None`` when there is no match. A missing or empty
32+
``client_token`` cannot replay because idempotency requires an
33+
explicit client identifier.
34+
"""
35+
if not client_token:
36+
return None
37+
cached: CheckpointIdempotencyRecord | None = execution.last_checkpoint
38+
if cached is None:
39+
return None
40+
if (
41+
cached.client_token != client_token
42+
or cached.inbound_checkpoint_token != checkpoint_token
43+
):
44+
return None
45+
return cached
Lines changed: 121 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,10 @@
1-
"""Main checkpoint processor that orchestrates operation transformations."""
1+
"""In-process checkpoint flow.
2+
3+
Orchestrates the full checkpoint request for in-process callers
4+
(``InMemoryServiceClient``). Mirrors the flow inside
5+
``Executor.checkpoint_execution`` used by the HTTP path, so both
6+
entry points share identical delta + watermark semantics.
7+
"""
28

39
from __future__ import annotations
410

@@ -7,37 +13,55 @@
713
from aws_durable_execution_sdk_python.lambda_service import (
814
CheckpointOutput,
915
CheckpointUpdatedExecutionState,
10-
OperationUpdate,
1116
StateOutput,
1217
)
1318

19+
from aws_durable_execution_sdk_python_testing.checkpoint.core import CheckpointCore
1420
from aws_durable_execution_sdk_python_testing.checkpoint.transformer import (
15-
OperationTransformer,
21+
CheckpointRequestDispatcher,
1622
)
1723
from aws_durable_execution_sdk_python_testing.checkpoint.validators.checkpoint import (
1824
CheckpointValidator,
1925
)
2026
from aws_durable_execution_sdk_python_testing.exceptions import (
2127
InvalidParameterValueException,
2228
)
29+
from aws_durable_execution_sdk_python_testing.execution import (
30+
CheckpointIdempotencyRecord,
31+
OperationPaginatorState,
32+
)
2333
from aws_durable_execution_sdk_python_testing.observer import ExecutionNotifier
2434
from aws_durable_execution_sdk_python_testing.token import CheckpointToken
2535

2636

2737
if TYPE_CHECKING:
38+
from aws_durable_execution_sdk_python.lambda_service import OperationUpdate
39+
2840
from aws_durable_execution_sdk_python_testing.execution import Execution
2941
from aws_durable_execution_sdk_python_testing.scheduler import Scheduler
3042
from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore
3143

3244

45+
# Canonical invocation-page byte budget, shared with the Executor.
46+
DEFAULT_MAX_INVOCATION_PAGE_BYTES = 5 * 1024 * 1024
47+
48+
3349
class CheckpointProcessor:
34-
"""Handle OperationUpdate transformations and execution state updates."""
50+
"""In-process checkpoint flow used by ``InMemoryServiceClient``.
51+
52+
Uses the same :class:`CheckpointRequestDispatcher` + paginator
53+
primitives as ``Executor.checkpoint_execution`` so observable
54+
behaviour is identical across the two entry points.
55+
"""
3556

36-
def __init__(self, store: ExecutionStore, scheduler: Scheduler):
57+
def __init__(
58+
self,
59+
store: ExecutionStore,
60+
scheduler: Scheduler, # noqa: ARG002 — kept for backward-compatible signature
61+
):
3762
self._store = store
38-
self._scheduler = scheduler
3963
self._notifier = ExecutionNotifier()
40-
self._transformer = OperationTransformer()
64+
self._dispatcher = CheckpointRequestDispatcher()
4165

4266
def add_execution_observer(self, observer) -> None:
4367
"""Add observer for execution events."""
@@ -47,41 +71,77 @@ def process_checkpoint(
4771
self,
4872
checkpoint_token: str,
4973
updates: list[OperationUpdate],
50-
client_token: str | None, # noqa: ARG002
74+
client_token: str | None,
5175
) -> CheckpointOutput:
52-
"""Process checkpoint updates and return result with updated execution state."""
53-
# 1. Get current execution state
54-
token: CheckpointToken = CheckpointToken.from_str(checkpoint_token)
76+
"""Apply ``updates`` and return the delta since the handler's
77+
last observation.
78+
79+
Advances ``token_sequence`` exactly once; bumps
80+
``seq_counter`` once per accepted update via
81+
``touch_operation``; advances ``handler_seen_seq`` only for
82+
operations actually returned in the response.
83+
"""
84+
token = CheckpointToken.from_str(checkpoint_token)
5585
execution: Execution = self._store.load(token.execution_arn)
5686

57-
# 2. Validate checkpoint token
58-
if execution.is_complete or token.token_sequence != execution.token_sequence:
59-
msg: str = "Invalid checkpoint token"
60-
87+
# Idempotency first, before the token-sequence check.
88+
cached = _maybe_replay_cached(execution, checkpoint_token, client_token)
89+
if cached is not None:
90+
return cached
91+
92+
if (
93+
execution.is_complete
94+
or token.token_sequence != execution.token_sequence
95+
or token.invocation_id != execution.current_invocation_id
96+
):
97+
msg = "Invalid checkpoint token"
6198
raise InvalidParameterValueException(msg)
6299

63-
# 3. Validate all updates, state transitions are valid, sizes etc.
64-
CheckpointValidator.validate_input(updates, execution)
65-
66-
# 4. Transform OperationUpdate -> Operation and schedule future replays
67-
updated_operations, all_updates = self._transformer.process_updates(
68-
updates=updates,
69-
current_operations=execution.operations,
70-
notifier=self._notifier,
71-
execution_arn=token.execution_arn,
100+
if updates:
101+
CheckpointValidator.validate_input(updates, execution)
102+
self._dispatcher.apply_updates(
103+
execution=execution,
104+
updates=updates,
105+
client_token=client_token,
106+
notifier=self._notifier,
107+
touch=execution.touch_operation,
108+
)
109+
110+
new_token_sequence = execution.advance_token_sequence()
111+
112+
paginator = OperationPaginatorState.pin(execution)
113+
# The checkpoint response returns the full unseen delta in a single
114+
# response. Advance handler_seen_seq to cover every returned op so
115+
# the next delta carries only operations touched after this response.
116+
response_ops = paginator.unseen_operations()
117+
if response_ops:
118+
highest_delivered_seq = max(
119+
execution.operation_last_touched_seq[op.operation_id]
120+
for op in response_ops
121+
)
122+
paginator.advance_handler_seen(highest_delivered_seq)
123+
124+
new_token = CheckpointToken(
125+
execution_arn=execution.durable_execution_arn,
126+
token_sequence=new_token_sequence,
127+
invocation_id=execution.current_invocation_id,
128+
).to_str()
129+
130+
execution.last_checkpoint = CheckpointIdempotencyRecord(
131+
client_token=client_token or "",
132+
inbound_checkpoint_token=checkpoint_token,
133+
outbound_checkpoint_token=new_token,
134+
operations=list(response_ops),
135+
next_marker=None,
72136
)
73137

74-
# 5. Generate a new checkpoint token and save updated operations
75-
new_checkpoint_token = execution.get_new_checkpoint_token()
76-
execution.operations = updated_operations
77-
execution.updates.extend(all_updates)
78138
self._store.update(execution)
79139

80-
# 6. Return checkpoint result
81140
return CheckpointOutput(
82-
checkpoint_token=new_checkpoint_token,
141+
checkpoint_token=new_token,
83142
new_execution_state=CheckpointUpdatedExecutionState(
84-
operations=execution.get_navigable_operations(), next_marker=None
143+
operations=response_ops,
144+
next_marker=None,
85145
),
86146
)
87147

@@ -91,11 +151,35 @@ def get_execution_state(
91151
next_marker: str, # noqa: ARG002
92152
max_items: int = 1000, # noqa: ARG002
93153
) -> StateOutput:
94-
"""Get current execution state."""
95-
token: CheckpointToken = CheckpointToken.from_str(checkpoint_token)
96-
execution: Execution = self._store.load(token.execution_arn)
97-
98-
# TODO: paging when size or max
154+
"""Get current execution state.
155+
156+
Returns the full navigable operation list with a null marker.
157+
Marker round-tripping and the invocation-state gate are
158+
enforced on the HTTP path in :class:`Executor`.
159+
"""
160+
token = CheckpointToken.from_str(checkpoint_token)
161+
execution = self._store.load(token.execution_arn)
99162
return StateOutput(
100-
operations=execution.get_navigable_operations(), next_marker=None
163+
operations=execution.get_navigable_operations(),
164+
next_marker=None,
101165
)
166+
167+
168+
def _maybe_replay_cached(
169+
execution: Execution,
170+
checkpoint_token: str,
171+
client_token: str | None,
172+
) -> CheckpointOutput | None:
173+
"""Replay cached CheckpointOutput when the incoming call matches
174+
the last checkpoint. Returns ``None`` when there is no match.
175+
"""
176+
cached = CheckpointCore.match_cached(execution, checkpoint_token, client_token)
177+
if cached is None:
178+
return None
179+
return CheckpointOutput(
180+
checkpoint_token=cached.outbound_checkpoint_token,
181+
new_execution_state=CheckpointUpdatedExecutionState(
182+
operations=list(cached.operations),
183+
next_marker=cached.next_marker,
184+
),
185+
)

packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processors/step.py

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -61,16 +61,8 @@ def process(
6161
new_step_details = StepDetails(
6262
attempt=current_attempt + 1,
6363
next_attempt_timestamp=next_attempt_time,
64-
result=(
65-
current_op.step_details.result
66-
if current_op and current_op.step_details
67-
else None
68-
),
69-
error=(
70-
current_op.step_details.error
71-
if current_op and current_op.step_details
72-
else None
73-
),
64+
result=update.payload,
65+
error=update.error,
7466
)
7567

7668
# Create new operation with updated step_details
@@ -99,12 +91,12 @@ def process(
9991
else None,
10092
)
10193

102-
# Schedule step retry timer to fire after delay
103-
notifier.notify_step_retry_scheduled(
104-
execution_arn=execution_arn,
105-
operation_id=update.operation_id,
106-
delay=delay,
107-
)
94+
# Timer scheduling is handled centrally by
95+
# Executor._schedule_earliest_pending. The
96+
# processor only needs to set
97+
# step_details.next_attempt_timestamp; the executor
98+
# finds the earliest pending wake across all ops and
99+
# arms a single timer.
108100
return retry_operation
109101
case OperationAction.SUCCEED:
110102
return self._translate_update_to_operation(

packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processors/wait.py

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -74,12 +74,11 @@ def process(
7474
chained_invoke_details=None,
7575
)
7676

77-
# Schedule wait timer to complete after delay
78-
notifier.notify_wait_timer_scheduled(
79-
execution_arn=execution_arn,
80-
operation_id=update.operation_id,
81-
delay=scaled_wait_seconds,
82-
)
77+
# Timer scheduling is handled centrally by
78+
# Executor._schedule_earliest_pending. The processor
79+
# only needs to set wait_details.scheduled_end_timestamp;
80+
# the executor finds the earliest pending wake across
81+
# all ops and arms a single timer.
8382
return wait_operation
8483
case OperationAction.CANCEL:
8584
# TODO: need to cancel the WAIT in the executor

0 commit comments

Comments
 (0)