From f88d7f9e7746c66989b834d0ce814bb4b3205b75 Mon Sep 17 00:00:00 2001 From: Kevin Kaspari Date: Mon, 20 Jan 2025 08:24:27 -0800 Subject: [PATCH] Fix setitem loop now that we use Dynamo retry config. Update number for retries in test assertion --- tests/serialize/runstate/dynamodb_state_store_test.py | 2 +- tron/serialize/runstate/dynamodb_state_store.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/serialize/runstate/dynamodb_state_store_test.py b/tests/serialize/runstate/dynamodb_state_store_test.py index 687b331bd..d2cbd3475 100644 --- a/tests/serialize/runstate/dynamodb_state_store_test.py +++ b/tests/serialize/runstate/dynamodb_state_store_test.py @@ -345,7 +345,7 @@ def test_retry_reading(self, store, small_object, large_object): ): store.restore(keys) except Exception: - assert_equal(mock_failed_read.call_count, 11) + assert_equal(mock_failed_read.call_count, 10) def test_restore_exception_propagation(self, store, small_object): # This test is to ensure that restore propagates exceptions upwards: see DAR-2328 diff --git a/tron/serialize/runstate/dynamodb_state_store.py b/tron/serialize/runstate/dynamodb_state_store.py index 929a7249b..6d69afe08 100644 --- a/tron/serialize/runstate/dynamodb_state_store.py +++ b/tron/serialize/runstate/dynamodb_state_store.py @@ -374,7 +374,9 @@ def __setitem__(self, key: str, value: Tuple[bytes, str]) -> None: items.append(item) - while len(items) == MAX_TRANSACT_WRITE_ITEMS or index == max_partitions - 1: + # We want to write the items when we've either reached the max number of items + # for a transaction, or when we're done processing all partitions + if len(items) == MAX_TRANSACT_WRITE_ITEMS or index == max_partitions - 1: try: self.client.transact_write_items(TransactItems=items) items = []