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,19 @@ 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 +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,
-                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

@@ -78,10 +78,10 @@ class TerminationMode(Enum):
 
 class NestingType(Enum):
     """
-    Defines how child operations inherit context from their parent.
+    How child operations should 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: Each child operation runs in its own isolated context
+    - FLAT: All child operations share the same parent context
     """
 
     NESTED = "NESTED"
@@ -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,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

@@ -5,7 +5,7 @@
 import logging
 from typing import TYPE_CHECKING, TypeVar
 
-from aws_durable_execution_sdk_python.config import ChildConfig, CheckpointMode
+from aws_durable_execution_sdk_python.config import ChildConfig
 from aws_durable_execution_sdk_python.exceptions import (
     InvocationError,
     SuspendExecution,
@@ -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 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,
@@ -161,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(
@@ -207,21 +211,18 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T:
                     else ""
                 )
 
-            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)
+            # 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",

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()),
     )