vectordb-bench 0.0.1__1-py3-none-any.whl

vectordb_bench/backend/runner/mp_runner.py

@@ -0,0 +1,124 @@
+import time
+import traceback
+import concurrent
+import multiprocessing as mp
+import logging
+from typing import Iterable
+import numpy as np
+from ..clients import api
+from .. import utils
+from ... import config
+
+
+NUM_PER_BATCH = config.NUM_PER_BATCH
+log = logging.getLogger(__name__)
+
+
+class MultiProcessingSearchRunner:
+    """ multiprocessing search runner
+
+    Args:
+        k(int): search topk, default to 100
+        concurrency(Iterable): concurrencies, default [1, 5, 10, 15, 20, 25, 30, 35]
+        duration(int): duration for each concurency, default to 30s
+    """
+    def __init__(
+        self,
+        db: api.VectorDB,
+        test_data: np.ndarray,
+        k: int = 100,
+        filters: dict | None = None,
+        concurrencies: Iterable[int] = (1, 5, 10, 15, 20, 25, 30, 35),
+        duration: int = 30,
+    ):
+        self.db = db
+        self.k = k
+        self.filters = filters
+        self.concurrencies = concurrencies
+        self.duration = duration
+
+        self.test_data = utils.SharedNumpyArray(test_data)
+        log.debug(f"test dataset columns: {len(test_data)}")
+
+    def search(self, test_np: utils.SharedNumpyArray) -> tuple[int, float]:
+        with self.db.init():
+            test_data = test_np.read().tolist()
+            num, idx = len(test_data), 0
+
+            start_time = time.perf_counter()
+            count = 0
+            while time.perf_counter() < start_time + self.duration:
+                s = time.perf_counter()
+                try:
+                    self.db.search_embedding(
+                        test_data[idx],
+                        self.k,
+                        self.filters,
+                    )
+                except Exception as e:
+                    log.warning(f"VectorDB search_embedding error: {e}")
+                    traceback.print_exc(chain=True)
+                    raise e from None
+
+                count += 1
+                # loop through the test data
+                idx = idx + 1 if idx < num - 1 else 0
+
+                if count % 500 == 0:
+                    log.debug(f"({mp.current_process().name:16}) search_count: {count}, latest_latency={time.perf_counter()-s}")
+
+        total_dur = round(time.perf_counter() - start_time, 4)
+        log.info(
+            f"{mp.current_process().name:16} search {self.duration}s: "
+            f"actual_dur={total_dur}s, count={count}, qps in this process: {round(count / total_dur, 4):3}"
+         )
+
+        return (count, total_dur)
+
+    @staticmethod
+    def get_mp_context():
+        mp_start_method = "forkserver" if "forkserver" in mp.get_all_start_methods() else "spawn"
+        log.debug(f"MultiProcessingSearchRunner get multiprocessing start method: {mp_start_method}")
+        return mp.get_context(mp_start_method)
+
+    def _run_all_concurrencies_mem_efficient(self) -> float:
+        max_qps = 0
+        try:
+            for conc in self.concurrencies:
+                with concurrent.futures.ProcessPoolExecutor(mp_context=self.get_mp_context(), max_workers=conc) as executor:
+                    start = time.perf_counter()
+                    log.info(f"start search {self.duration}s in concurrency {conc}, filters: {self.filters}")
+                    future_iter = executor.map(self.search, [self.test_data for i in range(conc)])
+                    all_count = sum([r[0] for r in future_iter])
+
+                    cost = time.perf_counter() - start
+                    qps = round(all_count / cost, 4)
+                    log.info(f"end search in concurrency {conc}: dur={cost}s, total_count={all_count}, qps={qps}")
+
+                if qps > max_qps:
+                    max_qps = qps
+                    log.info(f"update largest qps with concurrency {conc}: current max_qps={max_qps}")
+        except Exception as e:
+            log.warning(f"fail to search all concurrencies: {self.concurrencies}, max_qps before failure={max_qps}, reason={e}")
+            traceback.print_exc()
+
+            # No results available, raise exception
+            if max_qps == 0.0:
+                raise e from None
+
+        finally:
+            self.stop()
+
+        return max_qps
+
+    def run(self) -> float:
+        """
+        Returns:
+            float: largest qps
+        """
+        return self._run_all_concurrencies_mem_efficient()
+
+    def stop(self) -> None:
+        if self.test_data:
+            self.test_data.unlink()
+            self.test_data = None

