feat: add BatchResult serialization support with dedicated codec

Author: vip-amzn

src/aws_durable_execution_sdk_python/concurrency.py

@@ -367,12 +367,12 @@ class ExecutionCounters:
     def __init__(
         self,
         total_tasks: int,
-        min_successful: int,
+        min_successful: int | None,
         tolerated_failure_count: int | None,
         tolerated_failure_percentage: float | None,
     ):
         self.total_tasks: int = total_tasks
-        self.min_successful: int = min_successful
+        self.min_successful: int | None = min_successful
         self.tolerated_failure_count: int | None = tolerated_failure_count
         self.tolerated_failure_percentage: float | None = tolerated_failure_percentage
         self.success_count: int = 0
@@ -421,24 +421,26 @@ def is_complete(self) -> bool:
         """
         Check if execution should complete (based on completion criteria).
         Matches TypeScript isComplete() logic.
+
+        Note: This method only checks completion criteria (all done, or min_successful met).
+        Failure tolerance is enforced separately by should_continue() and combined in should_complete().
         """
         with self._lock:
             completed_count = self.success_count + self.failure_count
 
             # All tasks completed
             if completed_count == self.total_tasks:
-                # Complete if no failure tolerance OR no failures OR min successful reached
-                return (
-                    (
-                        self.tolerated_failure_count is None
-                        and self.tolerated_failure_percentage is None
-                    )
-                    or self.failure_count == 0
-                    or self.success_count >= self.min_successful
-                )
+                # If min_successful is explicitly set, check if we met it
+                # Otherwise, complete when all tasks are done
+                if self.min_successful is not None:
+                    return self.success_count >= self.min_successful
+                return True
 
