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

vectordb_bench/interface.py

@@ -0,0 +1,239 @@
+import traceback
+import pathlib
+import signal
+import logging
+import uuid
+import concurrent
+import multiprocessing as mp
+from multiprocessing.connection import Connection
+
+import psutil
+from enum import Enum
+
+from . import config
+from .metric import Metric
+from .models import (
+    TaskConfig,
+    TestResult,
+    CaseResult,
+    LoadTimeoutError,
+    PerformanceTimeoutError,
+    ResultLabel,
+)
+from .backend.result_collector import ResultCollector
+from .backend.assembler import Assembler
+from .backend.task_runner import TaskRunner
+
+log = logging.getLogger(__name__)
+
+global_result_future: concurrent.futures.Future | None = None
+
+class SIGNAL(Enum):
+    SUCCESS=0
+    ERROR=1
+    WIP=2
+
+
+class BenchMarkRunner:
+    def __init__(self):
+        self.running_task: TaskRunner | None = None
+        self.latest_error: str | None = None
+
+    def run(self, tasks: list[TaskConfig], task_label: str | None = None) -> bool:
+        """run all the tasks in the configs, write one result into the path"""
+        if self.running_task is not None:
+            log.warning("There're still tasks running in the background")
+            return False
+
+        if len(tasks) == 0:
+            log.warning("Empty tasks submitted")
+            return False
+
+        log.debug(f"tasks: {tasks}")
+
+        # Generate run_id
+        run_id = uuid.uuid4().hex
+        log.info(f"generated uuid for the tasks: {run_id}")
+        task_label = task_label if task_label else run_id
+
+        self.receive_conn, send_conn = mp.Pipe()
+        self.latest_error = ""
+        self.running_task = Assembler.assemble_all(run_id, task_label, tasks)
+        self.running_task.display()
+
+        return self._run_async(send_conn)
+
+    def get_results(self, result_dir: pathlib.Path | None = None) -> list[TestResult]:
+        """results of all runs, each TestResult represents one run."""
+        target_dir = result_dir if result_dir else config.RESULTS_LOCAL_DIR
+        return ResultCollector.collect(target_dir)
+
+    def _try_get_signal(self):
+        if self.receive_conn and self.receive_conn.poll():
+            sig, received = self.receive_conn.recv()
+            log.debug(f"Sigal received to process: {sig}, {received}")
+            if sig == SIGNAL.ERROR:
+                self.latest_error = received
+                self._clear_running_task()
+            elif sig == SIGNAL.SUCCESS:
+                global global_result_future
+                global_result_future = None
+                self.running_task = None
+                self.receive_conn = None
+            elif sig == SIGNAL.WIP:
+                self.running_task.set_finished(received)
+            else:
+                self._clear_running_task()
+
+    def has_running(self) -> bool:
+        """check if there're running benchmarks"""
+        if self.running_task:
+            self._try_get_signal()
+        return self.running_task is not None
+
+    def stop_running(self):
+        """force stop if ther're running benchmarks"""
+        self._clear_running_task()
+
+    def get_tasks_count(self) -> int:
+        """the count of all tasks"""
+        if self.running_task:
+            return self.running_task.num_cases()
+        return 0
+
+
+    def get_current_task_id(self) -> int:
+        """ the index of current running task
+        return -1 if not running
+        """
+        if not self.running_task:
+            return -1
+        return self.running_task.num_finished()
+
+    def _sync_running_task(self):
+        if not self.running_task:
+            return
+
+        global global_result_future
+        try:
+            if global_result_future:
+                global_result_future.result()
+        except Exception as e:
+            log.warning(f"task running failed: {e}", exc_info=True)
+        finally:
+            global_result_future = None
+            self.running_task = None
+
+    def _async_task_v2(self, running_task: TaskRunner, send_conn: Connection) -> None:
+        try:
+            if not running_task:
+                return
+
+            c_results = []
+            latest_runner, cached_load_duration = None, None
+            for idx, runner in enumerate(running_task.case_runners):
+                case_res = CaseResult(
+                    result_id=idx,
+                    metrics=Metric(),
+                    task_config=runner.config,
+                )
+
+                drop_old = False if latest_runner and runner == latest_runner else config.DROP_OLD
+                try:
+                    log.info(f"[{idx+1}/{running_task.num_cases()}] start case: {runner.display()}, drop_old={drop_old}")
+                    case_res.metrics = runner.run(drop_old)
+                    log.info(f"[{idx+1}/{running_task.num_cases()}] finish case: {runner.display()}, "
+                        f"result={case_res.metrics}, label={case_res.label}")
+
+                    # cache the latest succeeded runner
+                    latest_runner = runner
+
+                    # cache the latest drop_old=True load_duration of the latest succeeded runner
+                    cached_load_duration = case_res.metrics.load_duration if drop_old else cached_load_duration
+
+                    # use the cached load duration if this case didn't drop the existing collection
+                    if not drop_old:
+                        case_res.metrics.load_duration = cached_load_duration if cached_load_duration else 0.0
+                except (LoadTimeoutError, PerformanceTimeoutError) as e:
+                    log.warning(f"[{idx+1}/{running_task.num_cases()}] case {runner.display()} failed to run, reason={e}")
+                    case_res.label = ResultLabel.OUTOFRANGE
+                    continue
+
+                except Exception as e:
+                    log.warning(f"[{idx+1}/{running_task.num_cases()}] case {runner.display()} failed to run, reason={e}")
+                    traceback.print_exc()
+                    case_res.label = ResultLabel.FAILED
+                    continue
+
+                finally:
+                    c_results.append(case_res)
+                    send_conn.send((SIGNAL.WIP, idx))
+
+
+            test_result = TestResult(
+                run_id=running_task.run_id,
+                task_label=running_task.task_label,
+                results=c_results,
+            )
+            test_result.display()
+            test_result.write_file()
+
+            send_conn.send((SIGNAL.SUCCESS, None))
+            send_conn.close()
+            log.info(f"Succes to finish task: label={running_task.task_label}, run_id={running_task.run_id}")
+
+        except Exception as e:
+            err_msg = f"An error occurs when running task={running_task.task_label}, run_id={running_task.run_id}, err={e}"
+            traceback.print_exc()
+            log.warning(err_msg)
+            send_conn.send((SIGNAL.ERROR, err_msg))
+            send_conn.close()
+            return
+
+    def _clear_running_task(self):
+        global global_result_future
+        global_result_future = None
+
+        if self.running_task:
+            log.info(f"will force stop running task: {self.running_task.run_id}")
+            for r in self.running_task.case_runners:
+                r.stop()
+
+            self.kill_proc_tree(timeout=5)
+            self.running_task = None
+
+        if self.receive_conn:
+            self.receive_conn.close()
+            self.receive_conn = None
+
+
+    def _run_async(self, conn: Connection) -> bool:
+        log.info(f"task submitted: id={self.running_task.run_id}, {self.running_task.task_label}, case number: {len(self.running_task.case_runners)}")
+        global global_result_future
+        executor = concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=mp.get_context("spawn"))
+        global_result_future = executor.submit(self._async_task_v2, self.running_task, conn)
+
+        return True
+
+    def kill_proc_tree(self, sig=signal.SIGTERM, timeout=None, on_terminate=None):
+        """Kill a process tree (including grandchildren) with signal
+        "sig" and return a (gone, still_alive) tuple.
+        "on_terminate", if specified, is a callback function which is
+        called as soon as a child terminates.
+        """
+        children = psutil.Process().children(recursive=True)
+        for p in  children:
+            try:
+                log.warning(f"sending SIGTERM to child process: {p}")
+                p.send_signal(sig)
+            except psutil.NoSuchProcess:
+                pass
+        gone, alive = psutil.wait_procs(children, timeout=timeout,
+                                        callback=on_terminate)
+
+        for p in alive:
+            log.warning(f"force killing child process: {p}")
+            p.kill()
+
+
+benchMarkRunner = BenchMarkRunner()

