aws-durable-execution-sdk-python 1.4.0__py3-none-any.whl → 1.5.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.4.0"
+__version__ = "1.5.0"

aws_durable_execution_sdk_python/__init__.py

@@ -7,8 +7,10 @@ from aws_durable_execution_sdk_python.__about__ import __version__
 # Helper decorators - commonly used for step functions
 # Concurrency
 from aws_durable_execution_sdk_python.concurrency.models import BatchResult
+from aws_durable_execution_sdk_python.config import ParallelBranch
 from aws_durable_execution_sdk_python.context import (
     DurableContext,
+    durable_parallel_branch,
     durable_step,
     durable_wait_for_callback,
     durable_with_child_context,
@@ -27,15 +29,18 @@ from aws_durable_execution_sdk_python.execution import durable_execution
 # Essential context types - passed to user functions
 from aws_durable_execution_sdk_python.types import StepContext
 
+
 __all__ = [
     "BatchResult",
     "DurableContext",
     "DurableExecutionsError",
     "InvocationError",
+    "ParallelBranch",
     "StepContext",
     "ValidationError",
     "__version__",
     "durable_execution",
+    "durable_parallel_branch",
     "durable_step",
     "durable_wait_for_callback",
     "durable_with_child_context",

aws_durable_execution_sdk_python/concurrency/executor.py

@@ -20,7 +20,10 @@ from aws_durable_execution_sdk_python.concurrency.models import (
     ExecutionCounters,
     SuspendResult,
 )
-from aws_durable_execution_sdk_python.config import ChildConfig
+from aws_durable_execution_sdk_python.config import (
+    ChildConfig,
+    NestingType,
+)
 from aws_durable_execution_sdk_python.exceptions import (
     OrphanedChildException,
     SuspendExecution,
@@ -143,6 +146,7 @@ class ConcurrentExecutor(ABC, Generic[CallableType, ResultType]):
         serdes: SerDes | None,
         item_serdes: SerDes | None = None,
         summary_generator: SummaryGenerator | None = None,
+        nesting_type: NestingType = NestingType.NESTED,
     ):
         """Initialize ConcurrentExecutor.
 
@@ -160,6 +164,7 @@ class ConcurrentExecutor(ABC, Generic[CallableType, ResultType]):
         self.sub_type_iteration = sub_type_iteration
         self.name_prefix = name_prefix
         self.summary_generator = summary_generator
+        self.nesting_type = nesting_type
 
         # Event-driven state tracking for when the executor is done
         self._completion_event = threading.Event()
@@ -189,6 +194,14 @@ class ConcurrentExecutor(ABC, Generic[CallableType, ResultType]):
         """Execute a single executable in a child context and return the result."""
         raise NotImplementedError
 
+    def get_iteration_name(self, index: int) -> str:
+        """Get the display name for an iteration/branch at the given index.
+
+        Subclasses can override this to provide custom naming (e.g., from item_namer
+        or branch names). The default returns "{name_prefix}{index}".
+        """
+        return f"{self.name_prefix}{index}"
+
     def execute(
         self, execution_state: ExecutionState, executor_context: DurableContext
     ) -> BatchResult[ResultType]:
@@ -391,29 +404,36 @@ class ConcurrentExecutor(ABC, Generic[CallableType, ResultType]):
         """
         Execute a single item in a derived child context.
 
-        instead of relying on `executor_context.run_in_child_context`
-        we generate an operation_id for the child, and then call `child_handler`
-        directly. This avoids the hidden mutation of the context's internal counter.
-        we can do this because we explicitly control the generation of step_id and do it
-        using executable.index.
+        Instead of relying on `executor_context.run_in_child_context` we
+        generate an operation_id for the child, then call `child_handler`
+        directly. This avoids the hidden mutation of the context's
+        internal counter. We explicitly derive the child's operation_id
+        from `executable.index` so that the same input always produces
+        the same id regardless of the order branches actually run in.
 
-
-        invariant: `operation_id` for a given executable is deterministic,
-            and execution order invariant.
+        Invariant: `operation_id` for a given executable is deterministic
+        and execution-order invariant.
         """
 
-        operation_id = executor_context._create_step_id_for_logical_step(  # noqa: SLF001
+        operation_id: str = executor_context._create_step_id_for_logical_step(  # noqa: SLF001
             executable.index
         )
-        name = f"{self.name_prefix}{executable.index}"
-        child_context = executor_context.create_child_context(operation_id)
+        name: str = self.get_iteration_name(executable.index)
+        is_virtual: bool = self.nesting_type is NestingType.FLAT
+
+        child_context: DurableContext = executor_context.create_child_context(
+            operation_id, is_virtual=is_virtual
+        )
+        # For NESTED this is for branch's START/SUCCEED/FAIL checkpoints (not the children of the branch).
+        # For FLAT `child_handler` skips checkpoints, so not used.
+        # Construct it unconditionally to keep the call simple.
         operation_identifier = OperationIdentifier(
             operation_id,
             executor_context._parent_id,  # noqa: SLF001
             name,
         )
 
-        def run_in_child_handler():
+        def run_in_child_handler() -> ResultType:
             return self.execute_item(child_context, executable)
 
         result: ResultType = child_handler(
@@ -424,6 +444,7 @@ class ConcurrentExecutor(ABC, Generic[CallableType, ResultType]):
                 serdes=self.item_serdes or self.serdes,
                 sub_type=self.sub_type_iteration,
                 summary_generator=self.summary_generator,
+                is_virtual=is_virtual,
             ),
         )
         child_context.state.track_replay(operation_id=operation_id)

aws_durable_execution_sdk_python/config.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Generic, TypeVar
 
 from aws_durable_execution_sdk_python.exceptions import ValidationError
 
+
 P = TypeVar("P")  # Payload type
 R = TypeVar("R")  # Result type
 T = TypeVar("T")
@@ -76,6 +77,42 @@ class TerminationMode(Enum):
     ABANDON = "ABANDON"
 
 
+class NestingType(Enum):
+    """Control how child contexts are created for batch operations.
+
+    Applies to `map` and `parallel`. Each branch or iteration runs inside a
+    child context.
+
+        - NESTED: full checkpointed context
+        - FLAT: a virtual context that skips checkpoints for the branch/iteration.
+
+    """
+
+    NESTED = "NESTED"
+    """Create CONTEXT operations for each branch/iteration with full checkpointing.
+
+    Operations within each branch/iteration are wrapped in their own context.
+
+    - Observability: high — each branch/iteration appears as a separate
+      operation in execution history.
+    - Cost: higher — consumes more operations due to CONTEXT creation
+      overhead.
+    - Scale: lower maximum iterations due to operation limits.
+    """
+
+    FLAT = "FLAT"
+    """Skip CONTEXT operations for branches/iterations using virtual contexts.
+
+    Operations execute directly without individual context wrapping.
+
+    - Observability: lower — branches/iterations don't appear as separate
+      operations in execution history.
+    - Cost: ~30% lower — reduces operation consumption by skipping CONTEXT
+      overhead.
+    - Scale: higher maximum iterations possible within operation limits.
+    """
+
+
 @dataclass(frozen=True)
 class CompletionConfig:
     """Configuration for determining when parallel/map operations complete.
@@ -187,6 +224,10 @@ class ParallelConfig:
             Used internally by map/parallel operations to handle large BatchResult payloads.
             Signature: (result: T) -> str
 
+        nesting_type: How child operations should inherit context from their parent.
+            - NESTED: Each branch runs in its own isolated context (default)
+            - FLAT: All branches share the same parent context
+
     Example:
         # Run at most 3 branches concurrently, succeed if any one succeeds
         config = ParallelConfig(
@@ -202,6 +243,42 @@ class ParallelConfig:
     serdes: SerDes | None = None
     item_serdes: SerDes | None = None
     summary_generator: SummaryGenerator | None = None
+    nesting_type: NestingType = NestingType.NESTED
+
+
+@dataclass(frozen=True)
+class ParallelBranch(Generic[T]):
+    """A named branch for parallel execution.
+
+    Use this to provide custom names for parallel branches, improving
+    observability in execution history.
+
+    Type Parameters:
+        T: The return type of the branch function.
+
+    Args:
+        func: The callable to execute in this branch. Receives a DurableContext.
+        name: Optional custom name for this branch. When provided, replaces
+            the default "parallel-branch-{index}" naming in execution history.
+            This affects observability but not replay determinism.
+
+    Example:
+        context.parallel(
+            functions=[
+                ParallelBranch(func=lambda ctx: fetch_user(ctx), name="fetch-user-data"),
+                ParallelBranch(func=lambda ctx: fetch_orders(ctx), name="fetch-order-history"),
+            ],
+            name="load-data",
+            config=ParallelConfig(max_concurrency=2),
+        )
+    """
+
+    func: Callable
+    name: str | None = None
+
+    def __call__(self, *args, **kwargs):
+        """Delegate to the wrapped function, making ParallelBranch itself callable."""
+        return self.func(*args, **kwargs)
 
 
 class StepSemantics(Enum):
@@ -218,12 +295,6 @@ class StepConfig:
     serdes: SerDes | None = None
 
 
-class CheckpointMode(Enum):
-    NO_CHECKPOINT = ("NO_CHECKPOINT",)
-    CHECKPOINT_AT_FINISH = ("CHECKPOINT_AT_FINISH",)
-    CHECKPOINT_AT_START_AND_FINISH = "CHECKPOINT_AT_START_AND_FINISH"
-
-
 @dataclass(frozen=True)
 class ChildConfig(Generic[T]):
     """Configuration options for child context operations.
@@ -259,21 +330,23 @@ class ChildConfig(Generic[T]):
 
             Used internally by map/parallel operations to handle large BatchResult payloads.
             Signature: (result: T) -> str
-    Note:
-        checkpoint_mode field is commented out as it's not currently implemented.
-        When implemented, it will control when checkpoints are created:
-        - CHECKPOINT_AT_START_AND_FINISH: Checkpoint at both start and completion (default)
-        - CHECKPOINT_AT_FINISH: Only checkpoint when operation completes
-        - NO_CHECKPOINT: No automatic checkpointing
+
+        is_virtual: When True, skip all checkpoints (START, SUCCEED,
+            FAIL) for this child context and propagate the caller's reporting
+            parent id through to operations created inside the child. The
+            branch is a logical scope for step-id prefixing but does not
+            appear in the execution history. Used internally by
+            NestingType.FLAT branches. Use this to group operations without
+            adding a CONTEXT entry to the execution history.
 
     See TypeScript reference: aws-durable-execution-sdk-js/src/types/index.ts
     """
 
-    # checkpoint_mode: CheckpointMode = CheckpointMode.CHECKPOINT_AT_START_AND_FINISH
     serdes: SerDes | None = None
     item_serdes: SerDes | None = None
     sub_type: OperationSubType | None = None
     summary_generator: SummaryGenerator | None = None
+    is_virtual: bool = False
 
 
 class ItemsPerBatchUnit(Enum):
@@ -317,12 +390,15 @@ class ItemBatcher(Generic[T]):
 
 
 @dataclass(frozen=True)
-class MapConfig:
+class MapConfig(Generic[T]):
     """Configuration options for map operations over collections.
 
     This class configures how map operations process collections of items,
     including concurrency, batching, completion criteria, and serialization.
 
+    Type Parameters:
+        T: The type of items being processed in the map operation.
+
     Args:
         max_concurrency: Maximum number of items to process concurrently.
             If None, no limit is imposed and all items are processed concurrently.
@@ -361,6 +437,16 @@ class MapConfig:
             Used internally by map/parallel operations to handle large BatchResult payloads.
             Signature: (result: T) -> str
 
+        nesting_type: How child operations should inherit context from their parent.
+            - NESTED: Each item runs in its own isolated context (default)
+            - FLAT: All items share the same parent context
+
+        item_namer: Optional callable to generate custom names for each map iteration.
+            When provided, replaces the default "map-item-{index}" naming scheme.
+            Receives the item and its index, and returns a string name for that iteration.
+            This affects observability (execution history names) but not replay determinism.
+            If None, uses the default naming: "map-item-{index}".
+
     Example:
         # Process 5 items at a time, batch by count, require all to succeed
         config = MapConfig(
@@ -368,6 +454,12 @@ class MapConfig:
             item_batcher=ItemBatcher(max_items_per_batch=10),
             completion_config=CompletionConfig.all_successful()
         )
+
+        # With custom iteration names
+        config = MapConfig(
+            max_concurrency=5,
+            item_namer=lambda item, index: f"process-order-{item.id}"
+        )
     """
 
     max_concurrency: int | None = None
@@ -376,6 +468,8 @@ class MapConfig:
     serdes: SerDes | None = None
     item_serdes: SerDes | None = None
     summary_generator: SummaryGenerator | None = None
+    nesting_type: NestingType = NestingType.NESTED
+    item_namer: Callable[[T, int], str] | None = None
 
 
 @dataclass(frozen=True)

aws_durable_execution_sdk_python/context.py

@@ -12,6 +12,7 @@ from aws_durable_execution_sdk_python.config import (
     Duration,
     InvokeConfig,
     MapConfig,
+    ParallelBranch,
     ParallelConfig,
     StepConfig,
     WaitForCallbackConfig,
@@ -55,6 +56,7 @@ from aws_durable_execution_sdk_python.types import (
     WaitForConditionCheckContext,
 )
 
+
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
 
@@ -119,6 +121,52 @@ def durable_with_child_context(
     return wrapper
 
 
+def durable_parallel_branch(
+    name: str | None = None,
+) -> Callable[
+    [Callable[Concatenate[DurableContext, Params], T]],
+    Callable[Params, ParallelBranch[T]],
+]:
+    """Wrap your callable into a named ParallelBranch for use with context.parallel().
+
+    This is a decorator factory — call it with an optional name to produce
+    the actual decorator.
+
+    Args:
+        name: Optional custom name for this branch. When provided, replaces
+            the default "parallel-branch-{index}" naming in execution history.
+            If None, the function's __name__ is used.
+
+    Example:
+        @durable_parallel_branch(name="fetch-user-data")
+        def fetch_user(ctx: DurableContext, user_id: str) -> dict:
+            return ctx.step(lambda _: {"id": user_id, "name": "Jane"}, name="load_user")
+
+        @durable_parallel_branch(name="fetch-orders")
+        def fetch_orders(ctx: DurableContext, user_id: str) -> list:
+            return ctx.step(lambda _: ["order1", "order2"], name="load_orders")
+
+        # Usage in a durable handler:
+        results = context.parallel(
+            functions=[fetch_user(user_id), fetch_orders(user_id)],
+            name="load-data",
+        )
+    """
+
+    def decorator(
+        func: Callable[Concatenate[DurableContext, Params], T],
+    ) -> Callable[Params, ParallelBranch[T]]:
+        def wrapper(*args, **kwargs) -> ParallelBranch[T]:
+            def function_with_arguments(ctx: DurableContext) -> T:
+                return func(ctx, *args, **kwargs)
+
+            return ParallelBranch(func=function_with_arguments, name=name)
+
+        return wrapper
+
+    return decorator
+
+
 def durable_wait_for_callback(
     func: Callable[Concatenate[str, WaitForCallbackContext, Params], T],
 ) -> Callable[Params, Callable[[str, WaitForCallbackContext], T]]:
@@ -237,11 +285,20 @@ class DurableContext(DurableContextProtocol):
         lambda_context: LambdaContext | None = None,
         parent_id: str | None = None,
         logger: Logger | None = None,
+        step_id_prefix: str | None = None,
     ) -> None:
         self.state: ExecutionState = state
         self.execution_context: ExecutionContext = execution_context
         self.lambda_context = lambda_context
+        # operations inside this context use this id as their parent
         self._parent_id: str | None = parent_id
+        # child operations use this to generate deterministic step ids.
+        # differs from `parent_id` only for virtual contexts.
+        self._step_id_prefix: str | None = (
+            step_id_prefix if step_id_prefix is not None else parent_id
+        )
+        # cached at construction to make invariant even if parent/prefix mutates.
+        self._is_virtual: bool = self._parent_id != self._step_id_prefix
         self._step_counter: OrderedCounter = OrderedCounter()
 
         log_info = LogInfo(
@@ -254,6 +311,18 @@ class DurableContext(DurableContextProtocol):
             info=log_info,
         )
 
+    @property
+    def is_virtual(self) -> bool:
+        """True if this context does not checkpoint its own start and completion.
+
+        You create a virtual context by `create_child_context(..., is_virtual=True)`.
+        FLAT-mode `map`/`parallel` branches uses virtual contexts. Inner operations
+        use the grandfather as parent (enclosing non-virtual ancestor, skipping the branch level
+        in the hierarchy), while step ids are still prefixed with the branch's own
+        operation id so replay stays deterministic.
+        """
+        return self._is_virtual
+
     # region factories
     @staticmethod
     def from_lambda_context(
@@ -269,18 +338,43 @@ class DurableContext(DurableContextProtocol):
             parent_id=None,
         )
 
-    def create_child_context(self, parent_id: str) -> DurableContext:
-        """Create a child context from the given parent."""
-        logger.debug("Creating child context for parent %s", parent_id)
+    def create_child_context(
+        self, operation_id: str, *, is_virtual: bool = False
+    ) -> DurableContext:
+        """Create a child context for the given operation.
+
+        Args:
+            operation_id: The operation id that owns the child context. Used as
+                the child's step-id prefix in all cases.
+            is_virtual: When `True`, create a virtual child whose inner
+                operations report to this context's own `_parent_id` (one
+                level up the hierarchy). When `False` (default), produce a
+                regular child whose inner operations report to
+                `operation_id`.
+
+        Returns:
+            A new `DurableContext` child for the current context.
+        """
+        # For a virtual child, propagate the current `_parent_id` so its
+        # inner operations refer to the grandparent rather than the parent.
+        # For a regular non-virtual child, the child's own `operation_id` is
+        # the parent id for its inner operations (standard nesting).
+        child_parent_id: str | None = self._parent_id if is_virtual else operation_id
+        logger.debug(
+            "Creating child context for operation %s (is_virtual=%s)",
+            operation_id,
+            is_virtual,
+        )
         return DurableContext(
             state=self.state,
             execution_context=self.execution_context,
             lambda_context=self.lambda_context,
-            parent_id=parent_id,
+            parent_id=child_parent_id,
+            step_id_prefix=operation_id,
             logger=self.logger.with_log_info(
                 LogInfo(
                     execution_state=self.state,
-                    parent_id=parent_id,
+                    parent_id=child_parent_id,
                 )
             ),
         )
@@ -310,7 +404,8 @@ class DurableContext(DurableContextProtocol):
         This allows us to recover operation ids or even look
         forward without changing the internal state of this context.
         """
-        step_id = f"{self._parent_id}-{step}" if self._parent_id else str(step)
+        prefix: str | None = self._step_id_prefix
+        step_id: str = f"{prefix}-{step}" if prefix else str(step)
         return hashlib.blake2b(step_id.encode()).hexdigest()[:64]
 
     def _create_step_id(self) -> str:
@@ -347,7 +442,9 @@ class DurableContext(DurableContextProtocol):
         executor: CallbackOperationExecutor = CallbackOperationExecutor(
             state=self.state,
             operation_identifier=OperationIdentifier(
-                operation_id=operation_id, parent_id=self._parent_id, name=name
+                operation_id=operation_id,
+                parent_id=self._parent_id,
+                name=name,
             ),
             config=config,
         )
@@ -409,9 +506,11 @@ class DurableContext(DurableContextProtocol):
 
         operation_id = self._create_step_id()
         operation_identifier = OperationIdentifier(
-            operation_id=operation_id, parent_id=self._parent_id, name=map_name
+            operation_id=operation_id,
+            parent_id=self._parent_id,
+            name=map_name,
         )
-        map_context = self.create_child_context(parent_id=operation_id)
+        map_context = self.create_child_context(operation_id=operation_id)
 
         def map_in_child_context() -> BatchResult[R]:
             # map_context is a child_context of the context upon which `.map`
@@ -445,14 +544,14 @@ class DurableContext(DurableContextProtocol):
 
     def parallel(
         self,
-        functions: Sequence[Callable[[DurableContext], T]],
+        functions: Sequence[Callable[[DurableContext], T] | ParallelBranch[T]],
         name: str | None = None,
         config: ParallelConfig | None = None,
     ) -> BatchResult[T]:
         """Execute multiple callables in parallel."""
         # _create_step_id() is thread-safe. rest of method is safe, since using local copy of parent id
         operation_id = self._create_step_id()
-        parallel_context = self.create_child_context(parent_id=operation_id)
+        parallel_context = self.create_child_context(operation_id=operation_id)
         operation_identifier = OperationIdentifier(
             operation_id=operation_id, parent_id=self._parent_id, name=name
         )
@@ -508,14 +607,22 @@ class DurableContext(DurableContextProtocol):
         # _create_step_id() is thread-safe. rest of method is safe, since using local copy of parent id
         operation_id = self._create_step_id()
 
+        is_virtual: bool = config.is_virtual if config else False
+
         def callable_with_child_context():
-            return func(self.create_child_context(parent_id=operation_id))
+            return func(
+                self.create_child_context(
+                    operation_id=operation_id, is_virtual=is_virtual
+                )
+            )
 
         result: T = child_handler(
             func=callable_with_child_context,
             state=self.state,
             operation_identifier=OperationIdentifier(
-                operation_id=operation_id, parent_id=self._parent_id, name=step_name
+                operation_id=operation_id,
+                parent_id=self._parent_id,
+                name=step_name,
             ),
             config=config,
         )