aws-durable-execution-sdk-python 2.0.0__py3-none-any.whl → 2.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__ = "2.0.0"
+__version__ = "2.1.0"

aws_durable_execution_sdk_python/concurrency/executor.py

@@ -2,12 +2,14 @@
 
 from __future__ import annotations
 
+import contextvars
 import heapq
 import logging
 import queue
 import time
 from collections import deque
 from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Generic, TypeVar, cast
 
 from aws_durable_execution_sdk_python.concurrency.models import (
@@ -43,7 +45,7 @@ from aws_durable_execution_sdk_python.identifier import (
     OperationIdentifier,
     OperationIdNamespace,
 )
-from aws_durable_execution_sdk_python.lambda_service import ErrorObject
+from aws_durable_execution_sdk_python.lambda_service import ErrorObject, OperationType
 from aws_durable_execution_sdk_python.operation.child import child_handler
 
 
@@ -56,6 +58,7 @@ if TYPE_CHECKING:
     from aws_durable_execution_sdk_python.state import (
         CheckpointedResult,
         ExecutionState,
+        ScheduledRefresh,
     )
 
 
@@ -86,6 +89,20 @@ def _branch_error_object(err: Exception) -> ErrorObject:
     return ErrorObject.from_exception(err)
 
 
+@dataclass
+class ResumeWave(Generic[CallableType, ResultType]):
+    """Branches suspended until one time, and the refresh that resumes them.
+
+    The refresh is a delayed empty checkpoint requested when the first branch
+    suspends. Its response shows the waits complete, so the wave resumes one
+    round trip after its time.
+    """
+
+    resume_at: float
+    refresh: ScheduledRefresh
+    branches: list[Branch[CallableType, ResultType]] = field(default_factory=list)
+
+
 class ConcurrentExecutor(Generic[CallableType, ResultType]):
     """Execute durable operations concurrently. This contains the execution logic for Map and Parallel.
 
@@ -167,6 +184,7 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
                 executable.index
             ),
             sub_type=self.sub_type_iteration,
+            operation_type=OperationType.CONTEXT,
             parent_id=executor_context._parent_id,  # noqa: SLF001
             name=self.get_iteration_name(executable.index),
         )
@@ -252,7 +270,9 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
 
         events: queue.Queue[BranchEvent[ResultType]] = queue.Queue()
         pending: deque[Branch[CallableType, ResultType]] = deque(self.branches)
-        timed_resumes: list[tuple[float, int]] = []
+        # The heap orders resume times. The dict groups the branches under each.
+        resume_times: list[float] = []
+        waves: dict[float, ResumeWave[CallableType, ResultType]] = {}
         branch_by_index: dict[int, Branch[CallableType, ResultType]] = {
             branch.index: branch for branch in self.branches
         }
@@ -273,13 +293,25 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
 
         def submit(branch: Branch[CallableType, ResultType]) -> None:
             branch.start()
-            pool.submit(
-                self._branch_worker,
-                execution_state,
-                executor_context,
-                events,
-                branch.executable,
-            )
+            if execution_state._plugin_executor._plugins:  # noqa: SLF001
+                # Every admission/resume gets its own Context. Branch bindings
+                # cannot leak into siblings, the coordinator, or reused workers.
+                pool.submit(
+                    contextvars.copy_context().run,
+                    self._branch_worker,
+                    execution_state,
+                    executor_context,
+                    events,
+                    branch.executable,
+                )
+            else:
+                pool.submit(
+                    self._branch_worker,
+                    execution_state,
+                    executor_context,
+                    events,
+                    branch.executable,
+                )
 
         try:
             # Only rebuild the items snapshot after a terminal event changes
@@ -317,19 +349,20 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
                     running += 1
                     needs_snapshot_rebuild = True
 
-                # Resume due timed suspends in-process. One checkpoint
-                # refresh serves the whole due wave; a failure is terminal
+                # Resume due waves in-process. A refresh failure is terminal
                 # for the execution and propagates from this thread.
                 now: float = time.time()
-                due: list[Branch[CallableType, ResultType]] = []
-                while timed_resumes and timed_resumes[0][0] <= now:
-                    _, index = heapq.heappop(timed_resumes)
-                    due.append(branch_by_index[index])
-                if due:
-                    execution_state.create_checkpoint()
-                    for branch in due:
+                resumed = False
+                while resume_times and resume_times[0] <= now:
+                    wave = waves.pop(heapq.heappop(resume_times))
+                    wave.refresh.wait()
+                    # Branches joined the wave in event order. Resume in index
+                    # order so scheduling stays deterministic.
+                    for branch in sorted(wave.branches, key=lambda b: b.index):
                         submit(branch)
                         running += 1
+                    resumed = True
+                if resumed:
                     continue
 
                 if running == 0:
@@ -340,18 +373,18 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
                         raise retryable_error
                     # Every in-flight branch is suspended and no slot is
                     # free (or no work remains): suspend the parent.
-                    if timed_resumes:
+                    if resume_times:
                         raise TimedSuspendExecution(
                             "All concurrent work complete or suspended pending retry.",
-                            timed_resumes[0][0],
+                            resume_times[0],
                         )
                     raise SuspendExecution(
                         "All concurrent work complete or suspended and pending external callback."
                     )
 
                 timeout: float | None = None
-                if timed_resumes:
-                    timeout = max(timed_resumes[0][0] - time.time(), 0)
+                if resume_times:
+                    timeout = max(resume_times[0] - time.time(), 0)
                 try:
                     event: BranchEvent[ResultType] = events.get(timeout=timeout)
                 except queue.Empty:
@@ -383,7 +416,13 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
                         needs_snapshot_rebuild = True
                     case BranchEventKind.SUSPENDED_UNTIL if event.resume_at is not None:
                         applied.suspend_until(event.resume_at)
-                        heapq.heappush(timed_resumes, (event.resume_at, event.index))
+                        if event.resume_at not in waves:
+                            waves[event.resume_at] = ResumeWave(
+                                event.resume_at,
+                                execution_state.schedule_refresh(event.resume_at),
+                            )
+                            heapq.heappush(resume_times, event.resume_at)
+                        waves[event.resume_at].branches.append(applied)
                         running -= 1
                         needs_snapshot_rebuild = True
                     case BranchEventKind.ORPHANED:
@@ -402,6 +441,9 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
                         msg = f"Unhandled branch event: {event}"
                         raise InvalidStateError(msg)
         finally:
+            # Nothing will wait for a refresh whose wave never resumed.
+            for wave in waves.values():
+                wave.refresh.cancel()
             # Shutdown without waiting for running threads for early return
             # when completion criteria are met (e.g., min_successful).
             # Running threads continue in the background of this invocation

aws_durable_execution_sdk_python/context.py

@@ -51,7 +51,10 @@ from aws_durable_execution_sdk_python.operation.callback import (
     CallbackOperationExecutor,
     wait_for_callback_handler,
 )
-from aws_durable_execution_sdk_python.operation.child import child_handler
+from aws_durable_execution_sdk_python.operation.child import (
+    ChildOperationExecutor,
+    child_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
@@ -515,12 +518,18 @@ class DurableContext(DurableContextProtocol):
 
     @contextmanager
     def _replay_aware(
-        self, operation_identifier: OperationIdentifier | None = None
+        self,
+        operation_identifier: OperationIdentifier | None = None,
+        *,
+        is_virtual: bool = False,
     ) -> Iterator[None]:
         """Wrap a single operation with replay-boundary detection.
 
         The operation kind is inferred from its own checkpoint (when one
-        exists), so no per-call-site flag is needed.
+        exists). The one per-call-site input is `is_virtual`: a virtual child
+        context never checkpoints itself, so with `is_virtual=True` any
+        checkpoint at its id is rejected right after the identity check,
+        before replay and update hooks run.
 
         The boundary has these parts:
 