-            # when we breach min successful, we've completed
-            return self.success_count >= self.min_successful
+            # Early completion: when we breach min_successful (only if explicitly set)
+            return (
+                self.min_successful is not None
+                and self.success_count >= self.min_successful
+            )
 
     def should_complete(self) -> bool:
         """
@@ -455,7 +457,10 @@ def is_all_completed(self) -> bool:
     def is_min_successful_reached(self) -> bool:
         """True if minimum successful tasks reached."""
         with self._lock:
-            return self.success_count >= self.min_successful
+            return (
+                self.min_successful is not None
+                and self.success_count >= self.min_successful
+            )
 
     def is_failure_tolerance_exceeded(self) -> bool:
         """True if failure tolerance was exceeded."""
@@ -594,7 +599,9 @@ def __init__(
         self._suspend_exception: SuspendExecution | None = None
 
         # ExecutionCounters will keep track of completion criteria and on-going counters
-        min_successful = self.completion_config.min_successful or len(self.executables)
+        # Note: min_successful should remain None if not explicitly set
+        # When None, the operation completes when all tasks finish (respecting failure tolerance)
+        min_successful = self.completion_config.min_successful
         tolerated_failure_count = self.completion_config.tolerated_failure_count
         tolerated_failure_percentage = (
             self.completion_config.tolerated_failure_percentage

src/aws_durable_execution_sdk_python/operation/map.py

@@ -82,6 +82,7 @@ def from_items(
             name_prefix="map-item-",
             serdes=config.serdes,
             summary_generator=config.summary_generator,
+            item_serdes=config.item_serdes,
         )
 
     def execute_item(self, child_context, executable: Executable[Callable]) -> R:

src/aws_durable_execution_sdk_python/operation/parallel.py

@@ -69,6 +69,7 @@ def from_callables(
             name_prefix="parallel-branch-",
             serdes=config.serdes,
             summary_generator=config.summary_generator,
+            item_serdes=config.item_serdes,
         )
 
     def execute_item(self, child_context, executable: Executable[Callable]) -> R:  # noqa: PLR6301

src/aws_durable_execution_sdk_python/serdes.py

@@ -62,6 +62,7 @@ class TypeTag(StrEnum):
     TUPLE = "t"
     LIST = "l"
     DICT = "m"
+    BATCH_RESULT = "br"
 
 
 @dataclass(frozen=True)
@@ -206,7 +207,18 @@ def dispatcher(self):
 
     def encode(self, obj: Any) -> EncodedValue:
         """Encode container using dispatcher for recursive elements."""
+        # Import here to avoid circular dependency
+        from aws_durable_execution_sdk_python.concurrency import (
+            BatchResult,
+        )  # noqa: PLC0415
+
         match obj:
+            case BatchResult():
+                # Encode BatchResult as dict with special tag
+                return EncodedValue(
+                    TypeTag.BATCH_RESULT,
+                    self._wrap(obj.to_dict(), self.dispatcher).value,
+                )
             case list():
                 return EncodedValue(
                     TypeTag.LIST, [self._wrap(v, self.dispatcher) for v in obj]
@@ -230,7 +242,16 @@ def encode(self, obj: Any) -> EncodedValue:
 
     def decode(self, tag: TypeTag, value: Any) -> Any:
         """Decode container using dispatcher for recursive elements."""
+        # Import here to avoid circular dependency
+        from aws_durable_execution_sdk_python.concurrency import (
+            BatchResult,
+        )  # noqa: PLC0415
+
         match tag:
+            case TypeTag.BATCH_RESULT:
+                # Decode as dict (handles all recursive unwrapping) then reconstruct
+                decoded_dict = self.decode(TypeTag.DICT, value)
+                return BatchResult.from_dict(decoded_dict)
             case TypeTag.LIST:
                 if not isinstance(value, list):
                     msg = f"Expected list, got {type(value)}"
@@ -281,6 +302,9 @@ def __init__(self):
         self.container_codec.set_dispatcher(self)
 
     def encode(self, obj: Any) -> EncodedValue:
+        # Import here to avoid circular dependency
+        from aws_durable_execution_sdk_python.concurrency import BatchResult
+
         match obj:
             case None | str() | bool() | int() | float():
                 return self.primitive_codec.encode(obj)
@@ -292,7 +316,7 @@ def encode(self, obj: Any) -> EncodedValue:
                 return self.decimal_codec.encode(obj)
             case datetime() | date():
                 return self.datetime_codec.encode(obj)
-            case list() | tuple() | dict():
+            case BatchResult() | list() | tuple() | dict():
                 return self.container_codec.encode(obj)
             case _:
                 msg = f"Unsupported type: {type(obj)}"
@@ -301,11 +325,7 @@ def encode(self, obj: Any) -> EncodedValue:
     def decode(self, tag: TypeTag, value: Any) -> Any:
         match tag:
             case (
-                TypeTag.NONE
-                | TypeTag.STR
-                | TypeTag.BOOL
-                | TypeTag.INT
-                | TypeTag.FLOAT
+                TypeTag.NONE | TypeTag.STR | TypeTag.BOOL | TypeTag.INT | TypeTag.FLOAT
             ):
                 return self.primitive_codec.decode(tag, value)
             case TypeTag.BYTES:
@@ -316,7 +336,7 @@ def decode(self, tag: TypeTag, value: Any) -> Any:
                 return self.decimal_codec.decode(tag, value)
             case TypeTag.DATETIME | TypeTag.DATE:
                 return self.datetime_codec.decode(tag, value)
-            case TypeTag.LIST | TypeTag.TUPLE | TypeTag.DICT:
+            case TypeTag.BATCH_RESULT | TypeTag.LIST | TypeTag.TUPLE | TypeTag.DICT:
                 return self.container_codec.decode(tag, value)
             case _:
                 msg = f"Unknown type tag: {tag}"

tests/concurrency_test.py

@@ -809,6 +809,54 @@ def test_execution_counters_zero_total_tasks():
     assert not counters.is_failure_tolerance_exceeded()
 
 
+def test_execution_counters_none_min_successful():
+    """Test ExecutionCounters with None min_successful completes when all tasks done."""
+    counters = ExecutionCounters(5, None, None, None)
+
+    # Should not complete early
+    assert not counters.should_complete()
+
+    counters.complete_task()
+    counters.complete_task()
+    counters.complete_task()
+    assert not counters.should_complete()
+
+    # Should complete when all tasks are done
+    counters.complete_task()
+    counters.complete_task()
+    assert counters.should_complete()
+
+
+def test_execution_counters_none_min_successful_with_failures():
+    """Test ExecutionCounters with None min_successful and failure tolerance."""
+    counters = ExecutionCounters(5, None, 2, None)
+
+    # Should not complete early even with successes
+    counters.complete_task()
+    counters.complete_task()
+    assert not counters.should_complete()
+
+    # Should complete when failure tolerance exceeded
+    counters.fail_task()
+    counters.fail_task()
+    counters.fail_task()
+    assert counters.should_complete()
+
+
+def test_execution_counters_is_min_successful_reached_with_none():
+    """Test is_min_successful_reached returns False when min_successful is None."""
+    counters = ExecutionCounters(5, None, None, None)
+
+    counters.complete_task()
+    counters.complete_task()
+    counters.complete_task()
+    counters.complete_task()
+    counters.complete_task()
+
+    # Should always return False when min_successful is None
+    assert not counters.is_min_successful_reached()
+
+
 def test_execution_counters_failure_percentage_edge_case():
     """Test ExecutionCounters failure percentage at exact threshold."""
     counters = ExecutionCounters(10, 5, None, 20.0)

tests/operation/map_test.py

@@ -1,5 +1,6 @@
 """Tests for map operation."""
 
+import json
 from unittest.mock import Mock, patch
 
 # Mock the executor.execute method
@@ -750,3 +751,45 @@ def get_checkpoint_result(self, operation_id):
         # Verify replay was called, execute was not
         mock_replay.assert_called_once()
         mock_execute.assert_not_called()
+
+
+def test_map_result_serialization_roundtrip():
+    """Test that map operation BatchResult can be serialized and deserialized."""
+
+    items = ["a", "b", "c"]
+
+    def func(ctx, item, idx, items):
+        return {"item": item.upper(), "index": idx}
+
+    class MockExecutionState:
+        durable_execution_arn = "arn:test"
+
+        def get_checkpoint_result(self, operation_id):
+            mock_result = Mock()
+            mock_result.is_succeeded.return_value = False
+            return mock_result
+
+    execution_state = MockExecutionState()
+    map_context = Mock()
+    map_context._create_step_id_for_logical_step = Mock(side_effect=["1", "2", "3"])  # noqa SLF001
+    map_context.create_child_context = Mock(return_value=Mock())
+    operation_identifier = OperationIdentifier("test_op", "parent", "test_map")
+
+    # Execute map
+    result = map_handler(
+        items, func, MapConfig(), execution_state, map_context, operation_identifier
+    )
+
+    # Serialize the BatchResult
+    serialized = json.dumps(result.to_dict())
+
+    # Deserialize
+    deserialized = BatchResult.from_dict(json.loads(serialized))
+
+    # Verify all data preserved
+    assert len(deserialized.all) == 3
+    assert deserialized.all[0].result == {"item": "A", "index": 0}
+    assert deserialized.all[1].result == {"item": "B", "index": 1}
+    assert deserialized.all[2].result == {"item": "C", "index": 2}
+    assert deserialized.completion_reason == result.completion_reason
+    assert all(item.status == BatchItemStatus.SUCCEEDED for item in deserialized.all)

tests/operation/parallel_test.py

@@ -1,5 +1,6 @@
 """Tests for the parallel operation module."""
 
+import json
 from unittest.mock import Mock, patch
 
 import pytest
@@ -734,3 +735,57 @@ def get_checkpoint_result(self, operation_id):
         # Verify replay was called, execute was not
         mock_replay.assert_called_once()
         mock_execute.assert_not_called()
+
+
+def test_parallel_result_serialization_roundtrip():
+    """Test that parallel operation BatchResult can be serialized and deserialized."""
+
+    def func1(ctx):
+        return [1, 2, 3]
+
+    def func2(ctx):
+        return {"status": "complete", "count": 42}
+
+    def func3(ctx):
+        return "simple string"
+
+    callables = [func1, func2, func3]
+
+    class MockExecutionState:
+        durable_execution_arn = "arn:test"
+
+        def get_checkpoint_result(self, operation_id):
+            mock_result = Mock()
+            mock_result.is_succeeded.return_value = False
+            return mock_result
+
+    execution_state = MockExecutionState()
+    parallel_context = Mock()
+    parallel_context._create_step_id_for_logical_step = Mock(  # noqa SLF001
+        side_effect=["1", "2", "3"]
+    )
+    parallel_context.create_child_context = Mock(return_value=Mock())
+    operation_identifier = OperationIdentifier("test_op", "parent", "test_parallel")
+
+    # Execute parallel
+    result = parallel_handler(
+        callables,
+        ParallelConfig(),
+        execution_state,
+        parallel_context,
+        operation_identifier,
+    )
+
+    # Serialize the BatchResult
+    serialized = json.dumps(result.to_dict())
+
+    # Deserialize
+    deserialized = BatchResult.from_dict(json.loads(serialized))
+
+    # Verify all data preserved
+    assert len(deserialized.all) == 3
+    assert deserialized.all[0].result == [1, 2, 3]
+    assert deserialized.all[1].result == {"status": "complete", "count": 42}
+    assert deserialized.all[2].result == "simple string"
+    assert deserialized.completion_reason == result.completion_reason
+    assert all(item.status == BatchItemStatus.SUCCEEDED for item in deserialized.all)