aws-durable-execution-sdk-python 1.0.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

aws_durable_execution_sdk_python/__about__.py

@@ -1,4 +1,4 @@
 # SPDX-FileCopyrightText: 2025-present Amazon.com, Inc. or its affiliates.
 #
 # SPDX-License-Identifier: Apache-2.0
-__version__ = "1.0.0"
+__version__ = "1.1.0"

aws_durable_execution_sdk_python/__init__.py

@@ -5,6 +5,7 @@
 from aws_durable_execution_sdk_python.context import (
     DurableContext,
     durable_step,
+    durable_wait_for_callback,
     durable_with_child_context,
 )
 
@@ -30,5 +31,6 @@ __all__ = [
     "ValidationError",
     "durable_execution",
     "durable_step",
+    "durable_wait_for_callback",
     "durable_with_child_context",
 ]

aws_durable_execution_sdk_python/concurrency/executor.py

@@ -22,6 +22,7 @@ from aws_durable_execution_sdk_python.concurrency.models import (
 )
 from aws_durable_execution_sdk_python.config import ChildConfig
 from aws_durable_execution_sdk_python.exceptions import (
+    OrphanedChildException,
     SuspendExecution,
     TimedSuspendExecution,
 )
@@ -198,42 +199,47 @@ class ConcurrentExecutor(ABC, Generic[CallableType, ResultType]):
             execution_state.create_checkpoint()
             submit_task(executable_with_state)
 
-        with (
-            TimerScheduler(resubmitter) as scheduler,
-            ThreadPoolExecutor(max_workers=max_workers) as thread_executor,
-        ):
-
-            def submit_task(executable_with_state: ExecutableWithState) -> Future:
-                """Submit task to the thread executor and mark its state as started."""
-                future = thread_executor.submit(
-                    self._execute_item_in_child_context,
-                    executor_context,
-                    executable_with_state.executable,
-                )
-                executable_with_state.run(future)
+        thread_executor = ThreadPoolExecutor(max_workers=max_workers)
+        try:
+            with TimerScheduler(resubmitter) as scheduler:
+
+                def submit_task(executable_with_state: ExecutableWithState) -> Future:
+                    """Submit task to the thread executor and mark its state as started."""
+                    future = thread_executor.submit(
+                        self._execute_item_in_child_context,
+                        executor_context,
+                        executable_with_state.executable,
+                    )
+                    executable_with_state.run(future)
 
-                def on_done(future: Future) -> None:
-                    self._on_task_complete(executable_with_state, future, scheduler)
+                    def on_done(future: Future) -> None:
+                        self._on_task_complete(executable_with_state, future, scheduler)
 
-                future.add_done_callback(on_done)
-                return future
+                    future.add_done_callback(on_done)
+                    return future
 
-            # Submit initial tasks
-            futures = [
-                submit_task(exe_state) for exe_state in self.executables_with_state
-            ]
+                # Submit initial tasks
+                futures = [
+                    submit_task(exe_state) for exe_state in self.executables_with_state
+                ]
 
-            # Wait for completion
-            self._completion_event.wait()
+                # Wait for completion
+                self._completion_event.wait()
 
-            # Cancel remaining futures so
-            # that we don't wait for them to join.
-            for future in futures:
-                future.cancel()
+                # Cancel futures that haven't started yet
+                for future in futures:
+                    future.cancel()
 
-            # Suspend execution if everything done and at least one of the tasks raised a suspend exception.
-            if self._suspend_exception:
-                raise self._suspend_exception
+                # Suspend execution if everything done and at least one of the tasks raised a suspend exception.
+                if self._suspend_exception:
+                    raise self._suspend_exception
+
+        finally:
+            # Shutdown without waiting for running threads for early return when
+            # completion criteria are met (e.g., min_successful).
+            # Running threads will continue in background but they raise OrphanedChildException
+            # on the next attempt to checkpoint.
+            thread_executor.shutdown(wait=False, cancel_futures=True)
 
         # Build final result
         return self._create_result()
@@ -291,6 +297,15 @@ class ConcurrentExecutor(ABC, Generic[CallableType, ResultType]):
             result = future.result()
             exe_state.complete(result)
             self.counters.complete_task()
+        except OrphanedChildException:
+            # Parent already completed and returned.
+            # State is already RUNNING, which _create_result() marked as STARTED
+            # Just log and exit - no state change needed
+            logger.debug(
+                "Terminating orphaned branch %s without error because parent has completed already",
+                exe_state.index,
+            )
+            return
         except TimedSuspendExecution as tse:
             exe_state.suspend_with_timeout(tse.scheduled_timestamp)
             scheduler.schedule_resume(exe_state, tse.scheduled_timestamp)