vectordb_bench/log_util.py

@@ -0,0 +1,103 @@
+import logging
+from logging import config
+
+
+def init(log_level):
+    LOGGING = {
+        'version': 1,
+        'disable_existing_loggers': False,
+        'formatters': {
+            'default': {
+                'format': '%(asctime)s | %(levelname)s |%(message)s (%(filename)s:%(lineno)s)',
+            },
+            'colorful_console': {
+                'format': '%(asctime)s | %(levelname)s: %(message)s (%(filename)s:%(lineno)s) (%(process)s)',
+                '()': ColorfulFormatter,
+            },
+        },
+        'handlers': {
+            'console': {
+                'class': 'logging.StreamHandler',
+                'formatter': 'colorful_console',
+            },
+            'no_color_console': {
+                'class': 'logging.StreamHandler',
+                'formatter': 'default',
+            },
+        },
+        'loggers': {
+            'vectordb_bench': {
+                'handlers': ['console'],
+                'level': log_level,
+                'propagate': False
+            },
+            'no_color': {
+                'handlers': ['no_color_console'],
+                'level': log_level,
+                'propagate': False
+            },
+        },
+        'propagate': False,
+    }
+
+    config.dictConfig(LOGGING)
+
+class colors:
+    HEADER= '\033[95m'
+    INFO= '\033[92m'
+    DEBUG= '\033[94m'
+    WARNING= '\033[93m'
+    ERROR= '\033[95m'
+    CRITICAL= '\033[91m'
+    ENDC= '\033[0m'
+
+
+
+COLORS = {
+    'INFO': colors.INFO,
+    'INFOM': colors.INFO,
+    'DEBUG': colors.DEBUG,
+    'DEBUGM': colors.DEBUG,
+    'WARNING': colors.WARNING,
+    'WARNINGM': colors.WARNING,
+    'CRITICAL': colors.CRITICAL,
+    'CRITICALM': colors.CRITICAL,
+    'ERROR': colors.ERROR,
+    'ERRORM': colors.ERROR,
+    'ENDC': colors.ENDC,
+}
+
+
+class ColorFulFormatColMixin:
+    def format_col(self, message_str, level_name):
+        if level_name in COLORS.keys():
+            message_str = COLORS[level_name] + message_str + COLORS['ENDC']
+        return message_str
+
+    def formatTime(self, record, datefmt=None):
+        ret = super().formatTime(record, datefmt)
+        return ret
+
+
+class ColorfulLogRecordProxy(logging.LogRecord):
+    def __init__(self, record):
+        self._record = record
+        msg_level = record.levelname + 'M'
+        self.msg = f"{COLORS[msg_level]}{record.msg}{COLORS['ENDC']}"
+        self.filename = record.filename
+        self.lineno = f'{record.lineno}'
+        self.process = f'{record.process}'
+        self.levelname = f"{COLORS[record.levelname]}{record.levelname}{COLORS['ENDC']}"
+
+    def __getattr__(self, attr):
+        if attr not in self.__dict__:
+            return getattr(self._record, attr)
+        return getattr(self, attr)
+
+
+class ColorfulFormatter(ColorFulFormatColMixin, logging.Formatter):
+    def format(self, record):
+        proxy = ColorfulLogRecordProxy(record)
+        message_str = super().format(proxy)
+
+        return message_str

