diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py index e9a7588d2..a9883b3ea 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py @@ -154,6 +154,12 @@ def is_full(self) -> bool: return self.effective_count >= self._config.max_batch_operations def fits(self, size: int) -> bool: + # An empty batch accepts an operation of any size. An operation above + # the limit can only ever be sent alone, so refusing it here would + # return it to its queue on every pass, and its sync caller would block + # forever. The service, not the batcher, enforces the per-operation limit. + if self.is_empty(): + return True return self.total_size + size <= self._config.max_batch_size_bytes def add(self, op: QueuedOperation, size: int = 0) -> None: @@ -1487,8 +1493,7 @@ def _wait_for_first_operation(self, batch: _BatchAccumulator) -> bool: if isinstance(item, _Signal): self._consume_refresh_wake() continue - # The first operation is never refused for size. It could not be - # sent otherwise. + # The first operation is never refused for size; see fits(). batch.add(item, self._calculate_operation_size(item)) return True return False diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/large_step_result_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/large_step_result_int_test.py new file mode 100644 index 000000000..0d7eb3ac3 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/large_step_result_int_test.py @@ -0,0 +1,36 @@ +"""A step result above the checkpoint batch size limit completes (issue #871). + +The step result is 800 KB. That is above ``max_batch_size_bytes`` (750 KB) and +below the service's 1 MB per-operation limit. With the default +``AT_LEAST_ONCE_PER_RETRY`` semantics the START checkpoint is async, so the +batcher still holds it when the sync SUCCEED arrives, and SUCCEED takes the +overflow path. Before the fix the SUCCEED operation was never sent and the +handler thread blocked forever, so ``run()`` raised ``TimeoutError``. +""" + +from __future__ import annotations + +import json + +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.execution import InvocationStatus +from aws_durable_execution_sdk_python_testing.runner import ( + DurableFunctionTestResult, + DurableFunctionTestRunner, +) + +RESULT_SIZE_BYTES = 800 * 1024 + + +@durable_execution +def _large_step_handler(_event: dict[str, str], context: DurableContext) -> str: + return context.step(lambda _sc: "y" * RESULT_SIZE_BYTES, name="big_step") + + +def test_step_result_above_batch_size_limit_completes() -> None: + with DurableFunctionTestRunner(handler=_large_step_handler) as runner: + result: DurableFunctionTestResult = runner.run(input="{}", execution_timeout=30) + + assert result.status is InvocationStatus.SUCCEEDED + assert result.result is not None + assert len(json.loads(result.result)) == RESULT_SIZE_BYTES diff --git a/packages/aws-durable-execution-sdk-python/tests/state_test.py b/packages/aws-durable-execution-sdk-python/tests/state_test.py index 128aa33cd..45c41e53c 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -56,6 +56,7 @@ from aws_durable_execution_sdk_python.state import ( _BLOCKED_CALLER_WAIT_SECONDS, _STOP_SIGNAL_POLL_SECONDS, + _BatchAccumulator, _Signal, CheckpointBatcherConfig, CheckpointedResult, @@ -1367,6 +1368,96 @@ def test_collect_checkpoint_batch_uses_overflow_queue(): assert batch[1].operation_update.operation_id == "overflow_2" +def _oversized_queued_operation( + operation_id: str, max_batch_size_bytes: int +) -> QueuedOperation: + """A sync operation whose serialized size alone exceeds the batch limit.""" + return QueuedOperation( + OperationUpdate( + operation_id=operation_id, + operation_type=OperationType.STEP, + action=OperationAction.SUCCEED, + payload="y" * (2 * max_batch_size_bytes), + ), + CompletionEvent(), + ) + + +def test_collect_checkpoint_batch_sends_oversized_operation_from_overflow() -> None: + """An operation above max_batch_size_bytes that reached the overflow queue is + still sent, alone, in the next batch (issue #871). + + This is the default AT_LEAST_ONCE_PER_RETRY step shape: a small START is + already in the batch when the large SUCCEED arrives, so SUCCEED does not fit + and moves to the overflow queue. The next batch must send it rather than + return it to the queue forever. + """ + mock_lambda_client = Mock(spec=LambdaClient) + config = CheckpointBatcherConfig(max_batch_size_bytes=500) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + batcher_config=config, + ) + + small_op = QueuedOperation( + OperationUpdate( + operation_id="big_step", + operation_type=OperationType.STEP, + action=OperationAction.START, + ), + None, + ) + large_op = _oversized_queued_operation("big_step", config.max_batch_size_bytes) + assert state._calculate_operation_size(large_op) > config.max_batch_size_bytes + state._checkpoint_queue.put(small_op) + state._checkpoint_queue.put(large_op) + + first_batch = state._collect_checkpoint_batch() + assert first_batch == [small_op] + assert state._overflow_queue.qsize() == 1 + + # The handler thread is blocked on large_op, so no further operation will + # arrive. Stopping the batcher makes a regression return [] instead of + # blocking the test forever in _wait_for_first_operation. + state._checkpointing_stopped.set() + second_batch = state._collect_checkpoint_batch() + assert second_batch == [large_op] + assert state._overflow_queue.empty() + + +def test_drain_overflow_takes_oversized_operation_when_batch_is_empty() -> None: + """_drain_overflow adds an operation above the size limit to an empty batch, + and refuses it once the batch holds another operation. + """ + mock_lambda_client = Mock(spec=LambdaClient) + config = CheckpointBatcherConfig(max_batch_size_bytes=500) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + batcher_config=config, + ) + first_large = _oversized_queued_operation("large_1", config.max_batch_size_bytes) + second_large = _oversized_queued_operation("large_2", config.max_batch_size_bytes) + state._overflow_queue.put(first_large) + state._overflow_queue.put(second_large) + + batch = _BatchAccumulator(config) + state._drain_overflow(batch) + + # The first operation is taken despite its size. The second does not fit + # beside it and stays queued for the next batch. + assert batch.operations == [first_large] + assert state._overflow_queue.qsize() == 1 + assert state._overflow_queue.get_nowait() is second_large + + def test_collect_checkpoint_batch_handles_empty_checkpoint(): """Test batch collection with empty checkpoints (None operation_update).""" mock_lambda_client = Mock(spec=LambdaClient)