diff --git a/test/collection/test_batch.py b/test/collection/test_batch.py index 3837317ab..79053710e 100644 --- a/test/collection/test_batch.py +++ b/test/collection/test_batch.py @@ -2,6 +2,7 @@ import pytest +from weaviate.collections.batch.base import _RateLimitedBatching from weaviate.collections.batch.grpc_batch import _validate_props from weaviate.collections.classes.batch import ( MAX_STORED_RESULTS, @@ -35,6 +36,29 @@ def _error_reference(index: int) -> ErrorReference: ) +@pytest.mark.parametrize( + ("number_objects", "elapsed_time", "expected_sleep_time"), + [ + (500, 0, 31), + (100, 0, 6.2), + (100, 2, 4.2), + (100, -62, 68.2), + (100, 7, 0), + (0, 0, 0), + ], +) +def test_rate_limited_batching_sleep_time_uses_previous_batch_size( + number_objects: int, elapsed_time: float, expected_sleep_time: float +) -> None: + batching = _RateLimitedBatching(requests_per_minute=1000) + + assert batching.get_sleep_time( + number_objects=number_objects, + elapsed_time=elapsed_time, + base_time=62, + ) == pytest.approx(expected_sleep_time) + + def test_batch_object_return_add() -> None: lhs_uuids = [uuid.uuid4() for _ in range(MAX_STORED_RESULTS)] lhs = BatchObjectReturn( diff --git a/weaviate/collections/batch/base.py b/weaviate/collections/batch/base.py index af6a9ea49..e1d91763e 100644 --- a/weaviate/collections/batch/base.py +++ b/weaviate/collections/batch/base.py @@ -257,6 +257,10 @@ class _FixedSizeBatching: class _RateLimitedBatching: requests_per_minute: int + def get_sleep_time(self, number_objects: int, elapsed_time: float, base_time: float) -> float: + batch_interval = base_time * number_objects / self.requests_per_minute + return max(batch_interval - elapsed_time, 0) + @dataclass class _ServerSideBatching: @@ -346,6 +350,7 @@ def __init__( # fixed rate batching self.__time_stamp_last_request: float = 0 + self.__num_objects_in_previous_batch: int = 0 # do 62 secs to give us some buffer to the "per-minute" calculation self.__fix_rate_batching_base_time = 62 @@ -395,11 +400,13 @@ def __batch_send(self) -> None: and not self.__shut_background_thread_down.is_set() ): if isinstance(self.__batching_mode, _RateLimitedBatching): - if ( - time.time() - self.__time_stamp_last_request - < self.__fix_rate_batching_base_time // self.__concurrent_requests - ): - time.sleep(1) + sleep_time = self.__batching_mode.get_sleep_time( + number_objects=self.__num_objects_in_previous_batch, + elapsed_time=time.time() - self.__time_stamp_last_request, + base_time=self.__fix_rate_batching_base_time, + ) + if sleep_time > 0: + time.sleep(min(sleep_time, 1)) continue refresh_time = 0 elif isinstance(self.__batching_mode, _DynamicBatching) and self.__vectorizer_batching: @@ -415,7 +422,8 @@ def __batch_send(self) -> None: self.__active_requests < self.__concurrent_requests and len(self.__batch_objects) + len(self.__batch_references) > 0 ): - self.__time_stamp_last_request = time.time() + if not isinstance(self.__batching_mode, _RateLimitedBatching): + self.__time_stamp_last_request = time.time() self._batch_send = True with self.__active_requests_lock: @@ -444,6 +452,12 @@ def __batch_send(self) -> None: self.__recommended_num_refs, uuid_lookup=self.__uuid_lookup, ) + if isinstance(self.__batching_mode, _RateLimitedBatching): + # Preserve a future timestamp set by a concurrent rate-limit retry. + self.__time_stamp_last_request = max( + self.__time_stamp_last_request, time.time() + ) + self.__num_objects_in_previous_batch = len(objs) # do not block the thread - the results are written to a central (locked) list and we want to have multiple concurrent batch-requests ctx = contextvars.copy_context() self.__executor.submit(