aws_durable_execution_sdk_python/context.py

@@ -24,17 +24,17 @@ from aws_durable_execution_sdk_python.identifier import OperationIdentifier
 from aws_durable_execution_sdk_python.lambda_service import OperationSubType
 from aws_durable_execution_sdk_python.logger import Logger, LogInfo
 from aws_durable_execution_sdk_python.operation.callback import (
-    create_callback_handler,
+    CallbackOperationExecutor,
     wait_for_callback_handler,
 )
 from aws_durable_execution_sdk_python.operation.child import child_handler
-from aws_durable_execution_sdk_python.operation.invoke import invoke_handler
+from aws_durable_execution_sdk_python.operation.invoke import InvokeOperationExecutor
 from aws_durable_execution_sdk_python.operation.map import map_handler
 from aws_durable_execution_sdk_python.operation.parallel import parallel_handler
-from aws_durable_execution_sdk_python.operation.step import step_handler
-from aws_durable_execution_sdk_python.operation.wait import wait_handler
+from aws_durable_execution_sdk_python.operation.step import StepOperationExecutor
+from aws_durable_execution_sdk_python.operation.wait import WaitOperationExecutor
 from aws_durable_execution_sdk_python.operation.wait_for_condition import (
-    wait_for_condition_handler,
+    WaitForConditionOperationExecutor,
 )
 from aws_durable_execution_sdk_python.serdes import (
     PassThroughSerDes,
@@ -104,6 +104,52 @@ def durable_with_child_context(
     return wrapper
 
 
+def durable_wait_for_callback(
+    func: Callable[Concatenate[str, WaitForCallbackContext, Params], T],
+) -> Callable[Params, Callable[[str, WaitForCallbackContext], T]]:
+    """Wrap your callable into a wait_for_callback submitter function.
+
+    This decorator allows you to define a submitter function with additional
+    parameters that will be bound when called.
+
+    Args:
+        func: A callable that takes callback_id, context, and additional parameters
+
+    Returns:
+        A wrapper function that binds the additional parameters and returns
+        a submitter function compatible with wait_for_callback
+
+    Example:
+        @durable_wait_for_callback
+        def submit_to_external_system(
+            callback_id: str,
+            context: WaitForCallbackContext,
+            task_name: str,
+            priority: int
+        ):
+            context.logger.info(f"Submitting {task_name} with callback {callback_id}")
+            external_api.submit_task(
+                task_name=task_name,
+                priority=priority,
+                callback_id=callback_id
+            )
+
+        # Usage in durable handler:
+        result = context.wait_for_callback(
+            submit_to_external_system("my_task", priority=5)
+        )
+    """
+
+    def wrapper(*args, **kwargs):
+        def submitter_with_arguments(callback_id: str, context: WaitForCallbackContext):
+            return func(callback_id, context, *args, **kwargs)
+
+        submitter_with_arguments._original_name = func.__name__  # noqa: SLF001
+        return submitter_with_arguments
+
+    return wrapper
+
+
 class Callback(Generic[T], CallbackProtocol[T]):  # noqa: PYI059
     """A future that will block on result() until callback_id returns."""
 
@@ -136,7 +182,7 @@ class Callback(Generic[T], CallbackProtocol[T]):  # noqa: PYI059
 
         if not checkpointed_result.is_existent():
             msg = "Callback operation must exist"
-            raise CallbackError(msg)
+            raise CallbackError(message=msg, callback_id=self.callback_id)
 
         if (
             checkpointed_result.is_failed()
@@ -144,7 +190,12 @@ class Callback(Generic[T], CallbackProtocol[T]):  # noqa: PYI059
             or checkpointed_result.is_timed_out()
             or checkpointed_result.is_stopped()
         ):
-            checkpointed_result.raise_callable_error()
+            msg = (
+                checkpointed_result.error.message
+                if checkpointed_result.error and checkpointed_result.error.message
+                else "Callback failed"
+            )
+            raise CallbackError(message=msg, callback_id=self.callback_id)
 
         if checkpointed_result.is_succeeded():
             if checkpointed_result.result is None:
@@ -272,13 +323,14 @@ class DurableContext(DurableContextProtocol):
         if not config:
             config = CallbackConfig()
         operation_id: str = self._create_step_id()
-        callback_id: str = create_callback_handler(
+        executor: CallbackOperationExecutor = CallbackOperationExecutor(
             state=self.state,
             operation_identifier=OperationIdentifier(
                 operation_id=operation_id, parent_id=self._parent_id, name=name
             ),
             config=config,
         )
+        callback_id: str = executor.process()
         result: Callback = Callback(
             callback_id=callback_id,
             operation_id=operation_id,
@@ -306,8 +358,10 @@ class DurableContext(DurableContextProtocol):
         Returns:
             The result of the invoked function
         """
+        if not config:
+            config = InvokeConfig[P, R]()
         operation_id = self._create_step_id()
-        result: R = invoke_handler(
+        executor: InvokeOperationExecutor[R] = InvokeOperationExecutor(
             function_name=function_name,
             payload=payload,
             state=self.state,
@@ -318,6 +372,7 @@ class DurableContext(DurableContextProtocol):
             ),
             config=config,
         )
+        result: R = executor.process()
         self.state.track_replay(operation_id=operation_id)
         return result
 
@@ -454,8 +509,10 @@ class DurableContext(DurableContextProtocol):
     ) -> T:
         step_name = self._resolve_step_name(name, func)
         logger.debug("Step name: %s", step_name)
+        if not config:
+            config = StepConfig()
         operation_id = self._create_step_id()
-        result: T = step_handler(
+        executor: StepOperationExecutor[T] = StepOperationExecutor(
             func=func,
             config=config,
             state=self.state,
@@ -466,6 +523,7 @@ class DurableContext(DurableContextProtocol):
             ),
             context_logger=self.logger,
         )
+        result: T = executor.process()
         self.state.track_replay(operation_id=operation_id)
         return result
 
@@ -481,8 +539,9 @@ class DurableContext(DurableContextProtocol):
             msg = "duration must be at least 1 second"
             raise ValidationError(msg)
         operation_id = self._create_step_id()
-        wait_handler(
-            seconds=seconds,
+        wait_seconds = duration.seconds
+        executor: WaitOperationExecutor = WaitOperationExecutor(
+            seconds=wait_seconds,
             state=self.state,
             operation_identifier=OperationIdentifier(
                 operation_id=operation_id,
@@ -490,6 +549,7 @@ class DurableContext(DurableContextProtocol):
                 name=name,
             ),
         )
+        executor.process()
         self.state.track_replay(operation_id=operation_id)
 
     def wait_for_callback(
@@ -533,17 +593,20 @@ class DurableContext(DurableContextProtocol):
             raise ValidationError(msg)
 
         operation_id = self._create_step_id()
-        result: T = wait_for_condition_handler(
-            check=check,
-            config=config,
-            state=self.state,
-            operation_identifier=OperationIdentifier(
-                operation_id=operation_id,
-                parent_id=self._parent_id,
-                name=name,
-            ),
-            context_logger=self.logger,
+        executor: WaitForConditionOperationExecutor[T] = (
+            WaitForConditionOperationExecutor(
+                check=check,
+                config=config,
+                state=self.state,
+                operation_identifier=OperationIdentifier(
+                    operation_id=operation_id,
+                    parent_id=self._parent_id,
+                    name=name,
+                ),
+                context_logger=self.logger,
+            )
         )
+        result: T = executor.process()
         self.state.track_replay(operation_id=operation_id)
         return result
 

aws_durable_execution_sdk_python/exceptions.py

@@ -372,3 +372,32 @@ class CallableRuntimeErrorSerializableDetails:
 
 class SerDesError(DurableExecutionsError):
     """Raised when serialization fails."""
+
+
+class OrphanedChildException(BaseException):
+    """Raised when a child operation attempts to checkpoint after its parent context has completed.
+
+    This exception inherits from BaseException (not Exception) so that user-space doesn't
+    accidentally catch it with broad exception handlers like 'except Exception'.
+
+    This exception will happen when a parallel branch or map item tries to create a checkpoint
+    after its parent context (i.e the parallel/map operation) has already completed due to meeting
+    completion criteria (e.g., min_successful reached, failure tolerance exceeded).
+
+    Although you cannot cancel running futures in user-space, this will at least terminate the
+    child operation on the next checkpoint attempt, preventing subsequent operations in the
+    child scope from executing.
+
+    Attributes:
+        operation_id: Operation ID of the orphaned child
+    """
+
+    def __init__(self, message: str, operation_id: str):
+        """Initialize OrphanedChildException.
+
+        Args:
+            message: Human-readable error message
+            operation_id: Operation ID of the orphaned child (required)
+        """
+        super().__init__(message)
+        self.operation_id = operation_id

aws_durable_execution_sdk_python/execution.py

@@ -237,7 +237,7 @@ def durable_execution(
             service_client = (
                 LambdaClient(client=boto3_client)
                 if boto3_client is not None
-                else LambdaClient.initialize_from_env()
+                else LambdaClient.initialize_client()
             )
 
         raw_input_payload: str | None = (

aws_durable_execution_sdk_python/lambda_service.py

@@ -2,10 +2,8 @@ from __future__ import annotations
 
 import datetime
 import logging
-import os
 from dataclasses import dataclass, field
 from enum import Enum
-from pathlib import Path
 from typing import TYPE_CHECKING, Any, Protocol, TypeAlias
 
 import boto3  # type: ignore
@@ -751,7 +749,7 @@ class Operation:
             callback_details = CallbackDetails.from_dict(callback_details_input)
 
         chained_invoke_details = None
-        if chained_invoke_details := data.get("chained_invoke_details"):
+        if chained_invoke_details := data.get("ChainedInvokeDetails"):
             chained_invoke_details = ChainedInvokeDetails.from_dict(
                 chained_invoke_details
             )
@@ -943,41 +941,14 @@ class LambdaClient(DurableServiceClient):
         self.client = client
 
     @staticmethod
-    def load_preview_botocore_models() -> None:
-        """
-        Load boto3 models from the Python path for custom preview client.
-        """
-        os.environ["AWS_DATA_PATH"] = str(
-            Path(__file__).parent.joinpath("botocore", "data")
+    def initialize_client() -> LambdaClient:
+        client = boto3.client(
+            "lambda",
+            config=Config(
+                connect_timeout=5,
+                read_timeout=50,
+            ),
         )
-
-    @staticmethod
-    def initialize_from_env() -> LambdaClient:
-        LambdaClient.load_preview_botocore_models()
-
-        """
-        TODO - we can remove this when were using the actual lambda client,
-        but we need this with the preview model because boto won't match against lambdainternal.
-        """
-        endpoint_url = os.getenv("AWS_ENDPOINT_URL_LAMBDA", None)
-        if not endpoint_url:
-            client = boto3.client(
-                "lambdainternal",
-                config=Config(
-                    connect_timeout=5,
-                    read_timeout=50,
-                ),
-            )
-        else:
-            client = boto3.client(
-                "lambdainternal",
-                endpoint_url=endpoint_url,
-                config=Config(
-                    connect_timeout=5,
-                    read_timeout=50,
-                ),
-            )
-
         return LambdaClient(client=client)
 
     def checkpoint(

aws_durable_execution_sdk_python/operation/base.py

@@ -0,0 +1,187 @@
+"""Base classes for operation executors with checkpoint response handling."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+from aws_durable_execution_sdk_python.exceptions import InvalidStateError
+
+if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python.state import CheckpointedResult
+
+T = TypeVar("T")
+
+
+@dataclass(frozen=True)
+class CheckResult(Generic[T]):
+    """Result of checking operation checkpoint status.
+
+    Encapsulates the outcome of checking an operation's status and determines
+    the next action in the operation execution flow.
+
+    IMPORTANT: Do not construct directly. Use factory methods:
+    - create_is_ready_to_execute(checkpoint) - operation ready to execute
+    - create_started() - checkpoint created, check status again
+    - create_completed(result) - terminal result available
+
+    Attributes:
+        is_ready_to_execute: True if the operation is ready to execute its logic
+        has_checkpointed_result: True if a terminal result is already available
+        checkpointed_result: Checkpoint data for execute()
+        deserialized_result: Final result when operation completed
+    """
+
+    is_ready_to_execute: bool
+    has_checkpointed_result: bool
+    checkpointed_result: CheckpointedResult | None = None
+    deserialized_result: T | None = None
+
+    @classmethod
+    def create_is_ready_to_execute(
+        cls, checkpoint: CheckpointedResult
+    ) -> CheckResult[T]:
+        """Create a CheckResult indicating the operation is ready to execute.
+
+        Args:
+            checkpoint: The checkpoint data to pass to execute()
+
+        Returns:
+            CheckResult with is_ready_to_execute=True
+        """
+        return cls(
+            is_ready_to_execute=True,
+            has_checkpointed_result=False,
+            checkpointed_result=checkpoint,
+        )
+
+    @classmethod
+    def create_started(cls) -> CheckResult[T]:
+        """Create a CheckResult signaling that a checkpoint was created.
+
+        Signals that process() should verify checkpoint status again to detect
+        if the operation completed already during checkpoint creation.
+
+        Returns:
+            CheckResult indicating process() should check status again
+        """
+        return cls(is_ready_to_execute=False, has_checkpointed_result=False)
+
+    @classmethod
+    def create_completed(cls, result: T) -> CheckResult[T]:
+        """Create a CheckResult with a terminal result already deserialized.
+
+        Args:
+            result: The final deserialized result
+
+        Returns:
+            CheckResult with has_checkpointed_result=True and deserialized_result set
+        """
+        return cls(
+            is_ready_to_execute=False,
+            has_checkpointed_result=True,
+            deserialized_result=result,
+        )
+
+
+class OperationExecutor(ABC, Generic[T]):
+    """Base class for durable operations with checkpoint response handling.
+
+    Provides a framework for implementing operations that check status after
+    creating START checkpoints to handle synchronous completion, avoiding
+    unnecessary execution or suspension.
+
+    The common pattern:
+    1. Check operation status
+    2. Create START checkpoint if needed
+    3. Check status again (detects synchronous completion)
+    4. Execute operation logic when ready
+
+    Subclasses must implement:
+    - check_result_status(): Check status, create checkpoint if needed, return next action
+    - execute(): Execute the operation logic with checkpoint data
+    """
+
+    @abstractmethod
+    def check_result_status(self) -> CheckResult[T]:
+        """Check operation status and create START checkpoint if needed.
+
+        Called twice by process() when creating synchronous checkpoints: once before
+        and once after, to detect if the operation completed immediately.
+
+        This method should:
+        1. Get the current checkpoint result
+        2. Check for terminal statuses (SUCCEEDED, FAILED, etc.) and handle them
+        3. Check for pending statuses and suspend if needed
+        4. Create a START checkpoint if the operation hasn't started
+        5. Return a CheckResult indicating the next action
+
+        Returns:
+            CheckResult indicating whether to:
+            - Return a terminal result (has_checkpointed_result=True)
+            - Execute operation logic (is_ready_to_execute=True)
+            - Check status again (neither flag set - checkpoint was just created)
+
+        Raises:
+            Operation-specific exceptions for terminal failure states
+            SuspendExecution for pending states
+        """
+        ...  # pragma: no cover
+
+    @abstractmethod
+    def execute(self, checkpointed_result: CheckpointedResult) -> T:
+        """Execute operation logic with checkpoint data.
+
+        This method is called when the operation is ready to execute its core logic.
+        It receives the checkpoint data that was returned by check_result_status().
+
+        Args:
+            checkpointed_result: The checkpoint data containing operation state
+
+        Returns:
+            The result of executing the operation
+
+        Raises:
+            May raise operation-specific errors during execution
+        """
+        ...  # pragma: no cover
+
+    def process(self) -> T:
+        """Process operation with checkpoint response handling.
+
+        Orchestrates the double-check pattern:
+        1. Check status (handles replay and existing checkpoints)
+        2. If checkpoint was just created, check status again (detects synchronous completion)
+        3. Return terminal result if available
+        4. Execute operation logic if ready
+        5. Raise error for invalid states
+
+        Returns:
+            The final result of the operation
+
+        Raises:
+            InvalidStateError: If the check result is in an invalid state
+            May raise operation-specific errors from check_result_status() or execute()
+        """
+        # Check 1: Entry (handles replay and existing checkpoints)
+        result = self.check_result_status()
+
+        # If checkpoint was created, verify checkpoint response for immediate status change
+        if not result.is_ready_to_execute and not result.has_checkpointed_result:
+            result = self.check_result_status()
+
+        # Return terminal result if available (can be None for operations that return None)
+        if result.has_checkpointed_result:
+            return result.deserialized_result  # type: ignore[return-value]
+
+        # Execute operation logic
+        if result.is_ready_to_execute:
+            if result.checkpointed_result is None:
+                msg = "CheckResult is marked ready to execute but checkpointed result is not set."
+                raise InvalidStateError(msg)
+            return self.execute(result.checkpointed_result)
+
+        # Invalid state - neither terminal nor ready to execute
+        msg = "Invalid CheckResult state: neither terminal nor ready to execute"
+        raise InvalidStateError(msg)