Merge branch 'main' into exit-paths

Author: Astraea Quinn S

src/aws_durable_execution_sdk_python/concurrency.py

@@ -5,11 +5,11 @@
 import heapq
 import logging
 import threading
-import time
 from abc import ABC, abstractmethod
 from collections import Counter
 from concurrent.futures import Future, ThreadPoolExecutor
 from dataclasses import dataclass
+from datetime import UTC, datetime
 from enum import Enum
 from typing import TYPE_CHECKING, Generic, Self, TypeVar
 
@@ -258,7 +258,7 @@ def __init__(self, executable: Executable[CallableType]):
         self.executable = executable
         self._status = BranchStatus.PENDING
         self._future: Future | None = None
-        self._suspend_until: float | None = None
+        self._suspend_until: datetime | None = None
         self._result: ResultType = None  # type: ignore[assignment]
         self._is_result_set: bool = False
         self._error: Exception | None = None
@@ -293,7 +293,7 @@ def error(self) -> Exception:
         return self._error
 
     @property
-    def suspend_until(self) -> float | None:
+    def suspend_until(self) -> datetime | None:
         """Get suspend timestamp."""
         return self._suspend_until
 
@@ -308,7 +308,7 @@ def can_resume(self) -> bool:
         return self._status is BranchStatus.SUSPENDED or (
             self._status is BranchStatus.SUSPENDED_WITH_TIMEOUT
             and self._suspend_until is not None
-            and time.time() >= self._suspend_until
+            and datetime.now(UTC) >= self._suspend_until
         )
 
     @property
@@ -333,7 +333,7 @@ def suspend(self) -> None:
         self._status = BranchStatus.SUSPENDED
         self._suspend_until = None
 
-    def suspend_with_timeout(self, timestamp: float) -> None:
+    def suspend_with_timeout(self, timestamp: datetime) -> None:
         """Transition to SUSPENDED_WITH_TIMEOUT state."""
         self._status = BranchStatus.SUSPENDED_WITH_TIMEOUT
         self._suspend_until = timestamp
@@ -507,11 +507,11 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         self.shutdown()
 
     def schedule_resume(
-        self, exe_state: ExecutableWithState, resume_time: float
+        self, exe_state: ExecutableWithState, resume_time: datetime
     ) -> None:
         """Schedule a task to resume at the specified time."""
         with self._lock:
-            heapq.heappush(self._pending_resumes, (resume_time, exe_state))
+            heapq.heappush(self._pending_resumes, (resume_time.timestamp(), exe_state))
 
     def shutdown(self) -> None:
         """Shutdown the timer thread and cancel all pending resumes."""
@@ -534,7 +534,7 @@ def _timer_loop(self) -> None:
                 self._shutdown.wait(timeout=0.1)
                 continue
 
-            current_time = time.time()
+            current_time = datetime.now(UTC).timestamp()
             if current_time >= next_resume_time:
                 # Time to resume
                 with self._lock:
