Fix batch rate-limit scheduling for partial batches

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(

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(