async-durable-execution-runner 2.0.0a1__py3-none-any.whl

async_durable_execution_runner/checkpoint/processors/context.py

@@ -0,0 +1,59 @@
+"""Context operation processor for handling CONTEXT operation updates."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from async_durable_execution.lambda_service import (
+    Operation,
+    OperationAction,
+    OperationStatus,
+    OperationUpdate,
+)
+
+from async_durable_execution_runner.checkpoint.processors.base import (
+    OperationProcessor,
+)
+from async_durable_execution_runner.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+if TYPE_CHECKING:
+    from async_durable_execution_runner.observer import ExecutionNotifier
+
+
+class ContextProcessor(OperationProcessor):
+    """Processes CONTEXT operation updates for execution context management."""
+
+    def process(
+        self,
+        update: OperationUpdate,
+        current_op: Operation | None,
+        notifier: ExecutionNotifier,  # noqa: ARG002
+        execution_arn: str,  # noqa: ARG002
+    ) -> Operation:
+        """Process CONTEXT operation update for context state transitions."""
+        match update.action:
+            case OperationAction.START:
+                # TODO: check for "Cannot start a CONTEXT operation that already exists."
+                return self._translate_update_to_operation(
+                    update=update,
+                    current_operation=current_op,
+                    status=OperationStatus.STARTED,
+                )
+            case OperationAction.SUCCEED:
+                return self._translate_update_to_operation(
+                    update=update,
+                    current_operation=current_op,
+                    status=OperationStatus.SUCCEEDED,
+                )
+            case OperationAction.FAIL:
+                return self._translate_update_to_operation(
+                    update=update,
+                    current_operation=current_op,
+                    status=OperationStatus.FAILED,
+                )
+            case _:
+                msg: str = "Invalid action for CONTEXT operation."
+                raise InvalidParameterValueException(msg)

async_durable_execution_runner/checkpoint/processors/execution.py

@@ -0,0 +1,52 @@
+"""Execution operation processor for handling EXECUTION operation updates."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from async_durable_execution.lambda_service import (
+    ErrorObject,
+    Operation,
+    OperationAction,
+    OperationUpdate,
+)
+
+from async_durable_execution_runner.checkpoint.processors.base import (
+    OperationProcessor,
+)
+
+
+if TYPE_CHECKING:
+    from async_durable_execution_runner.observer import ExecutionNotifier
+
+
+class ExecutionProcessor(OperationProcessor):
+    """Processes EXECUTION operation updates for workflow completion."""
+
+    def process(
+        self,
+        update: OperationUpdate,
+        current_op: Operation | None,  # noqa: ARG002
+        notifier: ExecutionNotifier,
+        execution_arn: str,
+    ) -> Operation | None:
+        """Process EXECUTION operation update for workflow completion/failure."""
+        match update.action:
+            case OperationAction.SUCCEED:
+                notifier.notify_completed(
+                    execution_arn=execution_arn, result=update.payload
+                )
+            case _:
+                # intentional. actual service will fail any EXECUTION update that is not SUCCEED.
+                error = (
+                    update.error
+                    if update.error
+                    else ErrorObject.from_message(
+                        "There is no error details but EXECUTION checkpoint action is not SUCCEED."
+                    )
+                )
+                # All EXECUTION failures go through normal fail path
+                # Timeout/Stop status is set by executor based on the operation that caused it
+                notifier.notify_failed(execution_arn=execution_arn, error=error)
+        # TODO: Svc doesn't actually create checkpoint for EXECUTION. might have to for localrunner though.
+        return None

async_durable_execution_runner/checkpoint/processors/step.py