@@ -566,6 +566,7 @@ def __init__(
         sub_type_iteration: OperationSubType,
         name_prefix: str,
         serdes: SerDes | None,
+        item_serdes: SerDes | None = None,
         summary_generator: SummaryGenerator | None = None,
     ):
         """Initialize ConcurrentExecutor.
@@ -604,6 +605,7 @@ def __init__(
         )
         self.executables_with_state: list[ExecutableWithState] = []
         self.serdes = serdes
+        self.item_serdes = item_serdes
 
     @abstractmethod
     def execute_item(
@@ -673,7 +675,7 @@ def on_done(future: Future) -> None:
 
     def should_execution_suspend(self) -> SuspendResult:
         """Check if execution should suspend."""
-        earliest_timestamp: float = float("inf")
+        earliest_timestamp: datetime | None = None
         indefinite_suspend_task: (
             ExecutableWithState[CallableType, ResultType] | None
         ) = None
@@ -683,16 +685,16 @@ def should_execution_suspend(self) -> SuspendResult:
                 # Exit here! Still have tasks that can make progress, don't suspend.
                 return SuspendResult.do_not_suspend()
             if exe_state.status is BranchStatus.SUSPENDED_WITH_TIMEOUT:
-                if (
-                    exe_state.suspend_until
-                    and exe_state.suspend_until < earliest_timestamp
+                if exe_state.suspend_until and (
+                    earliest_timestamp is None
+                    or exe_state.suspend_until < earliest_timestamp
                 ):
                     earliest_timestamp = exe_state.suspend_until
             elif exe_state.status is BranchStatus.SUSPENDED:
                 indefinite_suspend_task = exe_state
 
         # All tasks are in final states and at least one of them is a suspend.
-        if earliest_timestamp != float("inf"):
+        if earliest_timestamp is not None:
             return SuspendResult.suspend(
                 TimedSuspendExecution(
                     "All concurrent work complete or suspended pending retry.",
@@ -797,7 +799,7 @@ def execute_in_child_context(child_context: DurableContext) -> ResultType:
             execute_in_child_context,
             f"{self.name_prefix}{executable.index}",
             ChildConfig(
-                serdes=self.serdes,
+                serdes=self.item_serdes or self.serdes,
                 sub_type=self.sub_type_iteration,
                 summary_generator=self.summary_generator,
             ),

src/aws_durable_execution_sdk_python/config.py

@@ -125,9 +125,21 @@ class ParallelConfig:
             Default is CompletionConfig.all_successful() which requires all branches
             to succeed. Other options include first_successful() and all_completed().
 
-        serdes: Custom serialization/deserialization configuration for parallel results.
-            If None, uses the default serializer. This allows custom handling of
-            complex result types or optimization for large result sets.
+        serdes: Custom serialization/deserialization configuration for BatchResult.
+            Applied at the handler level to serialize the entire BatchResult object.
+            If None, uses the default JSON serializer for BatchResult.
+
+            Backward Compatibility: If only 'serdes' is provided (no item_serdes),
+            it will be used for both individual functions AND BatchResult serialization
+            to maintain existing behavior.
+
+        item_serdes: Custom serialization/deserialization configuration for individual functions.
+            Applied to each function's result as tasks complete in child contexts.
+            If None, uses the default JSON serializer for individual function results.
+
+            When both 'serdes' and 'item_serdes' are provided:
+            - item_serdes: Used for individual function results in child contexts
+            - serdes: Used for the entire BatchResult at handler level
 
         summary_generator: Function to generate compact summaries for large results (>256KB).
             When the serialized result exceeds CHECKPOINT_SIZE_LIMIT, this generator
@@ -150,6 +162,7 @@ class ParallelConfig:
         default_factory=CompletionConfig.all_successful
     )
     serdes: SerDes | None = None
+    item_serdes: SerDes | None = None
     summary_generator: SummaryGenerator | None = None
 
 
@@ -181,9 +194,21 @@ class ChildConfig(Generic[T]):
     matching the TypeScript ChildConfig interface behavior.
 
     Args:
-        serdes: Custom serialization/deserialization configuration for the child context data.
-            If None, uses the default serializer. This allows different serialization
-            strategies for child operations vs parent operations.
+        serdes: Custom serialization/deserialization configuration for BatchResult.
+            Applied at the handler level to serialize the entire BatchResult object.
+            If None, uses the default JSON serializer for BatchResult.
+
+            Backward Compatibility: If only 'serdes' is provided (no item_serdes),
+            it will be used for both individual items AND BatchResult serialization
+            to maintain existing behavior.
+
+        item_serdes: Custom serialization/deserialization configuration for individual items.
+            Applied to each item's result as tasks complete in child contexts.
+            If None, uses the default JSON serializer for individual items.
+
+            When both 'serdes' and 'item_serdes' are provided:
+            - item_serdes: Used for individual item results in child contexts
+            - serdes: Used for the entire BatchResult at handler level
 
         sub_type: Operation subtype identifier used for tracking and debugging.
             Examples: OperationSubType.MAP_ITERATION, OperationSubType.PARALLEL_BRANCH.
@@ -208,6 +233,7 @@ class ChildConfig(Generic[T]):
 
     # 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
 
@@ -273,9 +299,21 @@ class MapConfig:
             Default allows any number of failures. Use CompletionConfig.all_successful()
             to require all items to succeed.
 
-        serdes: Custom serialization/deserialization configuration for map results.
-            If None, uses the default serializer. This allows custom handling of
-            complex item types or optimization for large result collections.
+        serdes: Custom serialization/deserialization configuration for BatchResult.
+            Applied at the handler level to serialize the entire BatchResult object.
+            If None, uses the default JSON serializer for BatchResult.
+
+            Backward Compatibility: If only 'serdes' is provided (no item_serdes),
+            it will be used for both individual items AND BatchResult serialization
+            to maintain existing behavior.
+
+        item_serdes: Custom serialization/deserialization configuration for individual items.
+            Applied to each item's result as tasks complete in child contexts.
+            If None, uses the default JSON serializer for individual items.
+
+            When both 'serdes' and 'item_serdes' are provided:
+            - item_serdes: Used for individual item results in child contexts
+            - serdes: Used for the entire BatchResult at handler level
 
         summary_generator: Function to generate compact summaries for large results (>256KB).
             When the serialized result exceeds CHECKPOINT_SIZE_LIMIT, this generator
@@ -298,6 +336,7 @@ class MapConfig:
     item_batcher: ItemBatcher = field(default_factory=ItemBatcher)
     completion_config: CompletionConfig = field(default_factory=CompletionConfig)
     serdes: SerDes | None = None
+    item_serdes: SerDes | None = None
     summary_generator: SummaryGenerator | None = None
 
 

src/aws_durable_execution_sdk_python/context.py

@@ -311,7 +311,7 @@ def map(
         """Execute a callable for each item in parallel."""
         map_name: str | None = self._resolve_step_name(name, func)
 
-        def map_in_child_context(child_context):
+        def map_in_child_context(child_context) -> BatchResult[R]:
             return map_handler(
                 items=inputs,
                 func=func,
@@ -323,7 +323,10 @@ def map_in_child_context(child_context):
         return self.run_in_child_context(
             func=map_in_child_context,
             name=map_name,
-            config=ChildConfig(sub_type=OperationSubType.MAP),
+            config=ChildConfig(
+                sub_type=OperationSubType.MAP,
+                serdes=config.serdes if config is not None else None,
+            ),
         )
 
     def parallel(
@@ -345,7 +348,10 @@ def parallel_in_child_context(child_context) -> BatchResult[T]:
         return self.run_in_child_context(
             func=parallel_in_child_context,
             name=name,
-            config=ChildConfig(sub_type=OperationSubType.PARALLEL),
+            config=ChildConfig(
+                sub_type=OperationSubType.PARALLEL,
+                serdes=config.serdes if config is not None else None,
+            ),
         )
 
     def run_in_child_context(

src/aws_durable_execution_sdk_python/exceptions.py

@@ -5,13 +5,12 @@
 
 from __future__ import annotations
 
-import time
 from dataclasses import dataclass
+
 from enum import Enum
 from typing import TYPE_CHECKING
+from datetime import UTC, datetime, timedelta
 
-if TYPE_CHECKING:
-    import datetime
 
 
 class TerminationReason(Enum):
@@ -150,10 +149,10 @@ class TimedSuspendExecution(SuspendExecution):
     This is a specialized form of SuspendExecution that includes a scheduled resume time.
 
     Attributes:
-        scheduled_timestamp (float): Unix timestamp in seconds at which to resume.
+        scheduled_timestamp (datetime): DateTime at which to resume.
     """
 
