[feature]: add nesting_type to concurrency operations (#351)

Author: zhongkechen

.github/workflows/integration-tests.yml

@@ -50,7 +50,7 @@ jobs:
       run: |
         echo "Running Testing SDK tests against Language SDK PR changes..."
         echo "Using Language SDK from: $AWS_DURABLE_SDK_URL"
-        python -m pip install -e .
+        hatch run -- test:pip install -e ../language-sdk
         hatch fmt --check
         hatch run types:check
         hatch run test:cov

src/aws_durable_execution_sdk_python/concurrency/executor.py

@@ -20,7 +20,10 @@
     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,
@@ -134,6 +137,7 @@ class ConcurrentExecutor(ABC, Generic[CallableType, ResultType]):
 
     def __init__(
         self,
+        operation_identifier: OperationIdentifier,
         executables: list[Executable[CallableType]],
         max_concurrency: int | None,
         completion_config: CompletionConfig,
@@ -143,6 +147,7 @@ def __init__(
         serdes: SerDes | None,
         item_serdes: SerDes | None = None,
         summary_generator: SummaryGenerator | None = None,
+        nesting_type: NestingType = NestingType.NESTED,
     ):
         """Initialize ConcurrentExecutor.
 
@@ -153,13 +158,15 @@ def __init__(
                 handle large BatchResult payloads efficiently. Matches TypeScript behavior in
                 run-in-child-context-handler.ts.
         """
+        self.operation_identifier = operation_identifier
         self.executables = executables
         self.max_concurrency = max_concurrency
         self.completion_config = completion_config
         self.sub_type_top = sub_type_top
         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()
@@ -406,7 +413,14 @@ def _execute_item_in_child_context(
             executable.index
         )
         name = f"{self.name_prefix}{executable.index}"
-        child_context = executor_context.create_child_context(operation_id)
+        non_virtual_parent_id = (
+            self.operation_identifier.operation_id
+            if self.nesting_type is NestingType.FLAT
+            else None
+        )
+        child_context = executor_context.create_child_context(
+            operation_id, non_virtual_parent_id
+        )
         operation_identifier = OperationIdentifier(
             operation_id,
             executor_context._parent_id,  # noqa: SLF001
@@ -424,6 +438,7 @@ def run_in_child_handler():
                 serdes=self.item_serdes or self.serdes,
                 sub_type=self.sub_type_iteration,
                 summary_generator=self.summary_generator,
+                is_virtual=self.nesting_type is NestingType.FLAT,
             ),
         )
         child_context.state.track_replay(operation_id=operation_id)

src/aws_durable_execution_sdk_python/config.py

@@ -76,6 +76,18 @@ class TerminationMode(Enum):
     ABANDON = "ABANDON"
 
 
+class NestingType(Enum):
+    """
+    How child operations should inherit context from their parent.
+
+    - NESTED: Each child operation runs in its own isolated context
+    - FLAT: All child operations share the same parent context
+    """
+
+    NESTED = "NESTED"
+    FLAT = "FLAT"
+
+
 @dataclass(frozen=True)
 class CompletionConfig:
     """Configuration for determining when parallel/map operations complete.
@@ -187,6 +199,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 +218,7 @@ class ParallelConfig:
     serdes: SerDes | None = None
     item_serdes: SerDes | None = None
     summary_generator: SummaryGenerator | None = None
+    nesting_type: NestingType = NestingType.NESTED
 
 
 class StepSemantics(Enum):
@@ -218,12 +235,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 +270,19 @@ 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: Whether the child operation is virtual (doesn't represent a real operation).
+            Virtual contexts are used for concurrency operations and don't appear in
+            the final execution history. Default is False.
 
     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):
@@ -361,6 +370,10 @@ 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
+
     Example:
         # Process 5 items at a time, batch by count, require all to succeed
         config = MapConfig(
@@ -376,6 +389,7 @@ class MapConfig:
     serdes: SerDes | None = None
     item_serdes: SerDes | None = None
     summary_generator: SummaryGenerator | None = None
+    nesting_type: NestingType = NestingType.NESTED
 
 
 @dataclass(frozen=True)

src/aws_durable_execution_sdk_python/context.py

@@ -237,12 +237,14 @@ def __init__(
         lambda_context: LambdaContext | None = None,
         parent_id: str | None = None,
         logger: Logger | None = None,
+        non_virtual_parent_id: str | None = None,
     ) -> None:
         self.state: ExecutionState = state
         self.execution_context: ExecutionContext = execution_context
         self.lambda_context = lambda_context
         self._parent_id: str | None = parent_id
         self._step_counter: OrderedCounter = OrderedCounter()
+        self._non_virtual_parent_id = non_virtual_parent_id or parent_id
 
         log_info = LogInfo(
             execution_state=state,
@@ -269,7 +271,9 @@ def from_lambda_context(
             parent_id=None,
         )
 
-    def create_child_context(self, parent_id: str) -> DurableContext:
+    def create_child_context(
+        self, parent_id: str, non_virtual_parent_id=None
+    ) -> DurableContext:
         """Create a child context from the given parent."""
         logger.debug("Creating child context for parent %s", parent_id)
         return DurableContext(
@@ -283,6 +287,7 @@ def create_child_context(self, parent_id: str) -> DurableContext:
                     parent_id=parent_id,
                 )
             ),
+            non_virtual_parent_id=non_virtual_parent_id,
         )
 
     # endregion factories
@@ -347,7 +352,9 @@ def create_callback(
         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._non_virtual_parent_id,
+                name=name,
             ),
             config=config,
         )
@@ -388,7 +395,7 @@ def invoke(
             state=self.state,
             operation_identifier=OperationIdentifier(
                 operation_id=operation_id,
-                parent_id=self._parent_id,
+                parent_id=self._non_virtual_parent_id,
                 name=name,
             ),
             config=config,
@@ -409,7 +416,9 @@ def map(
 
         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._non_virtual_parent_id,
+            name=map_name,
         )
         map_context = self.create_child_context(parent_id=operation_id)
 
@@ -454,7 +463,7 @@ def parallel(
         operation_id = self._create_step_id()
         parallel_context = self.create_child_context(parent_id=operation_id)
         operation_identifier = OperationIdentifier(
-            operation_id=operation_id, parent_id=self._parent_id, name=name
+            operation_id=operation_id, parent_id=self._non_virtual_parent_id, name=name
         )
 
         def parallel_in_child_context() -> BatchResult[T]:
@@ -515,7 +524,9 @@ def callable_with_child_context():
             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._non_virtual_parent_id,
+                name=step_name,
             ),
             config=config,
         )
@@ -539,7 +550,7 @@ def step(
             state=self.state,
             operation_identifier=OperationIdentifier(
                 operation_id=operation_id,
-                parent_id=self._parent_id,
+                parent_id=self._non_virtual_parent_id,
                 name=step_name,
             ),
             context_logger=self.logger,
@@ -566,7 +577,7 @@ def wait(self, duration: Duration, name: str | None = None) -> None:
             state=self.state,
             operation_identifier=OperationIdentifier(
                 operation_id=operation_id,
-                parent_id=self._parent_id,
+                parent_id=self._non_virtual_parent_id,
                 name=name,
             ),
         )
@@ -621,7 +632,7 @@ def wait_for_condition(
                 state=self.state,
                 operation_identifier=OperationIdentifier(
                     operation_id=operation_id,
-                    parent_id=self._parent_id,
+                    parent_id=self._non_virtual_parent_id,
                     name=name,
                 ),
                 context_logger=self.logger,

src/aws_durable_execution_sdk_python/operation/child.py

@@ -118,7 +118,7 @@ def check_result_status(self) -> CheckResult[T]:
             checkpointed_result.raise_callable_error()
 
         # Create START checkpoint if not exists
-        if not checkpointed_result.is_existent():
+        if not checkpointed_result.is_existent() and not self.config.is_virtual:
             start_operation: OperationUpdate = OperationUpdate.create_context_start(
                 identifier=self.operation_identifier,
                 sub_type=self.sub_type,
@@ -157,6 +157,14 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T:
         try:
             raw_result: T = self.func()
 
+            if self.config.is_virtual:
+                logger.debug(
+                    "Virtual context: Exiting child context without creating another checkpoint. id: %s, name: %s",
+                    self.operation_identifier.operation_id,
+                    self.operation_identifier.name,
+                )
+                return raw_result
+
             # If in replay_children mode, return without checkpointing
             if checkpointed_result.is_replay_children():
                 logger.debug(
@@ -270,8 +278,7 @@ def child_handler(
     Raises:
         May raise operation-specific errors during execution
     """
-    if not config:
-        config = ChildConfig()
-
-    executor = ChildOperationExecutor(func, state, operation_identifier, config)
+    executor = ChildOperationExecutor(
+        func, state, operation_identifier, config or ChildConfig()
+    )
     return executor.process()

src/aws_durable_execution_sdk_python/operation/map.py

@@ -12,7 +12,7 @@
     BatchResult,
     Executable,
 )