vectordb_bench/metric.py

@@ -0,0 +1,53 @@
+import logging
+import numpy as np
+
+from dataclasses import dataclass
+
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class Metric:
+    """result metrics"""
+
+    # for load cases
+    max_load_count: int = 0
+
+    # for performance cases
+    load_duration: float = 0.0  # duration to load all dataset into DB
+    qps: float = 0.0
+    serial_latency_p99: float = 0.0
+    recall: float = 0.0
+
+metricUnitMap = {
+    'load_duration': 's',
+    'serial_latency_p99': 'ms',
+    'max_load_count': 'K'
+}
+
+lowerIsBetterMetricList = [
+    "load_duration",
+    "serial_latency_p99",
+]
+
+metricOrder = [
+    "qps",
+    "recall",
+    "load_duration",
+    "serial_latency_p99",
+    "max_load_count",
+]
+
+
+def isLowerIsBetterMetric(metric: str) -> bool:
+    return metric in lowerIsBetterMetricList
+
+
+def calc_recall(count: int, ground_truth: list[int], got: list[int]) -> float:
+    recalls = np.zeros(count)
+    for i, result in enumerate(got):
+        if result in ground_truth:
+            recalls[i] = 1
+
+    return np.mean(recalls)

vectordb_bench/models.py

@@ -0,0 +1,234 @@
+import logging
+import pathlib
+from datetime import date
+from typing import Self
+from enum import Enum
+
+import ujson
+
+from .backend.clients import (
+    DB,
+    DBConfig,
+    DBCaseConfig,
+    IndexType,
+)
+from .base import BaseModel
+from . import config
+from .metric import Metric
+
+
+log = logging.getLogger(__name__)
+
+
+class LoadTimeoutError(TimeoutError):
+    pass
+
+class PerformanceTimeoutError(TimeoutError):
+    pass
+
+
+class CaseType(Enum):
+    """
+    Value will be displayed in UI
+    """
+
+    CapacitySDim = "Capacity Test (Large-dim)"
+    CapacityLDim = "Capacity Test (Small-dim)"
+
+    Performance100M = "Search Performance Test (XLarge Dataset)"
+    PerformanceLZero = "Search Performance Test (Large Dataset)"
+    PerformanceMZero = "Search Performance Test (Medium Dataset)"
+    PerformanceSZero = "Search Performance Test (Small Dataset)"
+
+    PerformanceLLow = (
+        "Filtering Search Performance Test (Large Dataset, Low Filtering Rate)"
+    )
+    PerformanceMLow = (
+        "Filtering Search Performance Test (Medium Dataset, Low Filtering Rate)"
+    )
+    PerformanceSLow = (
+        "Filtering Search Performance Test (Small Dataset, Low Filtering Rate)"
+    )
+    PerformanceLHigh = (
+        "Filtering Search Performance Test (Large Dataset, High Filtering Rate)"
+    )
+    PerformanceMHigh = (
+        "Filtering Search Performance Test (Medium Dataset, High Filtering Rate)"
+    )
+    PerformanceSHigh = (
+        "Filtering Search Performance Test (Small Dataset, High Filtering Rate)"
+    )
+
+
+class CaseConfigParamType(Enum):
+    """
+    Value will be the key of CaseConfig.params and displayed in UI
+    """
+
+    IndexType = "IndexType"
+    M = "M"
+    EFConstruction = "efConstruction"
+    EF = "ef"
+    SearchList = "search_list"
+    Nlist = "nlist"
+    Nprobe = "nprobe"
+    MaxConnections = "maxConnections"
+    numCandidates = "num_candidates"
+
+
+class CustomizedCase(BaseModel):
+    pass
+
+
+class CaseConfig(BaseModel):
+    """cases, dataset, test cases, filter rate, params"""
+
+    case_id: CaseType
+    custom_case: dict | None = None
+
+
+class TaskConfig(BaseModel):
+    db: DB
+    db_config: DBConfig
+    db_case_config: DBCaseConfig
+    case_config: CaseConfig
+
+    @property
+    def db_name(self):
+        db = self.db.value
+        db_label = self.db_config.db_label
+        return f"{db}-{db_label}" if db_label else db
+
+
+class ResultLabel(Enum):
+    NORMAL = ":)"
+    FAILED = "x"
+    OUTOFRANGE = "?"
+
+
+class CaseResult(BaseModel):
+    metrics: Metric
+    task_config: TaskConfig
+    label: ResultLabel = ResultLabel.NORMAL
+
+
+class TestResult(BaseModel):
+    """ROOT/result_{date.today()}_{task_label}.json"""
+
+    run_id: str
+    task_label: str
+    results: list[CaseResult]
+
+    def write_file(self):
+        result_dir = config.RESULTS_LOCAL_DIR
+        if not result_dir.exists():
+            log.info(f"local result directory not exist, creating it: {result_dir}")
+            result_dir.mkdir(parents=True)
+
+        file_name = f'result_{date.today().strftime("%Y%m%d")}_{self.task_label}.json'
+        result_file = result_dir.joinpath(file_name)
+        if result_file.exists():
+            log.warning(
+                f"Replacing existing result with the same file_name: {result_file}"
+            )
+
+        log.info(f"write results to disk {result_file}")
+        with open(result_file, "w") as f:
+            b = self.json(exclude={"db_config": {"password", "api_key"}})
+            f.write(b)
+
+    @classmethod
+    def read_file(cls, full_path: pathlib.Path, trans_unit: bool = False) -> Self:
+        if not full_path.exists():
+            raise ValueError(f"No such file: {full_path}")
+
+        with open(full_path) as f:
+            test_result = ujson.loads(f.read())
+            if "task_label" not in test_result:
+                test_result["task_label"] = test_result["run_id"]
+
+            for case_result in test_result["results"]:
+                task_config = case_result.get("task_config")
+                db = DB(task_config.get("db"))
+                dbcls = db.init_cls
+                task_config["db_config"] = dbcls.config_cls()(
+                    **task_config["db_config"]
+                )
+                task_config["db_case_config"] = dbcls.case_config_cls(
+                    index_type=task_config["db_case_config"].get("index", None),
+                )(**task_config["db_case_config"])
+
+                case_result["task_config"] = task_config
+
+                if trans_unit:
+                    cur_max_count = case_result["metrics"]["max_load_count"]
+                    case_result["metrics"]["max_load_count"] = (
+                        cur_max_count / 1000
+                        if int(cur_max_count) > 0
+                        else cur_max_count
+                    )
+
+                    cur_latency = case_result["metrics"]["serial_latency_p99"]
+                    case_result["metrics"]["serial_latency_p99"] = (
+                        cur_latency * 1000 if cur_latency > 0 else cur_latency
+                    )
+            c = TestResult.validate(test_result)
+
+            return c
+
+    def display(self, dbs: list[DB] | None = None):
+        filter_list = dbs if dbs and isinstance(dbs, list) else None
+        sorted_results = sorted(self.results, key=lambda x: (
+            x.task_config.db.name,
+            x.task_config.db_config.db_label,
+            x.task_config.case_config.case_id.name,
+        ), reverse=True)
+
+        filtered_results = [r for r in sorted_results  if not filter_list or r.task_config.db not in filter_list]
+
+        def append_return(x, y):
+            x.append(y)
+            return x
+
+        max_db = max(map(len, [f.task_config.db.name for f in filtered_results]))
+        max_db_labels = max(map(len, [f.task_config.db_config.db_label for f in filtered_results])) + 3
+        max_case = max(map(len, [f.task_config.case_config.case_id.name for f in filtered_results]))
+        max_load_dur = max(map(len, [str(f.metrics.load_duration) for f in filtered_results])) + 3
+        max_qps = max(map(len, [str(f.metrics.qps) for f in filtered_results])) + 3
+        max_recall = max(map(len, [str(f.metrics.recall) for f in filtered_results])) + 3
+
+        max_db_labels = 8 if max_db_labels == 0 else max_db_labels
+        max_load_dur = 11 if max_load_dur == 0 else max_load_dur + 3
+        max_qps = 10 if max_qps == 0 else max_load_dur + 3
+        max_recall = 13 if max_recall == 0 else max_recall + 3
+
+        LENGTH = (max_db, max_db_labels, max_case, len(self.task_label), max_load_dur, max_qps, 15, max_recall, 14)
+
+        DATA_FORMAT = (
+            f"%-{max_db}s | %-{max_db_labels}s %-{max_case}s %-{len(self.task_label)}s "
+            f"| %-{max_load_dur}s %-{max_qps}s %-15s %-{max_recall}s %-14s"
+        )
+
+        TITLE = DATA_FORMAT % (
+            "DB", "db_label", "case", "label", "load_dur", "qps", "latency(p99)", "recall", "max_load_count")
+        SPLIT = DATA_FORMAT%tuple(map(lambda x:"-"*x, LENGTH))
+        SUMMERY_FORMAT = ("Task summery: run_id=%s, task_label=%s") % (self.run_id[:5], self.task_label)
+        fmt = [SUMMERY_FORMAT, TITLE, SPLIT]
+
+
+        for f in filtered_results:
+            fmt.append(DATA_FORMAT%(
+                f.task_config.db.name,
+                f.task_config.db_config.db_label,
+                f.task_config.case_config.case_id.name,
+                self.task_label,
+                f.metrics.load_duration,
+                f.metrics.qps,
+                f.metrics.serial_latency_p99,
+                f.metrics.recall,
+                f.metrics.max_load_count,
+            ))
+
+        tmp_logger = logging.getLogger("no_color")
+        for f in fmt:
+            tmp_logger.info(f)