@@ -555,6 +564,10 @@ class DurableContext(DurableContextProtocol):
         if was_replaying and next_operation is not None:
             if operation_identifier is not None:
                 operation_identifier.validate_checkpoint(next_operation)
+                if is_virtual:
+                    ChildOperationExecutor.validate_virtual_checkpoint(
+                        operation_identifier, next_operation
+                    )
 
             # While replaying, an operation that already has a checkpoint was
             # observed in a prior invocation. If the backend says the operation
@@ -599,7 +612,9 @@ class DurableContext(DurableContextProtocol):
         self,
         sub_type: OperationSubType,
         name: str | None = None,
-        operation_type: OperationType | None = None,
+        *,
+        operation_type: OperationType,
+        is_virtual: bool = False,
     ) -> Iterator[OperationIdentifier]:
         """Allocate an operation and validate its replay identity before hooks."""
         operation_identifier = OperationIdentifier(
@@ -609,7 +624,7 @@ class DurableContext(DurableContextProtocol):
             name=name,
             operation_type=operation_type,
         )
-        with self._replay_aware(operation_identifier):
+        with self._replay_aware(operation_identifier, is_virtual=is_virtual):
             yield operation_identifier
 
     # endregion replay status
@@ -636,7 +651,7 @@ class DurableContext(DurableContextProtocol):
         if not config:
             config = CallbackConfig()
         with self._operation_replay_aware(
-            OperationSubType.CALLBACK, name
+            OperationSubType.CALLBACK, name, operation_type=OperationType.CALLBACK
         ) as operation_identifier:
             executor: CallbackOperationExecutor = CallbackOperationExecutor(
                 state=self.state,
@@ -672,7 +687,9 @@ class DurableContext(DurableContextProtocol):
         if not config:
             config = InvokeConfig[P, R]()
         with self._operation_replay_aware(
-            OperationSubType.CHAINED_INVOKE, name
+            OperationSubType.CHAINED_INVOKE,
+            name,
+            operation_type=OperationType.CHAINED_INVOKE,
         ) as operation_identifier:
             executor: InvokeOperationExecutor[R] = InvokeOperationExecutor(
                 function_name=function_name,
@@ -700,7 +717,7 @@ class DurableContext(DurableContextProtocol):
             config.completion_config._validate_for_total(len(inputs))
 
         with self._operation_replay_aware(
-            OperationSubType.MAP, map_name
+            OperationSubType.MAP, map_name, operation_type=OperationType.CONTEXT
         ) as operation_identifier:
             operation_id = operation_identifier.operation_id
             map_context = self.create_child_context(operation_id=operation_id)
@@ -751,7 +768,7 @@ class DurableContext(DurableContextProtocol):
             config.completion_config._validate_for_total(len(functions))
 
         with self._operation_replay_aware(
-            OperationSubType.PARALLEL, name
+            OperationSubType.PARALLEL, name, operation_type=OperationType.CONTEXT
         ) as operation_identifier:
             operation_id = operation_identifier.operation_id
             parallel_context = self.create_child_context(operation_id=operation_id)
@@ -811,14 +828,15 @@ class DurableContext(DurableContextProtocol):
             if config and config.sub_type
             else OperationSubType.RUN_IN_CHILD_CONTEXT
         )
+        is_virtual: bool = config.is_virtual if config else False
 
         with self._operation_replay_aware(
             sub_type,
             step_name,
             operation_type=OperationType.CONTEXT,
+            is_virtual=is_virtual,
         ) as operation_identifier:
             operation_id = operation_identifier.operation_id
-            is_virtual: bool = config.is_virtual if config else False
 
             def callable_with_child_context():
                 return func(
@@ -845,7 +863,7 @@ class DurableContext(DurableContextProtocol):
         if not config:
             config = StepConfig()
         with self._operation_replay_aware(
-            OperationSubType.STEP, step_name
+            OperationSubType.STEP, step_name, operation_type=OperationType.STEP
         ) as operation_identifier:
             executor: StepOperationExecutor[T] = StepOperationExecutor(
                 func=func,
@@ -868,7 +886,7 @@ class DurableContext(DurableContextProtocol):
             msg = "duration must be at least 1 second"
             raise ValidationError(msg)
         with self._operation_replay_aware(
-            OperationSubType.WAIT, name
+            OperationSubType.WAIT, name, operation_type=OperationType.WAIT
         ) as operation_identifier:
             wait_seconds = duration.seconds
             executor: WaitOperationExecutor = WaitOperationExecutor(
@@ -938,7 +956,7 @@ class DurableContext(DurableContextProtocol):
             raise ValidationError(msg)
 
         with self._operation_replay_aware(
-            OperationSubType.WAIT_FOR_CONDITION, name
+            OperationSubType.WAIT_FOR_CONDITION, name, operation_type=OperationType.STEP
         ) as operation_identifier:
             executor: WaitForConditionOperationExecutor[T] = (
                 WaitForConditionOperationExecutor(

aws_durable_execution_sdk_python/exceptions.py

@@ -13,7 +13,7 @@ BAD_REQUEST_ERROR: int = 400
 TOO_MANY_REQUESTS_ERROR: int = 429
 SERVICE_ERROR: int = 500
 INVALID_PARAMETER_VALUE_EXCEPTION: str = "InvalidParameterValueException"
-INVALID_CHECKPOINT_TOKEN_PREFIX: str = "Invalid Checkpoint Token"
+INVALID_CHECKPOINT_TOKEN_PREFIX: str = "Invalid checkpoint token"
 
 # Non-retryable customer error codes that arrive as non-4xx (e.g. HTTP 502) from Lambda.
 # Unlike typical 5xx errors, these require customer intervention (e.g., fixing
@@ -162,7 +162,7 @@ class BotoClientError(InvocationError):
           These arrive as HTTP 502 but require customer intervention to fix.
         - 4xx errors → EXECUTION, except:
           - 429 (TooManyRequests) → INVOCATION (throttling is transient)
-          - InvalidParameterValueException with "Invalid Checkpoint Token" → INVOCATION
+          - InvalidParameterValueException with "Invalid checkpoint token" (exact match) → INVOCATION
             (stale token from a concurrent checkpoint; next invocation gets a fresh token)
         - 5xx, network errors → INVOCATION
         """