fix parent id for operations in virtual context

Author: zhongkechen

src/aws_durable_execution_sdk_python/concurrency/executor.py

@@ -22,7 +22,6 @@
 )
 from aws_durable_execution_sdk_python.config import (
     ChildConfig,
-    CheckpointMode,
     NestingType,
 )
 from aws_durable_execution_sdk_python.exceptions import (
@@ -138,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,
@@ -158,6 +158,7 @@ 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
@@ -412,17 +413,13 @@ 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
             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)
@@ -435,7 +432,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,
+                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

@@ -235,11 +235,6 @@ class StepConfig:
     serdes: SerDes | None = None
 
 
-class CheckpointMode(Enum):
-    NO_CHECKPOINT = ("NO_CHECKPOINT",)
-    CHECKPOINT_AT_START_AND_FINISH = "CHECKPOINT_AT_START_AND_FINISH"
-
-
 @dataclass(frozen=True)
 class ChildConfig(Generic[T]):
     """Configuration options for child context operations.
@@ -276,19 +271,18 @@ class ChildConfig(Generic[T]):
             Used internally by map/parallel operations to handle large BatchResult payloads.
             Signature: (result: T) -> str
 
-        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 (not implemented)
-        - 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):

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,7 @@ 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 +285,7 @@ def create_child_context(self, parent_id: str) -> DurableContext:
                     parent_id=parent_id,
                 )
             ),
+            non_virtual_parent_id=non_virtual_parent_id,
         )
 
     # endregion factories

src/aws_durable_execution_sdk_python/operation/child.py

@@ -118,11 +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()
-            and self.config.checkpoint_mode
-            == CheckpointMode.CHECKPOINT_AT_START_AND_FINISH
-        ):
+        if not self.config.is_virtual:
             start_operation: OperationUpdate = OperationUpdate.create_context_start(
                 identifier=self.operation_identifier,
                 sub_type=self.sub_type,
@@ -207,7 +203,7 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T:
                     else ""
                 )
 
-            if self.config.checkpoint_mode != CheckpointMode.NO_CHECKPOINT:
+            if not self.config.is_virtual:
                 # Checkpoint SUCCEED
                 success_operation: OperationUpdate = (
                     OperationUpdate.create_context_succeed(

src/aws_durable_execution_sdk_python/operation/map.py

@@ -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,
@@ -49,6 +50,7 @@ def __init__(
         nesting_type: NestingType = NestingType.NESTED,
     ):
         super().__init__(
+            operation_identifier=operation_identifier,
             executables=executables,
             max_concurrency=max_concurrency,
             completion_config=completion_config,
@@ -65,6 +67,7 @@ def __init__(
     @classmethod
     def from_items(
         cls,
+        operation_identifier: OperationIdentifier,
         items: Sequence[T],
         func: Callable,
         config: MapConfig,
@@ -75,6 +78,7 @@ def from_items(
         ]
 
         return cls(
+            operation_identifier=operation_identifier,
             executables=executables,
             items=items,
             max_concurrency=config.max_concurrency,
@@ -112,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()),

src/aws_durable_execution_sdk_python/operation/parallel.py

@@ -29,6 +29,7 @@
 class ParallelExecutor(ConcurrentExecutor[Callable, R]):
     def __init__(
         self,
+        operation_identifier: OperationIdentifier,
         executables: list[Executable[Callable]],
         max_concurrency: int | None,
         completion_config,
@@ -41,6 +42,7 @@ def __init__(
         nesting_type: NestingType = NestingType.NESTED,
     ):
         super().__init__(
+            operation_identifier=operation_identifier,
             executables=executables,
             max_concurrency=max_concurrency,
             completion_config=completion_config,
@@ -56,6 +58,7 @@ def __init__(
     @classmethod
     def from_callables(
         cls,
+        operation_identifier: OperationIdentifier,
         callables: Sequence[Callable],
         config: ParallelConfig,
     ) -> ParallelExecutor:
@@ -64,6 +67,7 @@ def from_callables(
             Executable(index=i, func=func) for i, func in enumerate(callables)
         ]
         return cls(
+            operation_identifier=operation_identifier,
             executables=executables,
             max_concurrency=config.max_concurrency,
             completion_config=config.completion_config,
@@ -98,6 +102,7 @@ def parallel_handler(
     # See TypeScript reference: aws-durable-execution-sdk-js/src/handlers/parallel-handler/parallel-handler.ts (~line 112)
 
     executor = ParallelExecutor.from_callables(
+        operation_identifier,
         callables,
         config or ParallelConfig(summary_generator=ParallelSummaryGenerator()),
     )

tests/concurrency_test.py

@@ -881,7 +881,7 @@ def execute_item(self, child_context, executable):
         serdes=None,
         nesting_type=NestingType.NESTED,
     )
-    assert executor_nested.nesting_type == NestingType.NESTED
+    assert executor_nested.nesting_type is NestingType.NESTED
 
     # Test with FLAT
     executor_flat = TestExecutor(
@@ -894,7 +894,7 @@ def execute_item(self, child_context, executable):
         serdes=None,
         nesting_type=NestingType.FLAT,
     )
-    assert executor_flat.nesting_type == NestingType.FLAT
+    assert executor_flat.nesting_type is NestingType.FLAT
 
 
 def test_concurrent_executor_default_nesting_type():
@@ -916,7 +916,7 @@ def execute_item(self, child_context, executable):
         name_prefix="test_",
         serdes=None,
     )
-    assert executor.nesting_type == NestingType.NESTED
+    assert executor.nesting_type is NestingType.NESTED
 
 
 def test_concurrent_executor_full_execution_path():

tests/operation/map_test.py

@@ -67,7 +67,7 @@ def test_map_executor_init():
 
     assert executor.items == items
     assert executor.executables == executables
-    assert executor.nesting_type == NestingType.FLAT
+    assert executor.nesting_type is NestingType.FLAT
 
 
 def test_map_executor_from_items():
@@ -85,7 +85,7 @@ def callable_func(ctx, item, idx, items):
     assert executor.items == items
     assert all(exe.func == callable_func for exe in executor.executables)
     assert [exe.index for exe in executor.executables] == [0, 1, 2]
-    assert executor.nesting_type == NestingType.FLAT
+    assert executor.nesting_type is NestingType.FLAT
 
 
 def test_map_executor_from_items_default_config():
@@ -99,7 +99,7 @@ def callable_func(ctx, item, idx, items):
 
     assert len(executor.executables) == 1
     assert executor.items == items
-    assert executor.nesting_type == NestingType.NESTED
+    assert executor.nesting_type is NestingType.NESTED
 
 
 @patch("aws_durable_execution_sdk_python.operation.map.logger")

tests/operation/parallel_test.py

@@ -74,7 +74,7 @@ def test_parallel_executor_init():
     assert executor.sub_type_top == OperationSubType.PARALLEL
     assert executor.sub_type_iteration == OperationSubType.PARALLEL_BRANCH
     assert executor.name_prefix == "test-"
-    assert executor.nesting_type == NestingType.FLAT
+    assert executor.nesting_type is NestingType.FLAT
 
 
 def test_parallel_executor_from_callables():
@@ -100,7 +100,7 @@ def func2(ctx):
     assert executor.sub_type_top == OperationSubType.PARALLEL
     assert executor.sub_type_iteration == OperationSubType.PARALLEL_BRANCH
     assert executor.name_prefix == "parallel-branch-"
-    assert executor.nesting_type == NestingType.FLAT
+    assert executor.nesting_type is NestingType.FLAT
 
 
 def test_parallel_executor_from_callables_default_config():
@@ -117,7 +117,7 @@ def func1(ctx):
     assert len(executor.executables) == 1
     assert executor.max_concurrency is None
     assert executor.completion_config == CompletionConfig.all_successful()
-    assert executor.nesting_type == NestingType.NESTED
+    assert executor.nesting_type is NestingType.NESTED
 
 
 def test_parallel_executor_execute_item():