vectordb_bench/backend/runner/serial_runner.py

@@ -0,0 +1,164 @@
+import time
+import logging
+import traceback
+import concurrent
+import multiprocessing as mp
+import math
+import numpy as np
+import pandas as pd
+
+from ..clients import api
+from ...metric import calc_recall
+from ...models import LoadTimeoutError
+from .. import utils
+from ... import config
+
+NUM_PER_BATCH = config.NUM_PER_BATCH
+LOAD_TIMEOUT = 24 * 60 * 60
+
+log = logging.getLogger(__name__)
+
+
+class SerialInsertRunner:
+    def __init__(self, db: api.VectorDB, train_emb: list[list[float]], train_id: list[int]):
+        log.debug(f"Dataset shape: {len(train_emb)}")
+        self.db = db
+        self.shared_emb = train_emb
+        self.train_id = train_id
+
+        self.seq_batches = math.ceil(len(train_emb)/NUM_PER_BATCH)
+
+    def insert_data(self, left_id: int = 0) -> int:
+        with self.db.init():
+            all_embeddings = self.shared_emb
+
+            # unique id for endlessness insertion
+            all_metadata = [i+left_id for i in self.train_id]
+
+            num_conc_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}")
+            count = 0
+            for batch_id in range(self.seq_batches):
+                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_conc_batches}], Start inserting {len(metadata)} embeddings")
+                insert_count = self.db.insert_embeddings(
+                    embeddings=embeddings,
+                    metadata=metadata,
+                )
+                log.debug(f"({mp.current_process().name:16}) batch [{batch_id:3}/{num_conc_batches}], Finish inserting {len(metadata)} embeddings")
+
+                assert insert_count == len(metadata)
+                count += insert_count
+            log.info(f"({mp.current_process().name:16}) Finish inserting {len(all_embeddings)} embeddings in 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:
+            future = executor.submit(self.insert_data)
+            count = future.result()
+            return count
+
+    def run_endlessness(self) -> int:
+        """run forever util DB raises exception or crash"""
+        start_time = time.perf_counter()
+        max_load_count, times = 0, 0
+        try:
+            with self.db.init():
+                self.db.ready_to_load()
+            while time.perf_counter() - start_time < config.CASE_TIMEOUT_IN_SECOND:
+                count = self.insert_data(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}")
+            raise LoadTimeoutError("capacity case load timeout and stop")
+        except LoadTimeoutError as e:
+            log.info("load timetout, stop the load case")
+            raise e from None
+        except Exception as e:
+            log.info(f"Capacity case load reach limit, insertion counts={utils.numerize(max_load_count)}, {max_load_count}, err={e}")
+            traceback.print_exc()
+            return max_load_count
+
+    def run(self) -> int:
+        count, dur = self._insert_all_batches()
+        return count
+
+
+class SerialSearchRunner:
+    def __init__(
+        self,
+        db: api.VectorDB,
+        test_data: list[list[float]],
+        ground_truth: pd.DataFrame,
+        k: int = 100,
+        filters: dict | None = None,
+    ):
+        self.db = db
+        self.k = k
+        self.filters = filters
+
+        if isinstance(test_data[0], np.ndarray):
+            self.test_data = [query.tolist() for query in test_data]
+        else:
+            self.test_data = test_data
+        self.ground_truth = ground_truth
+
+    def search(self, args: tuple[list, pd.DataFrame]):
+        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
+
+            log.debug(f"test dataset size: {len(test_data)}")
+            log.info(f"ground truth size: {ground_truth.columns}, shape: {ground_truth.shape}")
+
+            latencies, recalls = [], []
+            for idx, emb in enumerate(test_data):
+                s = time.perf_counter()
+                try:
+                    results = self.db.search_embedding(
+                        emb,
+                        self.k,
+                        self.filters,
+                    )
+
+                except Exception as e:
+                    log.warning(f"VectorDB search_embedding error: {e}")
+                    traceback.print_exc(chain=True)
+                    raise e from None
+
+                latencies.append(time.perf_counter() - s)
+
+                gt = ground_truth['neighbors_id'][idx]
+                recalls.append(calc_recall(self.k, gt[:self.k], results))
+
+
+                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]}")
+
+        avg_latency = round(np.mean(latencies), 4)
+        avg_recall = round(np.mean(recalls), 4)
+        cost = round(np.sum(latencies), 4)
+        p99 = round(np.percentile(latencies, 99), 4)
+        log.info(
+            f"{mp.current_process().name:14} search entire test_data: "
+            f"cost={cost}s, "
+            f"queries={len(latencies)}, "
+            f"avg_recall={avg_recall}, "
+            f"avg_latency={avg_latency}, "
+            f"p99={p99}"
+         )
+        return (avg_recall, 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
+
+    def run(self) -> tuple[float, float]:
+        return self._run_in_subprocess()

vectordb_bench/backend/task_runner.py

@@ -0,0 +1,290 @@
+import logging
+import traceback
+import concurrent
+import numpy as np
+from enum import Enum, auto
+
+from . import utils
+from .cases import Case, CaseLabel
+from ..base import BaseModel
+from ..models import TaskConfig
+
+from .clients import (
+    api,
+    ZillizCloud,
+    Milvus,
+    MetricType
+)
+from ..metric import Metric
+from .runner import MultiProcessingSearchRunner
+from .runner import SerialSearchRunner, SerialInsertRunner
+
+
+log = logging.getLogger(__name__)
+
+
+class RunningStatus(Enum):
+    PENDING = auto()
+    FINISHED = auto()
+
+
+class CaseRunner(BaseModel):
+    """ DataSet, filter_rate, db_class with db config
+
+    Fields:
+        run_id(str): run_id of this case runner,
+            indicating which task does this case belong to.
+        config(TaskConfig): task configs of this case runner.
+        ca(Case): case for this case runner.
+        status(RunningStatus): RunningStatus of this case runner.
+
+        db(api.VectorDB): The vector database for this case runner.
+    """
+
+    run_id: str
+    config: TaskConfig
+    ca: Case
+    status: RunningStatus
+
+    db: api.VectorDB | None = None
+    test_emb: np.ndarray | None = None
+    search_runner: MultiProcessingSearchRunner | None = None
+    serial_search_runner: SerialSearchRunner | None = None
+
+    def __eq__(self, obj):
+        if isinstance(obj, CaseRunner):
+            return self.ca.label == CaseLabel.Performance and \
+                self.config.db == obj.config.db and \
+                self.config.db_case_config == obj.config.db_case_config and \
+                self.ca.dataset == obj.ca.dataset
+            return False
+
+    def display(self) -> dict:
+        c_dict = self.ca.dict(include={'label':True, 'filters': True,'dataset':{'data': True} })
+        c_dict['db'] = self.config.db_name
+        return c_dict
+
+    @property
+    def normalize(self) -> bool:
+        assert self.db
+        return isinstance(self.db, (Milvus, ZillizCloud)) and \
+            self.ca.dataset.data.metric_type == MetricType.COSINE
+
+    def init_db(self, drop_old: bool = True) -> None:
+        db_cls = self.config.db.init_cls
+
+        self.db = db_cls(
+            dim=self.ca.dataset.data.dim,
+            db_config=self.config.db_config.to_dict(),
+            db_case_config=self.config.db_case_config,
+            drop_old=drop_old,
+        )
+
+    def _pre_run(self, drop_old: bool = True):
+        try:
+            self.ca.dataset.prepare()
+            self.init_db(drop_old)
+        except Exception as e:
+            log.warning(f"pre run case error: {e}")
+            raise e from None
+
+    def run(self, drop_old: bool = True) -> Metric:
+        self._pre_run(drop_old)
+
+        if self.ca.label == CaseLabel.Load:
+            return self._run_load_case()
+        elif self.ca.label == CaseLabel.Performance:
+            return self._run_perf_case(drop_old)
+        else:
+            log.warning(f"unknown case type: {self.ca.label}")
+            raise ValueError(f"Unknown case type: {self.ca.label}")
+
+
+    def _run_load_case(self) -> Metric:
+        """ run load cases
+
+        Returns:
+            Metric: the max load count
+        """
+        log.info("Start capacity case")
+        # datasets for load tests are quite small, can fit into memory
+        # only 1 file
+        data_df = [data_df for data_df in self.ca.dataset][0]
+
+        all_embeddings, all_metadata = np.stack(data_df["emb"]).tolist(), data_df['id'].tolist()
+        runner = SerialInsertRunner(self.db, all_embeddings, all_metadata)
+        try:
+            count = runner.run_endlessness()
+            log.info(f"load reach limit: insertion counts={count}")
+            return Metric(max_load_count=count)
+        except Exception as e:
+            log.warning(f"run capacity case error: {e}")
+            raise e from None
+        log.info("End capacity case")
+
+
+    def _run_perf_case(self, drop_old: bool = True) -> Metric:
+        try:
+            m = Metric()
+            if drop_old:
+                _, load_dur = self._load_train_data()
+                build_dur = self._optimize()
+                m.load_duration = round(load_dur+build_dur, 4)
+
+            self._init_search_runner()
+            m.recall, m.serial_latency_p99 = self._serial_search()
+            m.qps = self._conc_search()
+
+            log.info(f"got results: {m}")
+            return m
+        except Exception as e:
+            log.warning(f"performance case run error: {e}")
+            traceback.print_exc()
+            raise e
+
+    @utils.time_it
+    def _load_train_data(self):
+        """Insert train data and get the insert_duration"""
+        for data_df in self.ca.dataset:
+            try:
+                all_metadata = data_df['id'].tolist()
+
+                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)
+                log.debug(f"normalized size: {len(all_embeddings)}, {len(all_metadata)}")
+
+                runner = SerialInsertRunner(self.db, all_embeddings, all_metadata)
+                runner.run()
+            except Exception as e:
+                raise e from None
+            finally:
+                runner = None
+
+
+    def _serial_search(self) -> tuple[float, float]:
+        """Performance serial tests, search the entire test data once,
+        calculate the recall, serial_latency_p99
+
+        Returns:
+            tuple[float, float]: recall, serial_latency_p99
+        """
+        try:
+            return self.serial_search_runner.run()
+        except Exception as e:
+            log.warning(f"search error: {str(e)}, {e}")
+            self.stop()
+            raise e from None
+
+    def _conc_search(self):
+        """Performance concurrency tests, search the test data endlessness
+        for 30s in several concurrencies
+
+        Returns:
+            float: the largest qps in all concurrencies
+        """
+        try:
+            return self.search_runner.run()
+        except Exception as e:
+            log.warning(f"search error: {str(e)}, {e}")
+            raise e from None
+        finally:
+            self.stop()
+
+    @utils.time_it
+    def _task(self) -> None:
+        """"""
+        with self.db.init():
+            self.db.ready_to_search()
+
+    def _optimize(self) -> float:
+        with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
+            future = executor.submit(self._task)
+            try:
+                return future.result()[1]
+            except Exception as e:
+                log.warning(f"VectorDB ready_to_search error: {e}")
+                raise e from None
+
+    def _init_search_runner(self):
+        test_emb = np.stack(self.ca.dataset.test_data["emb"])
+        if self.normalize:
+            test_emb = test_emb / np.linalg.norm(test_emb, axis=1)[:, np.newaxis]
+        self.test_emb = test_emb
+
+        gt_df = self.ca.dataset.get_ground_truth(self.ca.filter_rate)
+
+        self.serial_search_runner = SerialSearchRunner(
+            db=self.db,
+            test_data=self.test_emb.tolist(),
+            ground_truth=gt_df,
+            filters=self.ca.filters,
+        )
+
+        self.search_runner =  MultiProcessingSearchRunner(
+            db=self.db,
+            test_data=self.test_emb,
+            filters=self.ca.filters,
+        )
+
+    def stop(self):
+        if self.search_runner:
+            self.search_runner.stop()
+
+
+class TaskRunner(BaseModel):
+    run_id: str
+    task_label: str
+    case_runners: list[CaseRunner]
+
+    def num_cases(self) -> int:
+        return len(self.case_runners)
+
+    def num_finished(self) -> int:
+        return self._get_num_by_status(RunningStatus.FINISHED)
+
+    def set_finished(self, idx: int) -> None:
+        self.case_runners[idx].status = RunningStatus.FINISHED
+
+    def _get_num_by_status(self, status: RunningStatus) -> int:
+        return sum([1 for c in self.case_runners if c.status == status])
+
+    def display(self) -> None:
+        DATA_FORMAT = (" %-14s | %-12s %-20s %7s | %-10s")
+        TITLE_FORMAT = (" %-14s | %-12s %-20s %7s | %-10s") % (
+            "DB", "CaseType", "Dataset", "Filter", "task_label")
+
+        fmt = [TITLE_FORMAT]
+        fmt.append(DATA_FORMAT%(
+            "-"*11,
+            "-"*12,
+            "-"*20,
+            "-"*7,
+            "-"*7
+        ))
+
+        for f in self.case_runners:
+            if f.ca.filter_rate != 0.0:
+                filters = f.ca.filter_rate
+            elif f.ca.filter_size != 0:
+                filters = f.ca.filter_size
+            else:
+                filters = "None"
+
+            ds_str = f"{f.ca.dataset.data.name}-{f.ca.dataset.data.label}-{utils.numerize(f.ca.dataset.data.size)}"
+            fmt.append(DATA_FORMAT%(
+                f.config.db_name,
+                f.ca.label.name,
+                ds_str,
+                filters,
+                self.task_label,
+            ))
+
+        tmp_logger = logging.getLogger("no_color")
+        for f in fmt:
+            tmp_logger.info(f)