@@ -0,0 +1,124 @@
+"""Step operation processor for handling STEP operation updates."""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime, timedelta
+from typing import TYPE_CHECKING
+
+from async_durable_execution.lambda_service import (
+    Operation,
+    OperationAction,
+    OperationStatus,
+    OperationUpdate,
+    StepDetails,
+)
+
+from async_durable_execution_runner.checkpoint.processors.base import (
+    OperationProcessor,
+)
+from async_durable_execution_runner.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+if TYPE_CHECKING:
+    from async_durable_execution_runner.observer import ExecutionNotifier
+
+
+class StepProcessor(OperationProcessor):
+    """Processes STEP operation updates with retry scheduling."""
+
+    def process(
+        self,
+        update: OperationUpdate,
+        current_op: Operation | None,
+        notifier: ExecutionNotifier,
+        execution_arn: str,
+    ) -> Operation:
+        """Process STEP operation update with scheduler integration for retries."""
+        match update.action:
+            case OperationAction.START:
+                return self._translate_update_to_operation(
+                    update=update,
+                    current_operation=current_op,
+                    status=OperationStatus.STARTED,
+                )
+            case OperationAction.RETRY:
+                # set Status=PENDING, next attempt time, attempt count + 1
+                delay = (
+                    update.step_options.next_attempt_delay_seconds
+                    if update.step_options
+                    else 0
+                )
+                next_attempt_time = datetime.now(UTC) + timedelta(seconds=delay)
+
+                # Build new step_details with incremented attempt
+                current_attempt = (
+                    current_op.step_details.attempt
+                    if current_op and current_op.step_details
+                    else 0
+                )
+                new_step_details = StepDetails(
+                    attempt=current_attempt + 1,
+                    next_attempt_timestamp=next_attempt_time,
+                    result=(
+                        current_op.step_details.result
+                        if current_op and current_op.step_details
+                        else None
+                    ),
+                    error=(
+                        current_op.step_details.error
+                        if current_op and current_op.step_details
+                        else None
+                    ),
+                )
+
+                # Create new operation with updated step_details
+                retry_operation = Operation(
+                    operation_id=update.operation_id,
+                    operation_type=update.operation_type,
+                    status=OperationStatus.PENDING,
+                    parent_id=update.parent_id,
+                    name=update.name,
+                    start_timestamp=(
+                        current_op.start_timestamp if current_op else datetime.now(UTC)
+                    ),
+                    end_timestamp=None,
+                    sub_type=update.sub_type,
+                    execution_details=current_op.execution_details
+                    if current_op
+                    else None,
+                    context_details=current_op.context_details if current_op else None,
+                    step_details=new_step_details,
+                    wait_details=current_op.wait_details if current_op else None,
+                    callback_details=current_op.callback_details
+                    if current_op
+                    else None,
+                    chained_invoke_details=current_op.chained_invoke_details
+                    if current_op
+                    else None,
+                )
+
+                # Schedule step retry timer to fire after delay
+                notifier.notify_step_retry_scheduled(
+                    execution_arn=execution_arn,
+                    operation_id=update.operation_id,
+                    delay=delay,
+                )
+                return retry_operation
+            case OperationAction.SUCCEED:
+                return self._translate_update_to_operation(
+                    update=update,
+                    current_operation=current_op,
+                    status=OperationStatus.SUCCEEDED,
+                )
+            case OperationAction.FAIL:
+                return self._translate_update_to_operation(
+                    update=update,
+                    current_operation=current_op,
+                    status=OperationStatus.FAILED,
+                )
+            case _:
+                msg: str = "Invalid action for STEP operation."
+
+                raise InvalidParameterValueException(msg)

async_durable_execution_runner/checkpoint/processors/wait.py

@@ -0,0 +1,95 @@
+"""Wait operation processor for handling WAIT operation updates."""
+
+from __future__ import annotations
+
+import logging
+import os
+from datetime import UTC, datetime, timedelta
+from typing import TYPE_CHECKING
+
+from async_durable_execution.lambda_service import (
+    Operation,
+    OperationAction,
+    OperationStatus,
+    OperationUpdate,
+    WaitDetails,
+)
+
+from async_durable_execution_runner.checkpoint.processors.base import (
+    OperationProcessor,
+)
+from async_durable_execution_runner.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+if TYPE_CHECKING:
+    from async_durable_execution_runner.observer import ExecutionNotifier
+
+
+class WaitProcessor(OperationProcessor):
+    """Processes WAIT operation updates with timer scheduling."""
+
+    def process(
+        self,
+        update: OperationUpdate,
+        current_op: Operation | None,
+        notifier: ExecutionNotifier,
+        execution_arn: str,
+    ) -> Operation:
+        """Process WAIT operation update with scheduler integration for timers."""
+        match update.action:
+            case OperationAction.START:
+                wait_seconds = (
+                    update.wait_options.wait_seconds if update.wait_options else 0
+                )
+                time_scale = float(os.getenv("DURABLE_EXECUTION_TIME_SCALE", "1.0"))
+                logging.info("Using DURABLE_EXECUTION_TIME_SCALE: %f", time_scale)
+                scaled_wait_seconds = wait_seconds * time_scale
+
+                scheduled_end_timestamp = datetime.now(UTC) + timedelta(
+                    seconds=scaled_wait_seconds
+                )
+
+                # Create WaitDetails with scheduled timestamp
+                wait_details = WaitDetails(
+                    scheduled_end_timestamp=scheduled_end_timestamp
+                )
+
+                # Create new operation with wait details
+                wait_operation = Operation(
+                    operation_id=update.operation_id,
+                    operation_type=update.operation_type,
+                    status=OperationStatus.STARTED,
+                    parent_id=update.parent_id,
+                    name=update.name,
+                    start_timestamp=datetime.now(UTC),
+                    end_timestamp=None,
+                    sub_type=update.sub_type,
+                    execution_details=None,
+                    context_details=None,
+                    step_details=None,
+                    wait_details=wait_details,
+                    callback_details=None,
+                    chained_invoke_details=None,
+                )
+
+                # Schedule wait timer to complete after delay
+                notifier.notify_wait_timer_scheduled(
+                    execution_arn=execution_arn,
+                    operation_id=update.operation_id,
+                    delay=scaled_wait_seconds,
+                )
+                return wait_operation
+            case OperationAction.CANCEL:
+                # TODO: need to cancel the WAIT in the executor
+                # TODO: increase sequence id
+                return self._translate_update_to_operation(
+                    update=update,
+                    current_operation=current_op,
+                    status=OperationStatus.CANCELLED,
+                )
+            case _:
+                msg: str = "Invalid action for WAIT operation."
+
+                raise InvalidParameterValueException(msg)

