Coverage for async_durable_execution/_primitive/step.py: 88%
263 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +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 OperationExecutor
12from .._core import (
13 CallableRuntimeError,
14 Duration,
15 DurableContext,
16 ErrorObject,
17 ExecutionError,
18 ExecutionState,
19 InvocationError,
20 Operation,
21 OperationContext,
22 OperationIdentifier,
23 OperationStatus,
24 OperationType,
25 OperationUpdate,
26 RetryStrategy,
27 SerDes,
28 TerminationReason,
29 _encode_sdk_control_error_data,
30 _restore_sdk_control_error,
31 bind_current_context,
32 duration_to_seconds,
33 get_current_context,
34 suspend_with_optional_resume_delay,
35 suspend_with_optional_resume_timestamp,
36)
38if TYPE_CHECKING:
39 from collections.abc import Awaitable, Callable
41 from ..extension import (
42 ExtensionStepFunction,
43 ExtensionStepResult,
44 ExtensionStepRetryStrategy,
45 )
47logger = logging.getLogger(__name__)
49T = TypeVar("T")
52def _error_object_from_exception(
53 error: Exception,
54 *,
55 invocation_retryable: bool | None = None,
56) -> ErrorObject:
57 error_object = ErrorObject.from_exception(error)
58 sdk_error_data = _encode_sdk_control_error_data(
59 error,
60 invocation_retryable=invocation_retryable,
61 )
62 if sdk_error_data is None:
63 return error_object
64 return ErrorObject(
65 message=error_object.message,
66 type=error_object.type,
67 data=sdk_error_data,
68 stack_trace=error_object.stack_trace,
69 )
72class StepInterruptedError(InvocationError):
73 """Raised when a step is interrupted before it checkpointed at the end."""
75 def __init__(self, message: str, step_id: str | None = None) -> None:
76 super().__init__(message, TerminationReason.STEP_INTERRUPTED)
77 self.step_id = step_id
80class StepSemantics(Enum):
81 """Checkpoint timing guarantees for a durable step attempt."""
83 AT_MOST_ONCE_PER_RETRY = "AT_MOST_ONCE_PER_RETRY"
84 AT_LEAST_ONCE_PER_RETRY = "AT_LEAST_ONCE_PER_RETRY"
87class StepOperationExecutor(OperationExecutor[T]):
88 """Executor for step operations."""
90 SERDES_OPERATION_TYPE = OperationType.STEP
92 def __init__(
93 self,
94 func: Callable[[], Awaitable[T]],
95 state: ExecutionState,
96 operation_identifier: OperationIdentifier,
97 retry_strategy: Callable[[Exception, int], Duration | None] | None = None,
98 step_semantics: StepSemantics = StepSemantics.AT_LEAST_ONCE_PER_RETRY,
99 serdes: SerDes | None = None,
100 ) -> None:
101 """Initialize the step operation executor.
103 Args:
104 func: The step function to execute
105 state: The execution state
106 operation_identifier: The operation identifier
107 retry_strategy: Optional retry strategy for step failures
108 step_semantics: Checkpoint timing guarantee for the step attempt
109 serdes: Optional serializer/deserializer for the step result
110 """
111 super().__init__(state=state, operation_identifier=operation_identifier)
112 self.func = func
113 self.retry_strategy = retry_strategy
114 self.step_semantics = step_semantics
115 self.serdes = serdes
117 async def start(self) -> T:
118 """Start a new step operation."""
119 start_operation: OperationUpdate = OperationUpdate.create_step_start(
120 identifier=self.operation_identifier,
121 )
122 is_sync: bool = self.step_semantics is StepSemantics.AT_MOST_ONCE_PER_RETRY
123 await self.create_checkpoint(start_operation, is_sync=is_sync)
125 return await self.execute(None)
127 async def replay(self, operation: Operation) -> T:
128 """Replay an existing step operation from its checkpoint."""
129 if operation.status is OperationStatus.SUCCEEDED:
130 logger.debug(
131 "Step already completed, skipping execution for id: %s, name: %s",
132 self.operation_identifier.operation_id,
133 self.operation_name,
134 )
135 result_payload = (
136 operation.step_details.result if operation.step_details else None
137 )
138 if result_payload is None:
139 return cast("T", None)
141 result: T = await self.deserialize_value(
142 data=result_payload,
143 serdes=self.serdes,
144 operation=operation,
145 attempt=(
146 operation.step_details.attempt
147 if operation.step_details is not None
148 else None
149 ),
150 )
151 return result
153 if operation.status is OperationStatus.FAILED:
154 self._raise_callable_error(operation)
156 if operation.status is OperationStatus.PENDING:
157 scheduled_timestamp = (
158 operation.step_details.next_attempt_timestamp
159 if operation.step_details
160 else None
161 )
162 suspend_with_optional_resume_timestamp(
163 msg=f"Retry scheduled for {self.operation_name or self.operation_identifier.operation_id} will retry at timestamp {scheduled_timestamp}",
164 datetime_timestamp=scheduled_timestamp,
165 )
167 if (
168 operation.status is OperationStatus.STARTED
169 and self.step_semantics is StepSemantics.AT_MOST_ONCE_PER_RETRY
170 ):
171 msg: str = f"Step operation_id={self.operation_identifier.operation_id} name={self.operation_identifier.name} was previously interrupted"
172 await self.retry_handler(StepInterruptedError(msg), operation)
173 self._raise_callable_error(operation)
175 if (
176 operation.status is OperationStatus.STARTED
177 and self.step_semantics is StepSemantics.AT_LEAST_ONCE_PER_RETRY
178 ):
179 return await self.execute(operation)
181 if operation.status is OperationStatus.READY:
182 start_operation: OperationUpdate = OperationUpdate.create_step_start(
183 identifier=self.operation_identifier,
184 )
185 is_sync: bool = self.step_semantics is StepSemantics.AT_MOST_ONCE_PER_RETRY
186 await self.create_checkpoint(start_operation, is_sync=is_sync)
188 return await self.execute(operation)
190 return await self.execute(operation)
192 async def execute(self, operation: Operation | None) -> T:
193 """Execute step function with error handling and retry logic.
195 Args:
196 operation: The checkpointed operation state, if any
198 Returns:
199 The result of executing the step function
201 Raises:
202 ExecutionError: For fatal errors that should not be retried
203 May raise other exceptions that will be handled by retry_handler
204 """
205 # Get current attempt - checkpointed attempts + 1
206 attempt: int = 1
207 if operation and operation.step_details:
208 attempt = operation.step_details.attempt + 1
210 step_context: StepContext = StepContext(
211 attempt=attempt,
212 execution_state=self.state,
213 operation_identifier=self.operation_identifier,
214 )
216 try:
217 # This is the actual code provided by the caller to execute durably inside the step
218 with bind_current_context(step_context):
219 raw_result = await self.func()
221 serialized_result: str = await self.serialize_value(
222 value=raw_result,
223 serdes=self.serdes,
224 attempt=attempt,
225 )
227 success_operation: OperationUpdate = OperationUpdate.create_step_succeed(
228 identifier=self.operation_identifier,
229 payload=serialized_result,
230 )
232 # Checkpoint SUCCEED operation with blocking (is_sync=True, default).
233 # Must ensure the success state is persisted before returning the result to the caller.
234 # This guarantees the step result is durable and won't be lost if Lambda terminates.
235 await self.create_checkpoint(success_operation)
237 logger.debug(
238 "✅ Successfully completed step for id: %s, name: %s",
239 self.operation_identifier.operation_id,
240 self.operation_identifier.name,
241 )
242 except Exception as e:
243 if isinstance(e, ExecutionError):
244 # No retry on fatal - e.g checkpoint exception
245 logger.debug(
246 "💥 Fatal error for id: %s, name: %s",
247 self.operation_identifier.operation_id,
248 self.operation_identifier.name,
249 )
250 # This bubbles up to execution.durable_execution, where it will exit with FAILED
251 raise
253 logger.exception(
254 "❌ failed step for id: %s, name: %s",
255 self.operation_identifier.operation_id,
256 self.operation_identifier.name,
257 )
259 await self.retry_handler(e, operation)
260 # If we've failed to raise an exception from the retry_handler, then we are in a
261 # weird state, and should crash terminate the execution
262 msg = "retry handler should have raised an exception, but did not."
263 raise ExecutionError(msg) from None
265 # SUCCEED is already durable. A retryable read failure must retry the
266 # invocation and replay this completed step, not checkpoint RETRY/FAIL
267 # after the terminal operation transition.
268 return await self.deserialize_value(
269 data=serialized_result,
270 serdes=self.serdes,
271 attempt=attempt,
272 )
274 async def retry_handler(
275 self,
276 error: Exception,
277 operation: Operation | None,
278 ) -> None:
279 """Checkpoint and suspend for replay if retry required, otherwise raise error.
281 Args:
282 error: The exception that occurred during step execution
283 operation: The checkpointed operation state, if any
285 Raises:
286 SuspendExecution: If retry is scheduled
287 StepInterruptedError: If the error is a StepInterruptedError
288 CallableRuntimeError: If retry is exhausted or error is not retryable
289 """
290 error_object = _error_object_from_exception(error)
292 retry_strategy = self.retry_strategy or RetryStrategy.default()
294 retry_attempt: int = (
295 operation.step_details.attempt
296 if operation and operation.step_details
297 else 0
298 )
299 delay_seconds: int | None = None
300 try:
301 retry_delay = retry_strategy(error, retry_attempt + 1)
302 if retry_delay is not None:
303 delay_seconds = duration_to_seconds(retry_delay, "retry delay")
304 except Exception as retry_error:
305 fail_error_object = _error_object_from_exception(
306 retry_error,
307 invocation_retryable=False,
308 )
309 fail_operation: OperationUpdate = OperationUpdate.create_step_fail(
310 identifier=self.operation_identifier, error=fail_error_object
311 )
313 # Checkpoint FAIL operation with blocking (is_sync=True, default).
314 # Must ensure the failure state is persisted before raising the exception.
315 # This guarantees the error is durable and the step won't be retried on replay.
316 await self.create_checkpoint(fail_operation)
318 if isinstance(retry_error, StepInterruptedError):
319 raise retry_error
321 raise CallableRuntimeError.from_error_object(fail_error_object)
323 if retry_delay is None:
324 fail_error_object = _error_object_from_exception(
325 error,
326 invocation_retryable=False,
327 )
328 fail_operation = OperationUpdate.create_step_fail(
329 identifier=self.operation_identifier, error=fail_error_object
330 )
331 await self.create_checkpoint(fail_operation)
333 if isinstance(error, StepInterruptedError): 333 ↛ 334line 333 didn't jump to line 334 because the condition on line 333 was never true
334 raise error
336 raise CallableRuntimeError.from_error_object(fail_error_object)
338 assert delay_seconds is not None
340 logger.debug(
341 "Retrying step for id: %s, name: %s, attempt: %s",
342 self.operation_identifier.operation_id,
343 self.operation_identifier.name,
344 retry_attempt + 1,
345 )
347 # Because we are issuing a retry and create an OperationUpdate, enforce a
348 # minimum delay of one second here to match model behavior.
349 if delay_seconds < 1:
350 logger.warning(
351 (
352 "Retry delay_seconds step for id: %s, name: %s,"
353 "attempt: %s is %d < 1. Setting to minimum of 1 seconds."
354 ),
355 self.operation_identifier.operation_id,
356 self.operation_identifier.name,
357 retry_attempt + 1,
358 delay_seconds,
359 )
360 delay_seconds = 1
362 retry_operation: OperationUpdate = OperationUpdate.create_step_retry(
363 identifier=self.operation_identifier,
364 error=error_object,
365 next_attempt_delay_seconds=delay_seconds,
366 )
368 # Checkpoint RETRY operation with blocking (is_sync=True, default).
369 # Must ensure retry state is persisted before suspending execution.
370 # This guarantees the retry attempt count and next attempt timestamp are durable.
371 await self.create_checkpoint(retry_operation)
373 suspend_with_optional_resume_delay(
374 msg=(
375 f"Retry scheduled for {self.operation_identifier.operation_id} "
376 f"in {delay_seconds} seconds"
377 ),
378 delay_seconds=delay_seconds,
379 )
381 @staticmethod
382 def _raise_callable_error(operation: Operation) -> None:
383 error = operation.step_details.error if operation.step_details else None
384 if error is None:
385 msg = "Unknown error. No ErrorObject exists on the Checkpoint Operation."
386 raise CallableRuntimeError(
387 message=msg,
388 error_type=None,
389 data=None,
390 stack_trace=None,
391 )
393 raise CallableRuntimeError.from_error_object(error)
396class StatefulStepOperationExecutor(OperationExecutor[T]):
397 """Executor for stateful extension-authored step operations."""
399 SERDES_OPERATION_TYPE = OperationType.STEP
401 def __init__(
402 self,
403 func: ExtensionStepFunction[T],
404 state: ExecutionState,
405 operation_identifier: OperationIdentifier,
406 *,
407 initial_state: T | None,
408 retry_strategy: ExtensionStepRetryStrategy[T] | None,
409 step_semantics: StepSemantics,
410 serdes: SerDes[T] | None,
411 raise_original_error: bool = False,
412 ) -> None:
413 super().__init__(state=state, operation_identifier=operation_identifier)
414 self.func = func
415 self.initial_state = initial_state
416 self.retry_strategy = retry_strategy
417 self.step_semantics = step_semantics
418 self.serdes = serdes
419 self.raise_original_error = raise_original_error
421 async def start(self) -> T:
422 start = OperationUpdate.create_step_start(self.operation_identifier)
423 await self.create_checkpoint(
424 start,
425 is_sync=self.step_semantics is StepSemantics.AT_MOST_ONCE_PER_RETRY,
426 )
427 return await self._execute(None)
429 async def replay(self, operation: Operation) -> T:
430 if operation.status is OperationStatus.SUCCEEDED:
431 payload = operation.step_details.result if operation.step_details else None
432 if payload is None: 432 ↛ 433line 432 didn't jump to line 433 because the condition on line 432 was never true
433 return cast("T", None)
434 return await self.deserialize_value(
435 payload,
436 self.serdes,
437 operation=operation,
438 attempt=(
439 operation.step_details.attempt
440 if operation.step_details is not None
441 else None
442 ),
443 )
445 if operation.status is OperationStatus.FAILED:
446 self._raise_failed_operation(operation)
448 if operation.status is OperationStatus.PENDING: 448 ↛ 449line 448 didn't jump to line 449 because the condition on line 448 was never true
449 resume_at = (
450 operation.step_details.next_attempt_timestamp
451 if operation.step_details
452 else None
453 )
454 suspend_with_optional_resume_timestamp(
455 msg=f"Extension step {self.operation_name or self.operation_id} is pending",
456 datetime_timestamp=resume_at,
457 )
459 if (
460 operation.status is OperationStatus.STARTED
461 and self.step_semantics is StepSemantics.AT_MOST_ONCE_PER_RETRY
462 ):
463 msg = (
464 f"Extension step operation_id={self.operation_id} "
465 f"name={self.operation_name} was previously interrupted"
466 )
467 attempt = self._attempt(operation)
468 state = self.initial_state
469 try:
470 state = await self._load_state(operation)
471 except InvocationError as error:
472 if error.is_retryable():
473 raise
474 return await self._handle_failure(error, state, attempt)
475 except Exception as error:
476 return await self._handle_failure(error, state, attempt)
477 return await self._handle_failure(
478 StepInterruptedError(msg, self.operation_id),
479 state,
480 attempt,
481 )
483 if operation.status is OperationStatus.READY:
484 start = OperationUpdate.create_step_start(self.operation_identifier)
485 await self.create_checkpoint(
486 start,
487 is_sync=self.step_semantics is StepSemantics.AT_MOST_ONCE_PER_RETRY,
488 )
490 return await self._execute(operation)
492 async def _execute(self, operation: Operation | None) -> T:
493 from ..extension import ExtensionStepResult
495 state = self.initial_state
496 attempt = self._attempt(operation)
497 try:
498 state = await self._load_state(operation)
499 step_context = StepContext(
500 attempt=attempt,
501 execution_state=self.state,
502 operation_identifier=self.operation_identifier,
503 )
504 with bind_current_context(step_context):
505 outcome = await self.func(state)
507 if not isinstance(outcome, ExtensionStepResult): 507 ↛ 508line 507 didn't jump to line 508 because the condition on line 507 was never true
508 msg = (
509 "Extension step functions must return "
510 "ExtensionStepResult.succeed(...) or ExtensionStepResult.retry(...)"
511 )
512 raise TypeError(msg)
514 if outcome.is_retry:
515 delay_seconds, payload = await self._prepare_retry(outcome, attempt)
516 else:
517 payload = await self.serialize_value(
518 outcome.value,
519 self.serdes,
520 attempt=attempt,
521 )
522 except InvocationError as error:
523 if error.is_retryable(): 523 ↛ 524line 523 didn't jump to line 524 because the condition on line 523 was never true
524 raise
525 return await self._handle_failure(error, state, attempt)
526 except Exception as error:
527 return await self._handle_failure(error, state, attempt)
529 if outcome.is_retry:
530 return await self._schedule_retry(
531 delay_seconds=delay_seconds,
532 payload=payload,
533 )
535 await self.create_checkpoint(
536 OperationUpdate.create_step_succeed(
537 self.operation_identifier,
538 payload,
539 )
540 )
541 return await self.deserialize_value(
542 payload,
543 self.serdes,
544 attempt=attempt,
545 )
547 async def _load_state(self, operation: Operation | None) -> T | None:
548 if (
549 operation is not None
550 and operation.step_details is not None
551 and operation.step_details.result is not None
552 ):
553 return await self.deserialize_value(
554 operation.step_details.result,
555 self.serdes,
556 operation=operation,
557 attempt=operation.step_details.attempt,
558 )
559 return self.initial_state
561 @staticmethod
562 def _attempt(operation: Operation | None) -> int:
563 if operation is None or operation.step_details is None:
564 return 1
565 return operation.step_details.attempt + 1
567 async def _handle_failure(
568 self,
569 error: Exception,
570 state: T | None,
571 attempt: int,
572 ) -> T:
573 from ..extension import ExtensionStepResult
575 if self.retry_strategy is None:
576 return await self._fail(error)
578 try:
579 decision = self.retry_strategy(error, state, attempt)
580 except Exception as retry_error:
581 return await self._fail(retry_error)
583 if decision is None: 583 ↛ 584line 583 didn't jump to line 584 because the condition on line 583 was never true
584 return await self._fail(error)
585 if not isinstance(decision, ExtensionStepResult) or not decision.is_retry: 585 ↛ 586line 585 didn't jump to line 586 because the condition on line 585 was never true
586 msg = (
587 "Extension step retry_strategy must return "
588 "ExtensionStepResult.retry(...) or None"
589 )
590 return await self._fail(TypeError(msg))
592 try:
593 delay_seconds, payload = await self._prepare_retry(decision, attempt)
594 except InvocationError as retry_error:
595 if retry_error.is_retryable(): 595 ↛ 597line 595 didn't jump to line 597 because the condition on line 595 was always true
596 raise
597 return await self._fail(retry_error)
598 except Exception as retry_error:
599 return await self._fail(retry_error)
601 return await self._schedule_retry(
602 delay_seconds=delay_seconds,
603 payload=payload,
604 )
606 async def _prepare_retry(
607 self,
608 outcome: ExtensionStepResult[T],
609 attempt: int,
610 ) -> tuple[int, str]:
611 assert outcome.retry_delay is not None
612 delay_seconds = max(
613 1,
614 duration_to_seconds(outcome.retry_delay, "retry delay"),
615 )
616 payload = await self.serialize_value(
617 outcome.value,
618 self.serdes,
619 attempt=attempt,
620 )
621 return delay_seconds, payload
623 async def _schedule_retry(
624 self,
625 *,
626 delay_seconds: int,
627 payload: str,
628 ) -> T:
629 await self.create_checkpoint(
630 OperationUpdate.create_step_retry(
631 self.operation_identifier,
632 next_attempt_delay_seconds=delay_seconds,
633 payload=payload,
634 error=None,
635 )
636 )
637 suspend_with_optional_resume_delay(
638 msg=f"Extension step {self.operation_id} will retry",
639 delay_seconds=delay_seconds,
640 )
641 raise AssertionError("suspend_with_optional_resume_delay must raise")
643 async def _fail(self, error: Exception) -> T:
644 error_object = _error_object_from_exception(
645 error,
646 invocation_retryable=False,
647 )
648 await self.create_checkpoint(
649 OperationUpdate.create_step_fail(
650 self.operation_identifier,
651 error_object,
652 )
653 )
654 if self.raise_original_error:
655 raise error
656 control_error = _restore_sdk_control_error(
657 error_object.message or str(error),
658 error_object.type,
659 error_object.data,
660 )
661 if control_error is not None:
662 raise control_error
663 raise CallableRuntimeError.from_error_object(error_object)
665 @staticmethod
666 def _raise_failed_operation(operation: Operation) -> None:
667 error = operation.step_details.error if operation.step_details else None
668 if error is None: 668 ↛ 669line 668 didn't jump to line 669 because the condition on line 668 was never true
669 error = ErrorObject.from_message(
670 "Unknown error. No ErrorObject exists on the checkpoint operation."
671 )
672 control_error = _restore_sdk_control_error(
673 error.message or "Extension step failed",
674 error.type,
675 error.data,
676 )
677 if control_error is not None: 677 ↛ 679line 677 didn't jump to line 679 because the condition on line 677 was always true
678 raise control_error
679 raise CallableRuntimeError.from_error_object(error)
682def step(
683 func: Callable[[], Awaitable[T]],
684 *,
685 name: str | None = None,
686 retry_strategy: Callable[[Exception, int], Duration | None] | None = None,
687 step_semantics: StepSemantics = StepSemantics.AT_LEAST_ONCE_PER_RETRY,
688 serdes: SerDes | None = None,
689) -> asyncio.Task[T]:
690 """Compatibility import for the canonical operation-layer helper."""
691 from .._operation.step import step as operation_step
693 return operation_step(
694 func,
695 name=name,
696 retry_strategy=retry_strategy,
697 step_semantics=step_semantics,
698 serdes=serdes,
699 )
702async def _step(
703 func: Callable[[], Awaitable[T]],
704 *,
705 context: DurableContext,
706 operation_identifier: OperationIdentifier,
707 retry_strategy: Callable[[Exception, int], Duration | None] | None = None,
708 step_semantics: StepSemantics = StepSemantics.AT_LEAST_ONCE_PER_RETRY,
709 serdes: SerDes | None = None,
710) -> T:
711 executor: StepOperationExecutor[T] = StepOperationExecutor(
712 func=func,
713 state=context.execution_state,
714 operation_identifier=operation_identifier,
715 retry_strategy=retry_strategy,
716 step_semantics=step_semantics,
717 serdes=serdes,
718 )
719 return await executor.process()
722async def _stateful_step(
723 func: ExtensionStepFunction[T],
724 *,
725 context: DurableContext,
726 operation_identifier: OperationIdentifier,
727 initial_state: T | None,
728 retry_strategy: ExtensionStepRetryStrategy[T] | None,
729 step_semantics: StepSemantics,
730 serdes: SerDes[T] | None,
731 raise_original_error: bool = False,
732) -> T:
733 executor: StatefulStepOperationExecutor[T] = StatefulStepOperationExecutor(
734 func=func,
735 state=context.execution_state,
736 operation_identifier=operation_identifier,
737 initial_state=initial_state,
738 retry_strategy=retry_strategy,
739 step_semantics=step_semantics,
740 serdes=serdes,
741 raise_original_error=raise_original_error,
742 )
743 return await executor.process()
746@dataclass(frozen=True)
747class StepContext(OperationContext):
748 """Context exposed while a step function is executing."""
750 attempt: int | None = None
753def get_step_context() -> StepContext:
754 """Return the active `StepContext`."""
755 current_context = get_current_context()
756 if not isinstance(current_context, StepContext): 756 ↛ 757line 756 didn't jump to line 757 because the condition on line 756 was never true
757 msg = "get_step_context() can only be used while a step function is executing."
758 raise RuntimeError(msg)
759 return current_context