vectordb_bench/backend/utils.py

@@ -0,0 +1,85 @@
+import time
+from functools import wraps
+from multiprocessing.shared_memory import SharedMemory
+
+import numpy as np
+
+
+def numerize(n) -> str:
+    """display positive number n for readability
+
+    Examples:
+        >>> numerize(1_000)
+        '1K'
+        >>> numerize(1_000_000_000)
+        '1B'
+    """
+    sufix2upbound = {
+        "EMPTY": 1e3,
+        "K": 1e6,
+        "M": 1e9,
+        "B": 1e12,
+        "END": float('inf'),
+    }
+
+    display_n, sufix = n, ""
+    for s, base in sufix2upbound.items():
+        # number >= 1000B will alway have sufix 'B'
+        if s == "END":
+            display_n = int(n/1e9)
+            sufix = "B"
+            break
+
+        if n < base:
+            sufix = "" if s == "EMPTY" else s
+            display_n = int(n/(base/1e3))
+            break
+    return f"{display_n}{sufix}"
+
+
+def time_it(func):
+    @wraps(func)
+    def inner(*args, **kwargs):
+        pref = time.perf_counter()
+        result = func(*args, **kwargs)
+        delta = time.perf_counter() - pref
+        return result, delta
+    return inner
+
+
+class SharedNumpyArray:
+    ''' Wraps a numpy array so that it can be shared quickly among processes,
+    avoiding unnecessary copying and (de)serializing.
+    '''
+    def __init__(self, array: np.ndarray):
+        '''
+        Creates the shared memory and copies the array therein
+        '''
+        # create the shared memory location of the same size of the array
+        self._shared = SharedMemory(create=True, size=array.nbytes)
+
+        # save data type and shape, necessary to read the data correctly
+        self._dtype, self._shape = array.dtype, array.shape
+
+        # create a new numpy array that uses the shared memory we created.
+        # at first, it is filled with zeros
+        res = np.ndarray(
+            self._shape, dtype=self._dtype, buffer=self._shared.buf
+        )
+
+        # copy data from the array to the shared memory. numpy will
+        # take care of copying everything in the correct format
+        res[:] = array[:]
+
+    def read(self) -> np.ndarray:
+        '''Reads the array from the shared memory without unnecessary copying. '''
+        # simply create an array of the correct shape and type,
+        # using the shared memory location we created earlier
+        return np.ndarray(self._shape, self._dtype, buffer=self._shared.buf)
+
+    def unlink(self) -> None:
+        ''' Releases the allocated memory. Call when finished using the data,
+        or when the data was copied somewhere else.
+        '''
+        self._shared.close()
+        self._shared.unlink()

vectordb_bench/base.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel as PydanticBaseModel
+
+
+class BaseModel(PydanticBaseModel, arbitrary_types_allowed=True):
+    pass
+