async_durable_execution_runner/checkpoint/transformer.py

@@ -0,0 +1,104 @@
+"""Operation transformer for converting OperationUpdates to Operations."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from async_durable_execution.lambda_service import (
+    Operation,
+    OperationType,
+    OperationUpdate,
+)
+
+from async_durable_execution_runner.checkpoint.processors.callback import (
+    CallbackProcessor,
+)
+from async_durable_execution_runner.checkpoint.processors.context import (
+    ContextProcessor,
+)
+from async_durable_execution_runner.checkpoint.processors.execution import (
+    ExecutionProcessor,
+)
+from async_durable_execution_runner.checkpoint.processors.step import (
+    StepProcessor,
+)
+from async_durable_execution_runner.checkpoint.processors.wait import (
+    WaitProcessor,
+)
+from async_durable_execution_runner.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from async_durable_execution_runner.checkpoint.processors.base import (
+        OperationProcessor,
+    )
+
+from typing import ClassVar
+
+
+class OperationTransformer:
+    """Transforms OperationUpdates to Operations while maintaining order and triggering scheduler actions."""
+
+    _DEFAULT_PROCESSORS: ClassVar[dict[OperationType, OperationProcessor]] = {
+        OperationType.STEP: StepProcessor(),
+        OperationType.WAIT: WaitProcessor(),
+        OperationType.CONTEXT: ContextProcessor(),
+        OperationType.CALLBACK: CallbackProcessor(),
+        OperationType.EXECUTION: ExecutionProcessor(),
+    }
+
+    def __init__(
+        self,
+        processors: MutableMapping[OperationType, OperationProcessor] | None = None,
+    ):
+        self.processors = processors if processors else self._DEFAULT_PROCESSORS
+
+    def process_updates(
+        self,
+        updates: list[OperationUpdate],
+        current_operations: list[Operation],
+        notifier,
+        execution_arn: str,
+    ) -> tuple[list[Operation], list[OperationUpdate]]:
+        """Transform updates maintaining operation order and return (operations, updates)."""
+        op_map = {op.operation_id: op for op in current_operations}
+
+        # Start with copy of current operations list
+        result_operations = current_operations.copy()
+
+        for update in updates:
+            processor = self.processors.get(update.operation_type)
+            if processor:
+                current_op = op_map.get(update.operation_id)
+                updated_op = processor.process(
+                    update=update,
+                    current_op=current_op,
+                    notifier=notifier,
+                    execution_arn=execution_arn,
+                )
+
+                if updated_op is not None:
+                    if update.operation_id in op_map:
+                        # Update existing operation in-place
+                        for i, op in enumerate(result_operations):  # pragma: no branch
+                            # no branch coverage because result_operation empty not reachable here
+                            if op.operation_id == update.operation_id:
+                                result_operations[i] = updated_op
+                                break
+                    else:
+                        # Append new operation to end
+                        result_operations.append(updated_op)
+
+                    # Update map for future lookups
+                    op_map[update.operation_id] = updated_op
+            else:
+                msg: str = (
+                    f"Checkpoint for {update.operation_type} is not implemented yet."
+                )
+                raise InvalidParameterValueException(msg)
+
+        return result_operations, updates

async_durable_execution_runner/checkpoint/validators/__init__.py

@@ -0,0 +1 @@
+"""Checkpoint validation module."""

async_durable_execution_runner/checkpoint/validators/checkpoint.py

@@ -0,0 +1,242 @@
+"""Main checkpoint input validator."""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING
+
+from async_durable_execution.lambda_service import (
+    OperationAction,
+    OperationType,
+    OperationUpdate,
+)
+
+from async_durable_execution_runner.checkpoint.validators.operations.callback import (
+    CallbackOperationValidator,
+)
+from async_durable_execution_runner.checkpoint.validators.operations.context import (
+    ContextOperationValidator,
+)
+from async_durable_execution_runner.checkpoint.validators.operations.execution import (
+    ExecutionOperationValidator,
+)
+from async_durable_execution_runner.checkpoint.validators.operations.invoke import (
+    ChainedInvokeOperationValidator,
+)
+from async_durable_execution_runner.checkpoint.validators.operations.step import (
+    StepOperationValidator,
+)
+from async_durable_execution_runner.checkpoint.validators.operations.wait import (
+    WaitOperationValidator,
+)
+from async_durable_execution_runner.checkpoint.validators.transitions import (
+    ValidActionsByOperationTypeValidator,
+)
+from async_durable_execution_runner.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from async_durable_execution_runner.execution import Execution
+
+MAX_ERROR_PAYLOAD_SIZE_BYTES = 32768
+
+
+class CheckpointValidator:
+    """Validates checkpoint input based on current state."""
+
+    @staticmethod
+    def validate_input(updates: list[OperationUpdate], execution: Execution) -> None:
+        """Perform validation on the given input based on the current state."""
+        if not updates:
+            return
+
+        CheckpointValidator._validate_conflicting_execution_update(updates)
+        CheckpointValidator._validate_parent_id_and_duplicate_id(updates, execution)
+
+        for update in updates:
+            CheckpointValidator._validate_operation_update(update, execution)
+
+    @staticmethod
+    def _validate_conflicting_execution_update(updates: list[OperationUpdate]) -> None:
+        """Validate that there are no conflicting execution updates."""
+        execution_updates = [
+            update
+            for update in updates
+            if update.operation_type == OperationType.EXECUTION
+        ]
+
+        if len(execution_updates) > 1:
+            msg_multiple_exec: str = "Cannot checkpoint multiple EXECUTION updates."
+
+            raise InvalidParameterValueException(msg_multiple_exec)
+
+        if execution_updates and updates[-1].operation_type != OperationType.EXECUTION:
+            msg_exec_last: str = "EXECUTION checkpoint must be the last update."
+
+            raise InvalidParameterValueException(msg_exec_last)
+
+    @staticmethod
+    def _validate_operation_update(
+        update: OperationUpdate, execution: Execution
+    ) -> None:
+        """Validate a single operation update."""
+        CheckpointValidator._validate_inconsistent_operation_metadata(update, execution)
+        CheckpointValidator._validate_payload_sizes(update)
+        ValidActionsByOperationTypeValidator.validate(
+            update.operation_type, update.action
+        )
+        CheckpointValidator._validate_operation_status_transition(update, execution)
+
+    @staticmethod
+    def _validate_payload_sizes(update: OperationUpdate) -> None:
+        """Validate that operation payload sizes are not too large."""
+        if update.error is not None:
+            payload = json.dumps(update.error.to_dict())
+            if len(payload) > MAX_ERROR_PAYLOAD_SIZE_BYTES:
+                msg: str = f"Error object size must be less than {MAX_ERROR_PAYLOAD_SIZE_BYTES} bytes."
+                raise InvalidParameterValueException(msg)
+
+    @staticmethod
+    def _validate_operation_status_transition(
+        update: OperationUpdate, execution: Execution
+    ) -> None:
+        """Validate that the operation status transition is valid."""
+        current_state = None
+        for operation in execution.operations:
+            if operation.operation_id == update.operation_id:
+                current_state = operation
+                break
+
+        match update.operation_type:
+            case OperationType.STEP:
+                StepOperationValidator.validate(current_state, update)
+            case OperationType.CONTEXT:
+                ContextOperationValidator.validate(current_state, update)
+            case OperationType.WAIT:
+                WaitOperationValidator.validate(current_state, update)
+            case OperationType.CALLBACK:
+                CallbackOperationValidator.validate(current_state, update)
+            case OperationType.CHAINED_INVOKE:
+                ChainedInvokeOperationValidator.validate(current_state, update)
+            case OperationType.EXECUTION:
+                ExecutionOperationValidator.validate(update)
+            case _:  # pragma: no cover
+                msg: str = "Invalid operation type."
+
+                raise InvalidParameterValueException(msg)
+
+    @staticmethod
+    def _validate_inconsistent_operation_metadata(
+        update: OperationUpdate, execution: Execution
+    ) -> None:
+        """Validate that operation metadata is consistent with existing operation."""
+        current_state = None
+        for operation in execution.operations:
+            if operation.operation_id == update.operation_id:
+                current_state = operation
+                break
+
+        if current_state is not None:
+            if (
+                update.operation_type is not None
+                and update.operation_type != current_state.operation_type
+            ):
+                msg: str = "Inconsistent operation type."
+                raise InvalidParameterValueException(msg)
+
+            if (
+                update.sub_type is not None
+                and update.sub_type != current_state.sub_type
+            ):
+                msg_subtype: str = "Inconsistent operation subtype."
+                raise InvalidParameterValueException(msg_subtype)
+
+            if update.name is not None and update.name != current_state.name:
+                msg_name: str = "Inconsistent operation name."
+                raise InvalidParameterValueException(msg_name)
+
+            if (
+                update.parent_id is not None
+                and update.parent_id != current_state.parent_id
+            ):
+                msg_parent: str = "Inconsistent parent operation id."
+                raise InvalidParameterValueException(msg_parent)
+
+    @staticmethod
+    def _validate_parent_id_and_duplicate_id(
+        updates: list[OperationUpdate], execution: Execution
+    ) -> None:
+        """Validate parent IDs and check for duplicate operation IDs.
+
+        Validate that any provided parentId is valid, and also validate no duplicate operation is being
+        updated at the same time (unless it is a STEP/CONTEXT starting + performing one more non-START action).
+        """
+        operations_started: MutableMapping[str, OperationUpdate] = {}
+        last_updates_seen: MutableMapping[str, OperationUpdate] = {}
+
+        for update in updates:
+            if CheckpointValidator._is_invalid_duplicate_update(
+                update, last_updates_seen
+            ):
+                msg_duplicate: str = (
+                    "Cannot checkpoint multiple operations with the same ID."
+                )
+                raise InvalidParameterValueException(msg_duplicate)
+
+            if not CheckpointValidator._is_valid_parent_for_update(
+                execution, update, operations_started
+            ):
+                msg_parent: str = "Invalid parent operation id."
+                raise InvalidParameterValueException(msg_parent)
+
+            if update.action == OperationAction.START:
+                operations_started[update.operation_id] = update
+
+            last_updates_seen[update.operation_id] = update
+
+    @staticmethod
+    def _is_invalid_duplicate_update(
+        update: OperationUpdate, last_updates_seen: MutableMapping[str, OperationUpdate]
+    ) -> bool:
+        """Check if this is an invalid duplicate update."""
+        last_update = last_updates_seen.get(update.operation_id)
+        if last_update is None:
+            return False
+
+        if last_update.operation_type in (OperationType.STEP, OperationType.CONTEXT):
+            # Allow duplicate for STEP/CONTEXT if last was START and current is not START
+            allow_duplicate = (
+                last_update.action == OperationAction.START
+                and update.action != OperationAction.START
+            )
+            return not allow_duplicate
+
+        return True
+
+    @staticmethod
+    def _is_valid_parent_for_update(
+        execution: Execution,
+        update: OperationUpdate,
+        operations_started: MutableMapping[str, OperationUpdate],
+    ) -> bool:
+        """Check if the parent ID is valid for the update."""
+        parent_id = update.parent_id
+
+        if parent_id is None:
+            return True
+
+        # Check if parent is in operations started in this batch
+        if parent_id in operations_started:
+            parent_update = operations_started[parent_id]
+            return parent_update.operation_type == OperationType.CONTEXT
+
+        # Check if parent exists in current execution state
+        for operation in execution.operations:
+            if operation.operation_id == parent_id:
+                return operation.operation_type == OperationType.CONTEXT
+
+        return False

async_durable_execution_runner/checkpoint/validators/operations/__init__.py

@@ -0,0 +1 @@
+"""Operation-specific validators."""