[Feature]: concurrency performance improvement on nested waits/retries

[yaythomas]

What would you like?

Optimize Empty Checkpoint Coalescing in Checkpoint Batcher

Problem

Checkpoint batching coalesces empty checkpoints only AFTER batching rather than BEFORE.

When many map/parallel iterations resume from wait operations simultaneously, each queues a separate empty checkpoint operation. The current implementation collects these into batches of up to 250 operations, then filters out the empty checkpoints before making the API call. This results in unnecessary batches and API calls.

This impacts on map/parallel with wait, retry, or wait_for_condition inside it:

• High-concurrency scenarios (100+ concurrent branches/items)
• Operations that suspend (wait, wait_for_condition, step retries)
• Multiple resume cycles (polling, retries)

How It Happens in the Code

Step 1: Empty checkpoints are queued separately

When map iterations resume from wait, the resubmitter callback is invoked:

# In concurrency/executor.py
def resubmitter(executable_with_state: ExecutableWithState) -> None:
    execution_state.create_checkpoint()  # operation_update=None (empty checkpoint)
    submit_task(executable_with_state)

Each call to create_checkpoint() with no arguments queues a separate QueuedOperation:

# In state.py - create_checkpoint()
completion_event = CompletionEvent() if is_sync else None
queued_op = QueuedOperation(operation_update=None, completion_event)  # Separate object per call
self._checkpoint_queue.put(queued_op)  # Each queued individually

Step 2: Batcher collects up to 250 operations

The background thread collects operations into batches:

# In state.py - _collect_checkpoint_batch()
batch: list[QueuedOperation] = []
# ... collects up to max_batch_operations (250) ...
return batch  # Returns 250 QueuedOperation objects

Step 3: Empty checkpoints are filtered AFTER batching

Only when preparing the API call are empty checkpoints filtered:

# In state.py - checkpoint_batches_forever()
batch = self._collect_checkpoint_batch()  # 250 operations collected

# Filter happens HERE - after batching
updates = [q.operation_update for q in batch if q.operation_update is not None]

# API call with filtered updates
output = self._service_client.checkpoint(
    updates=updates,  # Empty list [] if all were empty checkpoints
    ...
)

The Inefficiency

For 999 empty checkpoints:

1. 999 separate QueuedOperation objects queued
2. Batcher collects them in 4 batches: 250 + 250 + 250 + 249
3. Each batch filters empties: updates = []
4. 4 API calls made with updates=[]
5. Each batch waits 1 second for batching window

Result: 4 API calls, ~4 seconds processing time

Optimal: 1 API call, ~1.2 seconds processing time

Impact

This affects high-concurrency map/parallel operations with wait operations:

• API overhead: 3-4x more API calls than necessary
• Memory usage: 999 queue objects instead of 1 marker + 999 events
• Processing time: 3x slower batch processing
• Latency: Items resume serially through batches instead of simultaneously

Possible Implementation

No response

Is this a breaking change?

No

Does this require an RFC?

No

Additional Context

No response

[yaythomas]

Solution: Modify _collect_checkpoint_batch() to track effective_operation_count separately from len(batch):