-    def __init__(self, message: str, scheduled_timestamp: float):
+    def __init__(self, message: str, scheduled_timestamp: datetime):
         super().__init__(message)
         self.scheduled_timestamp = scheduled_timestamp
 
@@ -172,23 +171,23 @@ def from_delay(cls, message: str, delay_seconds: int) -> TimedSuspendExecution:
             >>> exception = TimedSuspendExecution.from_delay("Waiting for callback", 30)
             >>> # Will suspend for 30 seconds from now
         """
-        resume_time = time.time() + delay_seconds
+        resume_time = datetime.now(UTC) + timedelta(seconds=delay_seconds)
         return cls(message, scheduled_timestamp=resume_time)
 
     @classmethod
     def from_datetime(
-        cls, message: str, datetime_timestamp: datetime.datetime
+        cls, message: str, datetime_timestamp: datetime
     ) -> TimedSuspendExecution:
         """Create a timed suspension with the delay calculated from now.
 
         Args:
             message: Descriptive message for the suspension
-            datetime_timestamp: Unix datetime timestamp in seconds at which to resume
+            datetime_timestamp: DateTime timestamp at which to resume
 
         Returns:
             TimedSuspendExecution: Instance with calculated resume time
         """
-        return cls(message, scheduled_timestamp=datetime_timestamp.timestamp())
+        return cls(message, scheduled_timestamp=datetime_timestamp)
 
 
 class OrderedLockError(DurableExecutionsError):
@@ -243,3 +242,7 @@ def __str__(self) -> str:
             A string in the format "type: message"
         """
         return f"{self.type}: {self.message}"
+
+
+class SerDesError(DurableExecutionsError):
+    """Raised when serialization fails."""

src/aws_durable_execution_sdk_python/execution.py

@@ -4,7 +4,9 @@
 import logging
 from dataclasses import dataclass
 from enum import Enum
+from functools import wraps
 from typing import TYPE_CHECKING, Any
+from warnings import deprecated
 
 from aws_durable_execution_sdk_python.context import DurableContext, ExecutionState
 from aws_durable_execution_sdk_python.exceptions import (
@@ -188,7 +190,7 @@ def create_succeeded(cls, result: str) -> DurableExecutionInvocationOutput:
 # endregion Invocation models
 
 
-def durable_handler(
+def durable_execution(
     func: Callable[[Any, DurableContext], Any],
 ) -> Callable[[Any, LambdaContext], Any]:
     logger.debug("Starting durable execution handler...")
@@ -315,3 +317,9 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
             ).to_dict()
 
     return wrapper
+
+
+@deprecated("Use `durable_execution` instead.")
+@wraps(durable_execution)
+def durable_handler(*args, **kwargs):
+    return durable_execution(*args, **kwargs)

src/aws_durable_execution_sdk_python/operation/map.py

@@ -21,7 +21,6 @@
     from aws_durable_execution_sdk_python.state import ExecutionState
     from aws_durable_execution_sdk_python.types import DurableContext, SummaryGenerator
 
-
 logger = logging.getLogger(__name__)
 
 # Input item type
@@ -42,6 +41,7 @@ def __init__(
         name_prefix: str,
         serdes: SerDes | None,
         summary_generator: SummaryGenerator | None = None,
+        item_serdes: SerDes | None = None,
     ):
         super().__init__(
             executables=executables,
@@ -52,6 +52,7 @@ def __init__(
             name_prefix=name_prefix,
             serdes=serdes,
             summary_generator=summary_generator,
+            item_serdes=item_serdes,
         )
         self.items = items