move batcher

Co-authored-by: thomas <18520168+yaythomas@users.noreply.github.com>

Author: yaythomas

src/aws_durable_execution_sdk_python/state.py

@@ -716,7 +716,9 @@ def _collect_checkpoint_batch(self) -> list[QueuedOperation]:
                 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
+                        effective_operation_count += (
+                            1  # First empty counts toward limit
+                        )
                         has_empty_checkpoint = True
                     # Subsequent empties don't count toward limit
                 else:
@@ -780,7 +782,9 @@ def _collect_checkpoint_batch(self) -> list[QueuedOperation]:
                 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
+                        effective_operation_count += (
+                            1  # First empty counts toward limit
+                        )
                         has_empty_checkpoint = True
                     # Subsequent empties don't count toward limit
                 else:

tests/e2e/map_with_concurrent_waits_int_test.py

@@ -76,7 +76,9 @@ def _make_tracking_client() -> tuple[Mock, list]:
     calls: list[list] = []
     mock_client = Mock(spec=LambdaClient)
 
-    def _checkpoint(durable_execution_arn, checkpoint_token, updates, client_token=None):
+    def _checkpoint(
+        durable_execution_arn, checkpoint_token, updates, client_token=None
+    ):
         calls.append(list(updates))
         return CheckpointOutput(
             checkpoint_token=f"token_{len(calls)}",
@@ -113,9 +115,9 @@ def test_map_with_concurrent_waits_coalesces_empty_checkpoints():
 
     def branch_work():
         try:
-            start_barrier.wait()          # all start simultaneously
-            state.create_checkpoint()     # empty checkpoint, synchronous
-        except Exception as e:            # noqa: BLE001
+            start_barrier.wait()  # all start simultaneously
+            state.create_checkpoint()  # empty checkpoint, synchronous
+        except Exception as e:  # noqa: BLE001
             errors.append(e)
 
     threads = [threading.Thread(target=branch_work) for _ in range(branch_count)]
@@ -156,7 +158,6 @@ def test_map_with_concurrent_waits_api_call_count_scales_with_real_ops_not_empti
     # limit = 1 (first empty) + 10 (real ops) = 11, so all fit in one batch
     state = _make_state(mock_client, batch_time=5.0, max_ops=11)
 
-
     completion_events: list[CompletionEvent] = []
 
     try:
@@ -173,13 +174,14 @@ def test_map_with_concurrent_waits_api_call_count_scales_with_real_ops_not_empti
                 operation_type=OperationType.STEP,
                 action=OperationAction.START,
             )
-        batcher = ThreadPoolExecutor(max_workers=1)
-        batcher.submit(state.checkpoint_batches_forever)
-   
+
             ev = CompletionEvent()
             completion_events.append(ev)
             state._checkpoint_queue.put(QueuedOperation(op, ev))  # noqa: SLF001
 
+        batcher = ThreadPoolExecutor(max_workers=1)
+        batcher.submit(state.checkpoint_batches_forever)
+
         # Wait for all 410 to be processed
         for ev in completion_events:
             ev.wait()

tests/state_test.py

@@ -3419,7 +3419,7 @@ def test_collect_checkpoint_batch_empty_checkpoints_with_real_ops_respects_limit
     real_in_batch = sum(1 for q in batch if q.operation_update is not None)
 
     assert empty_in_batch == 3  # All empty checkpoints coalesced
-    assert real_in_batch == 4   # 4 real ops (1 slot used by the first empty)
+    assert real_in_batch == 4  # 4 real ops (1 slot used by the first empty)
 
 
 def test_collect_checkpoint_batch_overflow_coalesces_empty_checkpoints():
@@ -3540,9 +3540,11 @@ def test_collect_checkpoint_batch_first_empty_counts_toward_limit():
         operation_type=OperationType.STEP,
         action=OperationAction.START,
     )
-    state._checkpoint_queue.put(QueuedOperation(None, None))   # empty — effective=1
-    state._checkpoint_queue.put(QueuedOperation(op1, None))    # real  — effective=2, limit hit
-    state._checkpoint_queue.put(QueuedOperation(op2, None))    # real  — stays in queue
+    state._checkpoint_queue.put(QueuedOperation(None, None))  # empty — effective=1
+    state._checkpoint_queue.put(
+        QueuedOperation(op1, None)
+    )  # real  — effective=2, limit hit
+    state._checkpoint_queue.put(QueuedOperation(op2, None))  # real  — stays in queue
 
     for _ in range(50):
         state._checkpoint_queue.put(QueuedOperation(None, None))  # trailing empties
@@ -3555,6 +3557,8 @@ def test_collect_checkpoint_batch_first_empty_counts_toward_limit():
     # The batch contains exactly: 1 leading empty + op_1 (limit=2 effective ops)
     assert len(real_in_batch) == 1
     assert real_in_batch[0].operation_update.operation_id == "op_1"
-    assert len(empty_in_batch) == 1  # Only the leading empty; trailing deferred to next batch
+    assert (
+        len(empty_in_batch) == 1
+    )  # Only the leading empty; trailing deferred to next batch
     # op_2 and trailing empties remain in the queue
     assert state._checkpoint_queue.qsize() == 51