-from aws_durable_execution_sdk_python.config import MapConfig
+from aws_durable_execution_sdk_python.config import MapConfig, NestingType
 from aws_durable_execution_sdk_python.lambda_service import OperationSubType
 
 if TYPE_CHECKING:
@@ -36,6 +36,7 @@
 class MapExecutor(Generic[T, R], ConcurrentExecutor[Callable, R]):  # noqa: PYI059
     def __init__(
         self,
+        operation_identifier: OperationIdentifier,
         executables: list[Executable[Callable]],
         items: Sequence[T],
         max_concurrency: int | None,
@@ -46,8 +47,10 @@ def __init__(
         serdes: SerDes | None,
         summary_generator: SummaryGenerator | None = None,
         item_serdes: SerDes | None = None,
+        nesting_type: NestingType = NestingType.NESTED,
     ):
         super().__init__(
+            operation_identifier=operation_identifier,
             executables=executables,
             max_concurrency=max_concurrency,
             completion_config=completion_config,
@@ -57,12 +60,14 @@ def __init__(
             serdes=serdes,
             summary_generator=summary_generator,
             item_serdes=item_serdes,
+            nesting_type=nesting_type,
         )
         self.items = items
 
     @classmethod
     def from_items(
         cls,
+        operation_identifier: OperationIdentifier,
         items: Sequence[T],
         func: Callable,
         config: MapConfig,
@@ -73,6 +78,7 @@ def from_items(
         ]
 
         return cls(
+            operation_identifier=operation_identifier,
             executables=executables,
             items=items,
             max_concurrency=config.max_concurrency,
@@ -83,6 +89,7 @@ def from_items(
             serdes=config.serdes,
             summary_generator=config.summary_generator,
             item_serdes=config.item_serdes,
+            nesting_type=config.nesting_type,
         )
 
     def execute_item(self, child_context, executable: Executable[Callable]) -> R:
@@ -109,6 +116,7 @@ def map_handler(
     # See TypeScript reference: aws-durable-execution-sdk-js/src/handlers/map-handler/map-handler.ts (~line 79)
 
     executor: MapExecutor[T, R] = MapExecutor.from_items(
+        operation_identifier=operation_identifier,
         items=items,
         func=func,
         config=config or MapConfig(summary_generator=MapSummaryGenerator()),