update the retry behavior for failed insert request

Signed-off-by: min.tian <min.tian.cn@gmail.com>

Author: alwayslove2013

vectordb_bench/backend/clients/chroma/chroma.py

@@ -65,7 +65,7 @@ def insert_embeddings(
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs: Any,
-    ) -> (int, Exception):
+    ) -> tuple[int, Exception]:
         """Insert embeddings into the database.
 
         Args:
@@ -78,8 +78,12 @@ def insert_embeddings(
         """
         ids = [str(i) for i in metadata]
         metadata = [{"id": int(i)} for i in metadata]
-        if len(embeddings) > 0:
-            self.collection.add(embeddings=embeddings, ids=ids, metadatas=metadata)
+        try:
+            if len(embeddings) > 0:
+                self.collection.add(embeddings=embeddings, ids=ids, metadatas=metadata)
+        except Exception as e:
+            log.warning(f"Failed to insert data: error: {e!s}")
+            return 0, e
         return len(embeddings), None
 
     def search_embedding(

vectordb_bench/backend/runner/mp_runner.py

@@ -258,7 +258,11 @@ def search_by_dur(self, dur: int, test_data: list[list[float]], q: mp.Queue, con
                     success_count += 1
                 except Exception as e:
                     failed_cnt += 1
-                    log.warning(f"VectorDB search_embedding error: {e}")
+                    # reduce log
+                    if failed_cnt <= 3:
+                        log.warning(f"VectorDB search_embedding error: {e}")
+                    else:
+                        log.debug(f"VectorDB search_embedding error: {e}")
 
                 # loop through the test data
                 idx = idx + 1 if idx < num - 1 else 0

vectordb_bench/backend/runner/rate_runner.py

@@ -34,17 +34,16 @@ def __init__(
         self.sig_idx = 0
 
     def send_insert_task(self, db: api.VectorDB, emb: list[list[float]], metadata: list[str], retry_idx: int = 0):
-        try:
-            db.insert_embeddings(emb, metadata)
-        except Exception as e:
-            log.warning(f"Insert Failed, try_idx={retry_idx}, Exception: {e}")
+        _, error = db.insert_embeddings(emb, metadata)
+        if error is not None:
+            log.warning(f"Insert Failed, try_idx={retry_idx}, Exception: {error}")
             retry_idx += 1
             if retry_idx <= config.MAX_INSERT_RETRY:
                 time.sleep(retry_idx)
                 self.send_insert_task(db, emb=emb, metadata=metadata, retry_idx=retry_idx)
             else:
                 msg = f"Insert failed and retried more than {config.MAX_INSERT_RETRY} times"
-                raise RuntimeError(msg) from e
+                raise RuntimeError(msg) from None
 
     @time_it
     def run_with_rate(self, q: mp.Queue):

vectordb_bench/backend/runner/serial_runner.py

@@ -38,6 +38,18 @@ def __init__(
         self.normalize = normalize
         self.filters = filters
 
+    def retry_insert(self, db: api.VectorDB, retry_idx: int = 0, **kwargs):
+        _, error = db.insert_embeddings(**kwargs)
+        if error is not None:
+            log.warning(f"Insert Failed, try_idx={retry_idx}, Exception: {error}")
+            retry_idx += 1
+            if retry_idx <= config.MAX_INSERT_RETRY:
+                time.sleep(retry_idx)
+                self.retry_insert(db, retry_idx=retry_idx, **kwargs)
+            else:
+                msg = f"Insert failed and retried more than {config.MAX_INSERT_RETRY} times"
+                raise RuntimeError(msg) from None
+
     def task(self) -> int:
         count = 0
         with self.db.init():
@@ -68,7 +80,12 @@ def task(self) -> int:
                     labels_data=labels_data,
                 )
                 if error is not None:
-                    raise error
+                    self.retry_insert(
+                        self.db,
+                        embeddings=all_embeddings,
+                        metadata=all_metadata,
+                        labels_data=labels_data,
+                    )
 
                 assert insert_count == len(all_metadata)
                 count += insert_count