vectordb-bench 0.0.19__py3-none-any.whl → 0.0.20__py3-none-any.whl

vectordb_bench/backend/runner/rate_runner.py

@@ -1,36 +1,36 @@
+import concurrent
 import logging
+import multiprocessing as mp
 import time
-import concurrent
 from concurrent.futures import ThreadPoolExecutor
-import multiprocessing as mp
-
 
+from vectordb_bench import config
 from vectordb_bench.backend.clients import api
 from vectordb_bench.backend.dataset import DataSetIterator
 from vectordb_bench.backend.utils import time_it
-from vectordb_bench import config
 
 from .util import get_data
+
 log = logging.getLogger(__name__)
 
 
 class RatedMultiThreadingInsertRunner:
     def __init__(
         self,
-        rate: int, # numRows per second
+        rate: int,  # numRows per second
         db: api.VectorDB,
         dataset_iter: DataSetIterator,
         normalize: bool = False,
         timeout: float | None = None,
     ):
-        self.timeout = timeout if isinstance(timeout, (int, float)) else None
+        self.timeout = timeout if isinstance(timeout, int | float) else None
         self.dataset = dataset_iter
         self.db = db
         self.normalize = normalize
         self.insert_rate = rate
         self.batch_rate = rate // config.NUM_PER_BATCH
 
-    def send_insert_task(self, db, emb: list[list[float]], metadata: list[str]):
+    def send_insert_task(self, db: api.VectorDB, emb: list[list[float]], metadata: list[str]):
         db.insert_embeddings(emb, metadata)
 
     @time_it
@@ -43,7 +43,9 @@ class RatedMultiThreadingInsertRunner:
                 rate = self.batch_rate
                 for data in self.dataset:
                     emb, metadata = get_data(data, self.normalize)
-                    executing_futures.append(executor.submit(self.send_insert_task, self.db, emb, metadata))
+                    executing_futures.append(
+                        executor.submit(self.send_insert_task, self.db, emb, metadata),
+                    )
                     rate -= 1
 
                     if rate == 0:
@@ -66,19 +68,26 @@ class RatedMultiThreadingInsertRunner:
                         done, not_done = concurrent.futures.wait(
                             executing_futures,
                             timeout=wait_interval,
-                            return_when=concurrent.futures.FIRST_EXCEPTION)
+                            return_when=concurrent.futures.FIRST_EXCEPTION,
+                        )
 
                         if len(not_done) > 0:
-                            log.warning(f"Failed to finish all tasks in 1s, [{len(not_done)}/{len(executing_futures)}] tasks are not done, waited={wait_interval:.2f}, trying to wait in the next round")
+                            log.warning(
+                                f"Failed to finish all tasks in 1s, [{len(not_done)}/{len(executing_futures)}] ",
+                                f"tasks are not done, waited={wait_interval:.2f}, trying to wait in the next round",
+                            )
                             executing_futures = list(not_done)
                         else:
-                            log.debug(f"Finished {len(executing_futures)} insert-{config.NUM_PER_BATCH} task in 1s, wait_interval={wait_interval:.2f}")
+                            log.debug(
+                                f"Finished {len(executing_futures)} insert-{config.NUM_PER_BATCH} ",
+                                f"task in 1s, wait_interval={wait_interval:.2f}",
+                            )
                             executing_futures = []
                     except Exception as e:
-                            log.warn(f"task error, terminating, err={e}")
-                            q.put(None, block=True)
-                            executor.shutdown(wait=True, cancel_futures=True)
-                            raise e
+                        log.warning(f"task error, terminating, err={e}")
+                        q.put(None, block=True)
+                        executor.shutdown(wait=True, cancel_futures=True)
+                        raise e from e
 
                     dur = time.perf_counter() - start_time
                     if dur < 1:
@@ -87,10 +96,12 @@ class RatedMultiThreadingInsertRunner:
                 # wait for all tasks in executing_futures to complete
                 if len(executing_futures) > 0:
                     try:
-                        done, _ = concurrent.futures.wait(executing_futures,
-                           return_when=concurrent.futures.FIRST_EXCEPTION)
+                        done, _ = concurrent.futures.wait(
+                            executing_futures,
+                            return_when=concurrent.futures.FIRST_EXCEPTION,
+                        )
                     except Exception as e:
-                        log.warn(f"task error, terminating, err={e}")
+                        log.warning(f"task error, terminating, err={e}")
                         q.put(None, block=True)
                         executor.shutdown(wait=True, cancel_futures=True)
-                        raise e
+                        raise e from e

vectordb_bench/backend/runner/read_write_runner.py

@@ -1,16 +1,18 @@
+import concurrent
 import logging
-from typing import Iterable
+import math
 import multiprocessing as mp
-import concurrent
+from collections.abc import Iterable
+
 import numpy as np
-import math
 
-from .mp_runner import MultiProcessingSearchRunner
-from .serial_runner import SerialSearchRunner
-from .rate_runner import RatedMultiThreadingInsertRunner
 from vectordb_bench.backend.clients import api
 from vectordb_bench.backend.dataset import DatasetManager
 
+from .mp_runner import MultiProcessingSearchRunner
+from .rate_runner import RatedMultiThreadingInsertRunner
+from .serial_runner import SerialSearchRunner
+
 log = logging.getLogger(__name__)
 
 
@@ -24,8 +26,14 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
         k: int = 100,
         filters: dict | None = None,
         concurrencies: Iterable[int] = (1, 15, 50),
-        search_stage: Iterable[float] = (0.5, 0.6, 0.7, 0.8, 0.9), # search from insert portion, 0.0 means search from the start
-        read_dur_after_write: int = 300, # seconds, search duration when insertion is done
+        search_stage: Iterable[float] = (
+            0.5,
+            0.6,
+            0.7,
+            0.8,
+            0.9,
+        ),  # search from insert portion, 0.0 means search from the start
+        read_dur_after_write: int = 300,  # seconds, search duration when insertion is done
         timeout: float | None = None,
     ):
         self.insert_rate = insert_rate
@@ -36,7 +44,10 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
         self.search_stage = sorted(search_stage)
         self.read_dur_after_write = read_dur_after_write
 
-        log.info(f"Init runner, concurencys={concurrencies}, search_stage={search_stage}, stage_search_dur={read_dur_after_write}")
+        log.info(
+            f"Init runner, concurencys={concurrencies}, search_stage={search_stage}, ",
+            f"stage_search_dur={read_dur_after_write}",
+        )
 
         test_emb = np.stack(dataset.test_data["emb"])
         if normalize:
@@ -76,8 +87,13 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
         log.info("Search after write - Serial search start")
         res, ssearch_dur = self.serial_search_runner.run()
         recall, ndcg, p99_latency = res
-        log.info(f"Search after write - Serial search - recall={recall}, ndcg={ndcg}, p99={p99_latency}, dur={ssearch_dur:.4f}")
-        log.info(f"Search after wirte - Conc search start, dur for each conc={self.read_dur_after_write}")
+        log.info(
+            f"Search after write - Serial search - recall={recall}, ndcg={ndcg}, p99={p99_latency}, ",
+            f"dur={ssearch_dur:.4f}",
+        )
+        log.info(
+            f"Search after wirte - Conc search start, dur for each conc={self.read_dur_after_write}",
+        )
         max_qps = self.run_by_dur(self.read_dur_after_write)
         log.info(f"Search after wirte - Conc search finished, max_qps={max_qps}")
 
@@ -86,7 +102,10 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
     def run_read_write(self):
         with mp.Manager() as m:
             q = m.Queue()
-            with concurrent.futures.ProcessPoolExecutor(mp_context=mp.get_context("spawn"), max_workers=2) as executor:
+            with concurrent.futures.ProcessPoolExecutor(
+                mp_context=mp.get_context("spawn"),
+                max_workers=2,
+            ) as executor:
                 read_write_futures = []
                 read_write_futures.append(executor.submit(self.run_with_rate, q))
                 read_write_futures.append(executor.submit(self.run_search_by_sig, q))
