Coverage for async-durable-execution/src/async_durable_execution/primitive/step.py: 98%
137 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 16:54 +0000
1"""Implement the Durable step operation."""
3from __future__ import annotations
5import asyncio
6import logging
7from dataclasses import dataclass
8from enum import Enum
9from typing import TYPE_CHECKING, TypeVar, cast
11from .base import (
12 OperationExecutor,
13 OperationContext,
14)
15from .child import get_durable_context
16from ..config import Duration, RetryStrategy, duration_to_seconds
17from ..context import bind_current_context, get_current_context
18from ..exceptions import (
19 CallableRuntimeError,
20 ExecutionError,
21 InvocationError,
22 TerminationReason,
23 suspend_with_optional_resume_delay,
24 suspend_with_optional_resume_timestamp,
25)
26from ..models import (
27 ErrorObject,
28 Operation,
29 OperationIdentifier,
30 OperationStatus,
31 OperationUpdate,
32 OperationSubType,
33)
34from ..task import create_eager_task
36if TYPE_CHECKING:
37 from collections.abc import Awaitable, Callable
39 from .child import DurableContext
40 from ..serdes import SerDes
41 from ..state import ExecutionState
43logger = logging.getLogger(__name__)
45T = TypeVar("T")
48class StepInterruptedError(InvocationError):
49 """Raised when a step is interrupted before it checkpointed at the end."""
51 def __init__(self, message: str, step_id: str | None = None):
52 super().__init__(message, TerminationReason.STEP_INTERRUPTED)
53 self.step_id = step_id
56class StepSemantics(Enum):
57 """Checkpoint timing guarantees for a durable step attempt."""
59 AT_MOST_ONCE_PER_RETRY = "AT_MOST_ONCE_PER_RETRY"
60 AT_LEAST_ONCE_PER_RETRY = "AT_LEAST_ONCE_PER_RETRY"
63class StepOperationExecutor(OperationExecutor[T]):
64 """Executor for step operations."""
66 def __init__(
67 self,
68 func: Callable[[], Awaitable[T]],
69 state: ExecutionState,
70 operation_identifier: OperationIdentifier,
71 retry_strategy: Callable[[Exception, int], Duration | None] | None = None,
72 step_semantics: StepSemantics = StepSemantics.AT_LEAST_ONCE_PER_RETRY,
73 serdes: SerDes | None = None,
74 ):
75 """Initialize the step operation executor.
77 Args:
78 func: The step function to execute
79 state: The execution state
80 operation_identifier: The operation identifier
81 retry_strategy: Optional retry strategy for step failures
82 step_semantics: Checkpoint timing guarantee for the step attempt
83 serdes: Optional serializer/deserializer for the step result
84 """
85 super().__init__(state=state, operation_identifier=operation_identifier)
86 self.func = func
87 self.retry_strategy = retry_strategy
88 self.step_semantics = step_semantics
89 self.serdes = serdes
91 async def start(self) -> T:
92 """Start a new step operation."""
93 start_operation: OperationUpdate = OperationUpdate.create_step_start(
94 identifier=self.operation_identifier,
95 )
96 is_sync: bool = self.step_semantics is StepSemantics.AT_MOST_ONCE_PER_RETRY
97 await self.create_checkpoint(start_operation, is_sync=is_sync)
99 return await self.execute(None)
101 async def replay(self, operation: Operation) -> T:
102 """Replay an existing step operation from its checkpoint."""
103 if operation.status is OperationStatus.SUCCEEDED:
104 logger.debug(
105 "Step already completed, skipping execution for id: %s, name: %s",
106 self.operation_identifier.operation_id,
107 self.operation_name,
108 )
109 result_payload = (
110 operation.step_details.result if operation.step_details else None
111 )
112 if result_payload is None:
113 return cast("T", None)
115 result: T = await self.deserialize_value(
116 data=result_payload,
117 serdes=self.serdes,
118 )
119 return result
121 if operation.status is OperationStatus.FAILED:
122 self._raise_callable_error(operation)
124 if operation.status is OperationStatus.PENDING:
125 scheduled_timestamp = (
126 operation.step_details.next_attempt_timestamp
127 if operation.step_details
128 else None
129 )
130 suspend_with_optional_resume_timestamp(
131 msg=f"Retry scheduled for {self.operation_name or self.operation_identifier.operation_id} will retry at timestamp {scheduled_timestamp}",
132 datetime_timestamp=scheduled_timestamp,
133 )
135 if (
136 operation.status is OperationStatus.STARTED
137 and self.step_semantics is StepSemantics.AT_MOST_ONCE_PER_RETRY
138 ):
139 msg: str = f"Step operation_id={self.operation_identifier.operation_id} name={self.operation_identifier.name} was previously interrupted"
140 await self.retry_handler(StepInterruptedError(msg), operation)
141 self._raise_callable_error(operation)
143 if (
144 operation.status is OperationStatus.STARTED
145 and self.step_semantics is StepSemantics.AT_LEAST_ONCE_PER_RETRY
146 ):
147 return await self.execute(operation)
149 if operation.status is OperationStatus.READY:
150 start_operation: OperationUpdate = OperationUpdate.create_step_start(
151 identifier=self.operation_identifier,
152 )
153 is_sync: bool = self.step_semantics is StepSemantics.AT_MOST_ONCE_PER_RETRY
154 await self.create_checkpoint(start_operation, is_sync=is_sync)
156 return await self.execute(operation)
158 return await self.execute(operation)
160 async def execute(self, operation: Operation | None) -> T:
161 """Execute step function with error handling and retry logic.
163 Args:
164 operation: The checkpointed operation state, if any
166 Returns:
167 The result of executing the step function
169 Raises:
170 ExecutionError: For fatal errors that should not be retried
171 May raise other exceptions that will be handled by retry_handler
172 """
173 # Get current attempt - checkpointed attempts + 1
174 attempt: int = 1
175 if operation and operation.step_details:
176 attempt = operation.step_details.attempt + 1
178 step_context: StepContext = StepContext(
179 attempt=attempt,
180 execution_state=self.state,
181 operation_identifier=self.operation_identifier,
182 )
184 try:
185 # This is the actual code provided by the caller to execute durably inside the step
186 with bind_current_context(step_context):
187 raw_result = await self.func()
189 serialized_result: str = await self.serialize_value(
190 value=raw_result,
191 serdes=self.serdes,
192 )
194 success_operation: OperationUpdate = OperationUpdate.create_step_succeed(
195 identifier=self.operation_identifier,
196 payload=serialized_result,
197 )
199 # Checkpoint SUCCEED operation with blocking (is_sync=True, default).
200 # Must ensure the success state is persisted before returning the result to the caller.
201 # This guarantees the step result is durable and won't be lost if Lambda terminates.
202 await self.create_checkpoint(success_operation)
204 logger.debug(
205 "✅ Successfully completed step for id: %s, name: %s",
206 self.operation_identifier.operation_id,
207 self.operation_identifier.name,
208 )
209 return await self.deserialize_value( # noqa: TRY300
210 data=serialized_result,
211 serdes=self.serdes,
212 )
213 except Exception as e:
214 if isinstance(e, ExecutionError):
215 # No retry on fatal - e.g checkpoint exception
216 logger.debug(
217 "💥 Fatal error for id: %s, name: %s",
218 self.operation_identifier.operation_id,
219 self.operation_identifier.name,
220 )
221 # This bubbles up to execution.durable_execution, where it will exit with FAILED
222 raise
224 logger.exception(
225 "❌ failed step for id: %s, name: %s",
226 self.operation_identifier.operation_id,
227 self.operation_identifier.name,
228 )
230 await self.retry_handler(e, operation)
231 # If we've failed to raise an exception from the retry_handler, then we are in a
232 # weird state, and should crash terminate the execution
233 msg = "retry handler should have raised an exception, but did not."
234 raise ExecutionError(msg) from None
236 async def retry_handler(
237 self,
238 error: Exception,
239 operation: Operation | None,
240 ):
241 """Checkpoint and suspend for replay if retry required, otherwise raise error.
243 Args:
244 error: The exception that occurred during step execution
245 operation: The checkpointed operation state, if any
247 Raises:
248 SuspendExecution: If retry is scheduled
249 StepInterruptedError: If the error is a StepInterruptedError
250 CallableRuntimeError: If retry is exhausted or error is not retryable
251 """
252 error_object = ErrorObject.from_exception(error)
254 retry_strategy = self.retry_strategy or RetryStrategy.default()
256 retry_attempt: int = (
257 operation.step_details.attempt
258 if operation and operation.step_details
259 else 0
260 )
261 delay_seconds: int | None = None
262 try:
263 retry_delay = retry_strategy(error, retry_attempt + 1)
264 if retry_delay is not None:
265 delay_seconds = duration_to_seconds(retry_delay, "retry delay")
266 except Exception as retry_error:
267 fail_error_object = ErrorObject.from_exception(retry_error)
268 fail_operation: OperationUpdate = OperationUpdate.create_step_fail(
269 identifier=self.operation_identifier, error=fail_error_object
270 )
272 # Checkpoint FAIL operation with blocking (is_sync=True, default).
273 # Must ensure the failure state is persisted before raising the exception.
274 # This guarantees the error is durable and the step won't be retried on replay.
275 await self.create_checkpoint(fail_operation)
277 if isinstance(retry_error, StepInterruptedError):
278 raise retry_error
280 raise CallableRuntimeError.from_error_object(fail_error_object)
282 if retry_delay is None:
283 fail_operation = OperationUpdate.create_step_fail(
284 identifier=self.operation_identifier, error=error_object
285 )
286 await self.create_checkpoint(fail_operation)
288 if isinstance(error, StepInterruptedError):
289 raise error
291 raise CallableRuntimeError.from_error_object(error_object)
293 assert delay_seconds is not None
295 logger.debug(
296 "Retrying step for id: %s, name: %s, attempt: %s",
297 self.operation_identifier.operation_id,
298 self.operation_identifier.name,
299 retry_attempt + 1,
300 )
302 # Because we are issuing a retry and create an OperationUpdate, enforce a
303 # minimum delay of one second here to match model behavior.
304 if delay_seconds < 1:
305 logger.warning(
306 (
307 "Retry delay_seconds step for id: %s, name: %s,"
308 "attempt: %s is %d < 1. Setting to minimum of 1 seconds."
309 ),
310 self.operation_identifier.operation_id,
311 self.operation_identifier.name,
312 retry_attempt + 1,
313 delay_seconds,
314 )
315 delay_seconds = 1
317 retry_operation: OperationUpdate = OperationUpdate.create_step_retry(
318 identifier=self.operation_identifier,
319 error=error_object,
320 next_attempt_delay_seconds=delay_seconds,
321 )
323 # Checkpoint RETRY operation with blocking (is_sync=True, default).
324 # Must ensure retry state is persisted before suspending execution.
325 # This guarantees the retry attempt count and next attempt timestamp are durable.
326 await self.create_checkpoint(retry_operation)
328 suspend_with_optional_resume_delay(
329 msg=(
330 f"Retry scheduled for {self.operation_identifier.operation_id} "
331 f"in {delay_seconds} seconds"
332 ),
333 delay_seconds=delay_seconds,
334 )
336 @staticmethod
337 def _raise_callable_error(operation: Operation) -> None:
338 error = operation.step_details.error if operation.step_details else None
339 if error is None:
340 msg = "Unknown error. No ErrorObject exists on the Checkpoint Operation."
341 raise CallableRuntimeError(
342 message=msg,
343 error_type=None,
344 data=None,
345 stack_trace=None,
346 )
348 raise CallableRuntimeError.from_error_object(error)
351def step(
352 func: Callable[[], Awaitable[T]],
353 *,
354 name: str | None = None,
355 retry_strategy: Callable[[Exception, int], Duration | None] | None = None,
356 step_semantics: StepSemantics = StepSemantics.AT_LEAST_ONCE_PER_RETRY,
357 serdes: SerDes | None = None,
358) -> asyncio.Task[T]:
359 """Run user code as a checkpointed durable step.
361 Durable steps are the main way to isolate non-deterministic work such as API
362 calls, clock reads, UUID generation, and database access from replayed code.
363 """
364 context = get_durable_context()
365 step_name = name if name is not None else getattr(func, "__name__", None)
366 logger.debug("Step name: %s", step_name)
368 with context._replay_aware(executes_user_code=True):
369 operation_id = context.step_counter.create_step_id()
370 operation_identifier = OperationIdentifier(
371 operation_id=operation_id,
372 sub_type=OperationSubType.STEP,
373 parent_id=context.parent_id,
374 name=step_name,
375 )
377 return create_eager_task(
378 lambda: _step(
379 func=func,
380 context=context,
381 operation_identifier=operation_identifier,
382 retry_strategy=retry_strategy,
383 step_semantics=step_semantics,
384 serdes=serdes,
385 ),
386 )
389async def _step(
390 func: Callable[[], Awaitable[T]],
391 *,
392 context: DurableContext,
393 operation_identifier: OperationIdentifier,
394 retry_strategy: Callable[[Exception, int], Duration | None] | None = None,
395 step_semantics: StepSemantics = StepSemantics.AT_LEAST_ONCE_PER_RETRY,
396 serdes: SerDes | None = None,
397) -> T:
398 executor: StepOperationExecutor[T] = StepOperationExecutor(
399 func=func,
400 state=context.execution_state,
401 operation_identifier=operation_identifier,
402 retry_strategy=retry_strategy,
403 step_semantics=step_semantics,
404 serdes=serdes,
405 )
406 return await executor.process()
409@dataclass(frozen=True)
410class StepContext(OperationContext):
411 """Context exposed while a step function is executing."""
413 attempt: int | None = None
416def get_step_context() -> StepContext:
417 """Return the active `StepContext`."""
418 current_context = get_current_context()
419 if not isinstance(current_context, StepContext):
420 msg = "get_step_context() can only be used while a step function is executing."
421 raise RuntimeError(msg)
422 return current_context