First empty checkpoint: Counts toward the 250 operation limit
Subsequent empty checkpoints: Added to batch but DON'T count toward limit
All empty checkpoints remain in batch for completion signaling
When sending to API, deduplicate empties (they're excluded from updates list)

Implementation: Two methods in state.py:

two key changes:

• _collect_checkpoint_batch(): Track effective_operation_count and only count first empty toward limit
• checkpoint_batches_forever(): Deduplicate empties when extracting updates, signal all operations in batch. Count empties and log clearly what's being sent vs excluded

something like this (untested):

diff --git a/src/aws_durable_execution_sdk_python/state.py b/src/aws_durable_execution_sdk_python/state.py
index abc123..def456 100644
--- a/src/aws_durable_execution_sdk_python/state.py
+++ b/src/aws_durable_execution_sdk_python/state.py
@@ -685,12 +685,14 @@ class ExecutionState:
         """
         batch: list[QueuedOperation] = []
+        has_empty_checkpoint = False
         total_size = 0
+        effective_operation_count = 0  # Operations that count toward batch limit
 
         # First, drain overflow queue (FIFO order preserved)
         try:
-            while len(batch) < self._batcher_config.max_batch_operations:
+            while effective_operation_count < self._batcher_config.max_batch_operations:
                 overflow_op = self._overflow_queue.get_nowait()
-                op_size = self._calculate_operation_size(overflow_op)
 
+                if overflow_op.operation_update is None:  # Empty checkpoint
+                    batch.append(overflow_op)
+                    if not has_empty_checkpoint:
+                        effective_operation_count += 1  # First empty counts toward limit
+                        has_empty_checkpoint = True
+                    # Subsequent empties don't count toward limit
+                else:
+                    op_size = self._calculate_operation_size(overflow_op)
                 if total_size + op_size > self._batcher_config.max_batch_size_bytes:
                     self._overflow_queue.put(overflow_op)
                     break
                 batch.append(overflow_op)
                 total_size += op_size
+                    effective_operation_count += 1
         except queue.Empty:
             pass
 
@@ -702,7 +710,12 @@ class ExecutionState:
                     timeout=0.1
                 )
                 self._checkpoint_queue.task_done()
                 batch.append(first_op)
-                total_size += self._calculate_operation_size(first_op)
+                
+                if first_op.operation_update is None:
+                    has_empty_checkpoint = True
+                else:
+                    total_size += self._calculate_operation_size(first_op)
+                
+                effective_operation_count = 1
                 break
             except queue.Empty:
                 continue
@@ -715,7 +728,7 @@ class ExecutionState:
 
         while (
             time.time() < batch_deadline
-            and len(batch) < self._batcher_config.max_batch_operations
+            and effective_operation_count < self._batcher_config.max_batch_operations
             and not self._checkpointing_stopped.is_set()
         ):
             remaining_time = min(
@@ -728,7 +741,16 @@ class ExecutionState:
             try:
                 additional_op = self._checkpoint_queue.get(timeout=remaining_time)
                 self._checkpoint_queue.task_done()
-                op_size = self._calculate_operation_size(additional_op)
+                
+                if additional_op.operation_update is None:  # Empty checkpoint
+                    batch.append(additional_op)
+                    if not has_empty_checkpoint:
+                        effective_operation_count += 1  # First empty counts toward limit
+                        has_empty_checkpoint = True
+                    # Subsequent empties don't count toward limit
+                else:
+                    op_size = self._calculate_operation_size(additional_op)
 
                 if total_size + op_size > self._batcher_config.max_batch_size_bytes:
                     self._overflow_queue.put(additional_op)
@@ -737,6 +759,7 @@ class ExecutionState:
                     break
                 batch.append(additional_op)
                 total_size += op_size
+                    effective_operation_count += 1
 
             except queue.Empty:
                 break
@@ -744,6 +767,7 @@ class ExecutionState:
+        empty_count = sum(1 for q in batch if q.operation_update is None)
         logger.debug(
-            "Collected batch of %d operations, total size: %d bytes",
+            "Collected batch of %d operations (%d effective, %d non-empty, %d empty), total size: %d bytes",
             len(batch),
+            effective_operation_count,
+            len(batch) - empty_count,
+            empty_count,
             total_size,
         )
@@ -590,14 +615,21 @@ class ExecutionState:
             batch: list[QueuedOperation] = self._collect_checkpoint_batch()
 
             if batch:
-                # Extract OperationUpdates from QueuedOperations for API call
-                updates: list[OperationUpdate] = [
-                    q.operation_update for q in batch if q.operation_update is not None
-                ]
+                # Extract OperationUpdates, deduplicating empty checkpoints
+                updates: list[OperationUpdate] = []
+                empty_count = 0
+                
+                for q in batch:
+                    if q.operation_update is not None:
+                        updates.append(q.operation_update)
+                    else:
+                        empty_count += 1
 
                 logger.debug(
-                    "Processing checkpoint batch with %d operations (%d non-empty)",
+                    "Sending %d OperationUpdates out of %d operations, excluding %d empty checkpoints",
+                    len(updates),
                     len(batch),
-                    len(updates),
+                    empty_count,
                 )