@@ -107,10 +126,10 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
                 except Exception as e:
                     log.warning(f"Read and write error: {e}")
                     executor.shutdown(wait=True, cancel_futures=True)
-                    raise e
+                    raise e from e
         log.info("Concurrent read write all done")
 
-    def run_search_by_sig(self, q):
+    def run_search_by_sig(self, q: mp.Queue):
         """
         Args:
             q: multiprocessing queue
@@ -122,15 +141,14 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
         total_batch = math.ceil(self.data_volume / self.insert_rate)
         recall, ndcg, p99_latency = None, None, None
 
-        def wait_next_target(start, target_batch) -> bool:
+        def wait_next_target(start: int, target_batch: int) -> bool:
             """Return False when receive True or None"""
             while start < target_batch:
                 sig = q.get(block=True)
 
                 if sig is None or sig is True:
                     return False
-                else:
-                    start += 1
+                start += 1
             return True
 
         for idx, stage in enumerate(self.search_stage):
@@ -139,19 +157,24 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
 
             got = wait_next_target(start_batch, target_batch)
             if got is False:
-                log.warning(f"Abnormal exit, target_batch={target_batch}, start_batch={start_batch}")
-                return
+                log.warning(
+                    f"Abnormal exit, target_batch={target_batch}, start_batch={start_batch}",
+                )
+                return None
 
             log.info(f"Insert {perc}% done, total batch={total_batch}")
             log.info(f"[{target_batch}/{total_batch}] Serial search - {perc}% start")
             res, ssearch_dur = self.serial_search_runner.run()
             recall, ndcg, p99_latency = res
-            log.info(f"[{target_batch}/{total_batch}] Serial search - {perc}% done, recall={recall}, ndcg={ndcg}, p99={p99_latency}, dur={ssearch_dur:.4f}")
+            log.info(
+                f"[{target_batch}/{total_batch}] Serial search - {perc}% done, recall={recall}, ",
+                f"ndcg={ndcg}, p99={p99_latency}, dur={ssearch_dur:.4f}",
+            )
 
             # Search duration for non-last search stage is carefully calculated.
             # If duration for each concurrency is less than 30s, runner will raise error.
             if idx < len(self.search_stage) - 1:
-                total_dur_between_stages = self.data_volume  * (self.search_stage[idx + 1] - stage) // self.insert_rate
+                total_dur_between_stages = self.data_volume * (self.search_stage[idx + 1] - stage) // self.insert_rate
                 csearch_dur = total_dur_between_stages - ssearch_dur
 
                 # Try to leave room for init process executors
@@ -159,14 +182,19 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
 
                 each_conc_search_dur = csearch_dur / len(self.concurrencies)
                 if each_conc_search_dur < 30:
-                    warning_msg = f"Results might be inaccurate, duration[{csearch_dur:.4f}] left for conc-search is too short, total available dur={total_dur_between_stages}, serial_search_cost={ssearch_dur}."
+                    warning_msg = (
+                        f"Results might be inaccurate, duration[{csearch_dur:.4f}] left for conc-search is too short, ",
+                        f"total available dur={total_dur_between_stages}, serial_search_cost={ssearch_dur}.",
+                    )
                     log.warning(warning_msg)
 
             # The last stage
             else:
                 each_conc_search_dur = 60
 
-            log.info(f"[{target_batch}/{total_batch}] Concurrent search - {perc}% start, dur={each_conc_search_dur:.4f}")
+            log.info(
+                f"[{target_batch}/{total_batch}] Concurrent search - {perc}% start, dur={each_conc_search_dur:.4f}",
+            )
             max_qps = self.run_by_dur(each_conc_search_dur)
             result.append((perc, max_qps, recall, ndcg, p99_latency))
 

vectordb_bench/backend/runner/serial_runner.py

@@ -1,20 +1,21 @@
-import time
-import logging
-import traceback
 import concurrent
-import multiprocessing as mp
+import logging
 import math
-import psutil
+import multiprocessing as mp
+import time
+import traceback
 
 import numpy as np
 import pandas as pd
+import psutil
 
-from ..clients import api
+from vectordb_bench.backend.dataset import DatasetManager
+
+from ... import config
 from ...metric import calc_ndcg, calc_recall, get_ideal_dcg
 from ...models import LoadTimeoutError, PerformanceTimeoutError
 from .. import utils
-from ... import config
-from vectordb_bench.backend.dataset import DatasetManager
+from ..clients import api
 
 NUM_PER_BATCH = config.NUM_PER_BATCH
 LOAD_MAX_TRY_COUNT = 10
@@ -22,9 +23,16 @@ WAITTING_TIME = 60
 
 log = logging.getLogger(__name__)
 
+
 class SerialInsertRunner:
-    def __init__(self, db: api.VectorDB, dataset: DatasetManager, normalize: bool, timeout: float | None = None):
-        self.timeout = timeout if isinstance(timeout, (int, float)) else None
+    def __init__(
+        self,
+        db: api.VectorDB,
+        dataset: DatasetManager,
+        normalize: bool,
+        timeout: float | None = None,
+    ):
+        self.timeout = timeout if isinstance(timeout, int | float) else None
         self.dataset = dataset
         self.db = db
         self.normalize = normalize
@@ -32,18 +40,20 @@ class SerialInsertRunner:
     def task(self) -> int:
         count = 0
         with self.db.init():
-            log.info(f"({mp.current_process().name:16}) Start inserting embeddings in batch {config.NUM_PER_BATCH}")
+            log.info(
+                f"({mp.current_process().name:16}) Start inserting embeddings in batch {config.NUM_PER_BATCH}",
+            )
             start = time.perf_counter()
             for data_df in self.dataset:
-                all_metadata = data_df['id'].tolist()
+                all_metadata = data_df["id"].tolist()
 
-                emb_np = np.stack(data_df['emb'])
+                emb_np = np.stack(data_df["emb"])
                 if self.normalize:
                     log.debug("normalize the 100k train data")
                     all_embeddings = (emb_np / np.linalg.norm(emb_np, axis=1)[:, np.newaxis]).tolist()
                 else:
                     all_embeddings = emb_np.tolist()
-                del(emb_np)
+                del emb_np
                 log.debug(f"batch dataset size: {len(all_embeddings)}, {len(all_metadata)}")
 
                 insert_count, error = self.db.insert_embeddings(
@@ -56,30 +66,41 @@ class SerialInsertRunner:
                 assert insert_count == len(all_metadata)
                 count += insert_count
                 if count % 100_000 == 0:
-                    log.info(f"({mp.current_process().name:16}) Loaded {count} embeddings into VectorDB")
+                    log.info(
+                        f"({mp.current_process().name:16}) Loaded {count} embeddings into VectorDB",
+                    )
 
-            log.info(f"({mp.current_process().name:16}) Finish loading all dataset into VectorDB, dur={time.perf_counter()-start}")
+            log.info(
+                f"({mp.current_process().name:16}) Finish loading all dataset into VectorDB, ",
+                f"dur={time.perf_counter()-start}",
+            )
             return count
 
-    def endless_insert_data(self, all_embeddings, all_metadata, left_id: int = 0) -> int:
+    def endless_insert_data(self, all_embeddings: list, all_metadata: list, left_id: int = 0) -> int:
         with self.db.init():
             # unique id for endlessness insertion
-            all_metadata = [i+left_id for i in all_metadata]
+            all_metadata = [i + left_id for i in all_metadata]
 
-            NUM_BATCHES = math.ceil(len(all_embeddings)/NUM_PER_BATCH)
-            log.info(f"({mp.current_process().name:16}) Start inserting {len(all_embeddings)} embeddings in batch {NUM_PER_BATCH}")
+            num_batches = math.ceil(len(all_embeddings) / NUM_PER_BATCH)
+            log.info(
+                f"({mp.current_process().name:16}) Start inserting {len(all_embeddings)} ",
+                f"embeddings in batch {NUM_PER_BATCH}",
+            )
             count = 0
-            for batch_id in range(NUM_BATCHES):
+            for batch_id in range(num_batches):
                 retry_count = 0
                 already_insert_count = 0
-                metadata = all_metadata[batch_id*NUM_PER_BATCH : (batch_id+1)*NUM_PER_BATCH]
-                embeddings = all_embeddings[batch_id*NUM_PER_BATCH : (batch_id+1)*NUM_PER_BATCH]
+                metadata = all_metadata[batch_id * NUM_PER_BATCH : (batch_id + 1) * NUM_PER_BATCH]
+                embeddings = all_embeddings[batch_id * NUM_PER_BATCH : (batch_id + 1) * NUM_PER_BATCH]
 
-                log.debug(f"({mp.current_process().name:16}) batch [{batch_id:3}/{NUM_BATCHES}], Start inserting {len(metadata)} embeddings")
+                log.debug(
+                    f"({mp.current_process().name:16}) batch [{batch_id:3}/{num_batches}], ",
+                    f"Start inserting {len(metadata)} embeddings",
+                )
                 while retry_count < LOAD_MAX_TRY_COUNT:
                     insert_count, error = self.db.insert_embeddings(
-                        embeddings=embeddings[already_insert_count :],
-                        metadata=metadata[already_insert_count :],
+                        embeddings=embeddings[already_insert_count:],
+                        metadata=metadata[already_insert_count:],
                     )
                     already_insert_count += insert_count
                     if error is not None:
@@ -91,17 +112,26 @@ class SerialInsertRunner:
                             raise error
                     else:
                         break
-                log.debug(f"({mp.current_process().name:16}) batch [{batch_id:3}/{NUM_BATCHES}], Finish inserting {len(metadata)} embeddings")
+                log.debug(
+                    f"({mp.current_process().name:16}) batch [{batch_id:3}/{num_batches}], ",
+                    f"Finish inserting {len(metadata)} embeddings",
+                )
 
                 assert already_insert_count == len(metadata)
                 count += already_insert_count
-            log.info(f"({mp.current_process().name:16}) Finish inserting {len(all_embeddings)} embeddings in batch {NUM_PER_BATCH}")
+            log.info(
+                f"({mp.current_process().name:16}) Finish inserting {len(all_embeddings)} embeddings in ",
+                f"batch {NUM_PER_BATCH}",
+            )
         return count
 
     @utils.time_it
     def _insert_all_batches(self) -> int:
         """Performance case only"""
-        with concurrent.futures.ProcessPoolExecutor(mp_context=mp.get_context('spawn'), max_workers=1) as executor:
+        with concurrent.futures.ProcessPoolExecutor(
+            mp_context=mp.get_context("spawn"),
+            max_workers=1,
+        ) as executor:
             future = executor.submit(self.task)
             try:
                 count = future.result(timeout=self.timeout)
@@ -121,8 +151,11 @@ class SerialInsertRunner:
         """run forever util DB raises exception or crash"""
         # datasets for load tests are quite small, can fit into memory
         # only 1 file
-        data_df = [data_df for data_df in self.dataset][0]
-        all_embeddings, all_metadata = np.stack(data_df["emb"]).tolist(), data_df['id'].tolist()
+        data_df = next(iter(self.dataset))
+        all_embeddings, all_metadata = (
+            np.stack(data_df["emb"]).tolist(),
+            data_df["id"].tolist(),
+        )
 
         start_time = time.perf_counter()
         max_load_count, times = 0, 0
@@ -130,18 +163,26 @@ class SerialInsertRunner:
             with self.db.init():
                 self.db.ready_to_load()
             while time.perf_counter() - start_time < self.timeout:
-                count = self.endless_insert_data(all_embeddings, all_metadata, left_id=max_load_count)
+                count = self.endless_insert_data(
+                    all_embeddings,
+                    all_metadata,
+                    left_id=max_load_count,
+                )
                 max_load_count += count
                 times += 1
-                log.info(f"Loaded {times} entire dataset, current max load counts={utils.numerize(max_load_count)}, {max_load_count}")
+                log.info(
+                    f"Loaded {times} entire dataset, current max load counts={utils.numerize(max_load_count)}, ",
+                    f"{max_load_count}",
+                )
         except Exception as e:
-            log.info(f"Capacity case load reach limit, insertion counts={utils.numerize(max_load_count)}, {max_load_count}, err={e}")
+            log.info(
+                f"Capacity case load reach limit, insertion counts={utils.numerize(max_load_count)}, ",
+                f"{max_load_count}, err={e}",
+            )
             traceback.print_exc()
             return max_load_count
         else:
-            msg = f"capacity case load timeout in {self.timeout}s"
-            log.info(msg)
-            raise LoadTimeoutError(msg)
+            raise LoadTimeoutError(self.timeout)
 
     def run(self) -> int:
         count, dur = self._insert_all_batches()
@@ -168,7 +209,9 @@ class SerialSearchRunner:
         self.ground_truth = ground_truth
 
     def search(self, args: tuple[list, pd.DataFrame]) -> tuple[float, float, float]:
-        log.info(f"{mp.current_process().name:14} start search the entire test_data to get recall and latency")
+        log.info(
+            f"{mp.current_process().name:14} start search the entire test_data to get recall and latency",
+        )
         with self.db.init():
             test_data, ground_truth = args
             ideal_dcg = get_ideal_dcg(self.k)
@@ -193,13 +236,15 @@ class SerialSearchRunner:
 
                 latencies.append(time.perf_counter() - s)
 
-                gt = ground_truth['neighbors_id'][idx]
-                recalls.append(calc_recall(self.k, gt[:self.k], results))
-                ndcgs.append(calc_ndcg(gt[:self.k], results, ideal_dcg))
-
+                gt = ground_truth["neighbors_id"][idx]
+                recalls.append(calc_recall(self.k, gt[: self.k], results))
+                ndcgs.append(calc_ndcg(gt[: self.k], results, ideal_dcg))
 
                 if len(latencies) % 100 == 0:
-                    log.debug(f"({mp.current_process().name:14}) search_count={len(latencies):3}, latest_latency={latencies[-1]}, latest recall={recalls[-1]}")
+                    log.debug(
+                        f"({mp.current_process().name:14}) search_count={len(latencies):3}, ",
+                        f"latest_latency={latencies[-1]}, latest recall={recalls[-1]}",
+                    )
 
         avg_latency = round(np.mean(latencies), 4)
         avg_recall = round(np.mean(recalls), 4)
@@ -213,16 +258,14 @@ class SerialSearchRunner:
             f"avg_recall={avg_recall}, "
             f"avg_ndcg={avg_ndcg},"
             f"avg_latency={avg_latency}, "
-            f"p99={p99}"
-         )
+            f"p99={p99}",
+        )
         return (avg_recall, avg_ndcg, p99)
 
-
     def _run_in_subprocess(self) -> tuple[float, float]:
         with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
             future = executor.submit(self.search, (self.test_data, self.ground_truth))
-            result = future.result()
-            return result
+            return future.result()
 
     @utils.time_it
     def run(self) -> tuple[float, float, float]:

vectordb_bench/backend/runner/util.py

@@ -1,13 +1,14 @@
 import logging
 
-from pandas import DataFrame
 import numpy as np
+from pandas import DataFrame
 
 log = logging.getLogger(__name__)
 
+
 def get_data(data_df: DataFrame, normalize: bool) -> tuple[list[list[float]], list[str]]:
-    all_metadata = data_df['id'].tolist()
-    emb_np = np.stack(data_df['emb'])
+    all_metadata = data_df["id"].tolist()
+    emb_np = np.stack(data_df["emb"])
     if normalize:
         log.debug("normalize the 100k train data")
         all_embeddings = (emb_np / np.linalg.norm(emb_np, axis=1)[:, np.newaxis]).tolist()