[feature]: add nesting_type to concurrency operations

Author: zhongkechen

src/aws_durable_execution_sdk_python/concurrency/executor.py

@@ -20,7 +20,7 @@
     ExecutionCounters,
     SuspendResult,
 )
-from aws_durable_execution_sdk_python.config import ChildConfig
+from aws_durable_execution_sdk_python.config import ChildConfig, CheckpointMode, NestingType
 from aws_durable_execution_sdk_python.exceptions import (
     OrphanedChildException,
     SuspendExecution,
@@ -143,6 +143,7 @@ def __init__(
         serdes: SerDes | None,
         item_serdes: SerDes | None = None,
         summary_generator: SummaryGenerator | None = None,
+        nesting_type: NestingType = NestingType.NESTED,
     ):
         """Initialize ConcurrentExecutor.
 
@@ -160,6 +161,7 @@ def __init__(
         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()
@@ -412,6 +414,9 @@ def _execute_item_in_child_context(
             executor_context._parent_id,  # noqa: SLF001
             name,
         )
+        checkpoint_mode = CheckpointMode.NO_CHECKPOINT \
+            if self.nesting_type == NestingType.FLAT \
+            else CheckpointMode.CHECKPOINT_AT_START_AND_FINISH
 
         def run_in_child_handler():
             return self.execute_item(child_context, executable)
@@ -424,6 +429,7 @@ def run_in_child_handler():
                 serdes=self.item_serdes or self.serdes,
                 sub_type=self.sub_type_iteration,
                 summary_generator=self.summary_generator,
+                checkpoint_mode=checkpoint_mode,
             ),
         )
         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):
+    """
+    Defines how child operations inherit context from their parent.
+
+    NESTED: Child operations are executed in their own nested context.
+    FLAT: Child operations are executed within the same context as their parent.
+    """
+    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):
@@ -220,7 +237,6 @@ class StepConfig:
 
 class CheckpointMode(Enum):
     NO_CHECKPOINT = ("NO_CHECKPOINT",)
-    CHECKPOINT_AT_FINISH = ("CHECKPOINT_AT_FINISH",)
     CHECKPOINT_AT_START_AND_FINISH = "CHECKPOINT_AT_START_AND_FINISH"
 
 
@@ -259,17 +275,16 @@ 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_mode: controls when checkpoints are created
         - CHECKPOINT_AT_START_AND_FINISH: Checkpoint at both start and completion (default)
-        - CHECKPOINT_AT_FINISH: Only checkpoint when operation completes
+        - CHECKPOINT_AT_FINISH: Only checkpoint when operation completes (not implemented)
         - NO_CHECKPOINT: No automatic checkpointing
 
     See TypeScript reference: aws-durable-execution-sdk-js/src/types/index.ts
     """
 
-    # checkpoint_mode: CheckpointMode = CheckpointMode.CHECKPOINT_AT_START_AND_FINISH
+    checkpoint_mode: CheckpointMode = CheckpointMode.CHECKPOINT_AT_START_AND_FINISH
     serdes: SerDes | None = None
     item_serdes: SerDes | None = None
     sub_type: OperationSubType | None = None
@@ -361,6 +376,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 +395,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/operation/child.py

@@ -5,7 +5,7 @@
 import logging
 from typing import TYPE_CHECKING, TypeVar
 
-from aws_durable_execution_sdk_python.config import ChildConfig
+from aws_durable_execution_sdk_python.config import ChildConfig, CheckpointMode
 from aws_durable_execution_sdk_python.exceptions import (
     InvocationError,
     SuspendExecution,
@@ -118,7 +118,8 @@ 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
+                self.config.checkpoint_mode == CheckpointMode.CHECKPOINT_AT_START_AND_FINISH):
             start_operation: OperationUpdate = OperationUpdate.create_context_start(
                 identifier=self.operation_identifier,
                 sub_type=self.sub_type,
@@ -203,18 +204,19 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T:
                     else ""
                 )
 
-            # Checkpoint SUCCEED
-            success_operation: OperationUpdate = OperationUpdate.create_context_succeed(
-                identifier=self.operation_identifier,
-                payload=serialized_result,
-                sub_type=self.sub_type,
-                context_options=ContextOptions(replay_children=replay_children),
-            )
-            # Checkpoint child context SUCCEED with blocking (is_sync=True, default).
-            # Must ensure the child context result is persisted before returning to the parent.
-            # This guarantees the result is durable and child operations won't be re-executed on replay
-            # (unless replay_children=True for large payloads).
-            self.state.create_checkpoint(operation_update=success_operation)
+            if self.config.checkpoint_mode != CheckpointMode.NO_CHECKPOINT:
+                # Checkpoint SUCCEED
+                success_operation: OperationUpdate = OperationUpdate.create_context_succeed(
+                    identifier=self.operation_identifier,
+                    payload=serialized_result,
+                    sub_type=self.sub_type,
+                    context_options=ContextOptions(replay_children=replay_children),
+                )
+                # Checkpoint child context SUCCEED with blocking (is_sync=True, default).
+                # Must ensure the child context result is persisted before returning to the parent.
+                # This guarantees the result is durable and child operations won't be re-executed on replay
+                # (unless replay_children=True for large payloads).
+                self.state.create_checkpoint(operation_update=success_operation)
 
             logger.debug(
                 "✅ Successfully completed child context for id: %s, name: %s",
@@ -270,8 +272,5 @@ 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:
@@ -46,6 +46,7 @@ def __init__(
         serdes: SerDes | None,
         summary_generator: SummaryGenerator | None = None,
         item_serdes: SerDes | None = None,
+        nesting_type: NestingType = NestingType.NESTED,
     ):
         super().__init__(
             executables=executables,
@@ -57,6 +58,7 @@ def __init__(
             serdes=serdes,
             summary_generator=summary_generator,
             item_serdes=item_serdes,
+            nesting_type=nesting_type,
         )
         self.items = items
 
@@ -83,6 +85,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:

src/aws_durable_execution_sdk_python/operation/parallel.py

@@ -9,7 +9,7 @@
 
 from aws_durable_execution_sdk_python.concurrency.executor import ConcurrentExecutor
 from aws_durable_execution_sdk_python.concurrency.models import Executable
-from aws_durable_execution_sdk_python.config import ParallelConfig
+from aws_durable_execution_sdk_python.config import ParallelConfig, NestingType
 from aws_durable_execution_sdk_python.lambda_service import OperationSubType
 
 if TYPE_CHECKING:
@@ -38,6 +38,7 @@ def __init__(
         serdes: SerDes | None,
         summary_generator: SummaryGenerator | None = None,
         item_serdes: SerDes | None = None,
+        nesting_type: NestingType = NestingType.NESTED,
     ):
         super().__init__(
             executables=executables,
@@ -49,6 +50,7 @@ def __init__(
             serdes=serdes,
             summary_generator=summary_generator,
             item_serdes=item_serdes,
+            nesting_type=nesting_type,
         )
 
     @classmethod
@@ -71,6 +73,7 @@ def from_callables(
             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:  # noqa: PLR6301

tests/concurrency_test.py

@@ -25,15 +25,16 @@
     ExecutableWithState,
     ExecutionCounters,
 )
-from aws_durable_execution_sdk_python.config import CompletionConfig, MapConfig
+from aws_durable_execution_sdk_python.config import CompletionConfig, MapConfig, NestingType, CheckpointMode, \
+    ChildConfig
 from aws_durable_execution_sdk_python.exceptions import (
     CallableRuntimeError,
     InvalidStateError,
     SuspendExecution,
     TimedSuspendExecution,
 )
 from aws_durable_execution_sdk_python.lambda_service import (
-    ErrorObject,
+    ErrorObject, OperationSubType,
 )
 from aws_durable_execution_sdk_python.operation.map import MapExecutor
 
@@ -853,36 +854,63 @@ def test_batch_result_failed_with_none_error():
     assert failed[0].error is not None
 
 
-def test_concurrent_executor_properties():
-    """Test ConcurrentExecutor basic properties."""
+def test_concurrent_executor_nesting_type_parameter():
+    """Test ConcurrentExecutor nesting_type parameter."""
 
     class TestExecutor(ConcurrentExecutor):
         def execute_item(self, child_context, executable):
             return f"result_{executable.index}"
 
-    executables = [Executable(0, lambda: "test"), Executable(1, lambda: "test2")]
-    completion_config = CompletionConfig(
-        min_successful=1,
-        tolerated_failure_count=None,
-        tolerated_failure_percentage=None,
+    executables = [Executable(0, lambda: "test")]
+    completion_config = CompletionConfig(min_successful=1)
+
+    # Test with NESTED (default)
+    executor_nested = TestExecutor(
+        executables=executables,
+        max_concurrency=1,
+        completion_config=completion_config,
+        sub_type_top="TOP",
+        sub_type_iteration="ITER",
+        name_prefix="test_",
+        serdes=None,
+        nesting_type=NestingType.NESTED,
     )
-    executor = TestExecutor(
+    assert executor_nested.nesting_type == NestingType.NESTED
+
+    # Test with FLAT
+    executor_flat = TestExecutor(
         executables=executables,
-        max_concurrency=2,
+        max_concurrency=1,
         completion_config=completion_config,
         sub_type_top="TOP",
         sub_type_iteration="ITER",
         name_prefix="test_",
         serdes=None,
+        nesting_type=NestingType.FLAT,
     )
+    assert executor_flat.nesting_type == NestingType.FLAT
 
-    # Test basic properties
-    assert executor.executables == executables
-    assert executor.max_concurrency == 2
-    assert executor.completion_config == completion_config
-    assert executor.sub_type_top == "TOP"
-    assert executor.sub_type_iteration == "ITER"
-    assert executor.name_prefix == "test_"
+
+def test_concurrent_executor_default_nesting_type():
+    """Test ConcurrentExecutor uses NESTED as default nesting_type."""
+
+    class TestExecutor(ConcurrentExecutor):
+        def execute_item(self, child_context, executable):
+            return f"result_{executable.index}"
+
+    executables = [Executable(0, lambda: "test")]
+    completion_config = CompletionConfig(min_successful=1)
+
+    executor = TestExecutor(
+        executables=executables,
+        max_concurrency=1,
+        completion_config=completion_config,
+        sub_type_top="TOP",
+        sub_type_iteration="ITER",
+        name_prefix="test_",
+        serdes=None,
+    )
+    assert executor.nesting_type == NestingType.NESTED
 
 
 def test_concurrent_executor_full_execution_path():
@@ -2474,8 +2502,10 @@ def execute_item(self, child_context, executable):
     # Track operation_id -> result associations
     captured_associations = []
 
-    def patched_child_handler(func, execution_state, operation_identifier, config):
+    def patched_child_handler(func, execution_state, operation_identifier, config: ChildConfig):
         """Patched child handler that captures operation_id -> result mapping."""
+        assert config.checkpoint_mode == CheckpointMode.NO_CHECKPOINT
+        assert config.sub_type == "TEST_ITER"
         result = func()  # Execute the function
         captured_associations.append((operation_identifier.operation_id, result))
         return result
@@ -2504,6 +2534,7 @@ def patched_child_handler(func, execution_state, operation_identifier, config):
             sub_type_iteration="TEST_ITER",
             name_prefix="test_",
             serdes=None,
+            nesting_type=NestingType.FLAT,
         )
 
         # Create executor context mock

tests/config_test.py

@@ -148,7 +148,6 @@ def test_step_config_with_values():
 def test_checkpoint_mode_enum():
     """Test CheckpointMode enum."""
     assert CheckpointMode.NO_CHECKPOINT.value == ("NO_CHECKPOINT",)
-    assert CheckpointMode.CHECKPOINT_AT_FINISH.value == ("CHECKPOINT_AT_FINISH",)
     assert (
         CheckpointMode.CHECKPOINT_AT_START_AND_FINISH.value
         == "CHECKPOINT_AT_START_AND_FINISH"