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

vectordb_bench/interface.py

@@ -5,6 +5,7 @@ import pathlib
 import signal
 import traceback
 import uuid
+from collections.abc import Callable
 from enum import Enum
 from multiprocessing.connection import Connection
 
@@ -16,8 +17,15 @@ from .backend.data_source import DatasetSource
 from .backend.result_collector import ResultCollector
 from .backend.task_runner import TaskRunner
 from .metric import Metric
-from .models import (CaseResult, LoadTimeoutError, PerformanceTimeoutError,
-                     ResultLabel, TaskConfig, TaskStage, TestResult)
+from .models import (
+    CaseResult,
+    LoadTimeoutError,
+    PerformanceTimeoutError,
+    ResultLabel,
+    TaskConfig,
+    TaskStage,
+    TestResult,
+)
 
 log = logging.getLogger(__name__)
 
@@ -37,11 +45,9 @@ class BenchMarkRunner:
         self.drop_old: bool = True
         self.dataset_source: DatasetSource = DatasetSource.S3
 
-
     def set_drop_old(self, drop_old: bool):
         self.drop_old = drop_old
 
-
     def set_download_address(self, use_aliyun: bool):
         if use_aliyun:
             self.dataset_source = DatasetSource.AliyunOSS
@@ -59,7 +65,9 @@ class BenchMarkRunner:
             log.warning("Empty tasks submitted")
             return False
 
-        log.debug(f"tasks: {tasks}, task_label: {task_label}, dataset source: {self.dataset_source}")
+        log.debug(
+            f"tasks: {tasks}, task_label: {task_label}, dataset source: {self.dataset_source}",
+        )
 
         # Generate run_id
         run_id = uuid.uuid4().hex
@@ -70,7 +78,12 @@ class BenchMarkRunner:
         self.latest_error = ""
 
         try:
-            self.running_task = Assembler.assemble_all(run_id, task_label, tasks, self.dataset_source)
+            self.running_task = Assembler.assemble_all(
+                run_id,
+                task_label,
+                tasks,
+                self.dataset_source,
+            )
             self.running_task.display()
         except ModuleNotFoundError as e:
             msg = f"Please install client for database, error={e}"
@@ -119,7 +132,7 @@ class BenchMarkRunner:
         return 0
 
     def get_current_task_id(self) -> int:
-        """ the index of current running task
+        """the index of current running task
         return -1 if not running
         """
         if not self.running_task:
@@ -153,18 +166,18 @@ class BenchMarkRunner:
                     task_config=runner.config,
                 )
 
-                # drop_old = False if latest_runner and runner == latest_runner else config.DROP_OLD
-                # drop_old = config.DROP_OLD
                 drop_old = TaskStage.DROP_OLD in runner.config.stages
-                if latest_runner and runner == latest_runner:
-                    drop_old = False
-                elif not self.drop_old:
+                if (latest_runner and runner == latest_runner) or not self.drop_old:
                     drop_old = False
                 try:
-                    log.info(f"[{idx+1}/{running_task.num_cases()}] start case: {runner.display()}, drop_old={drop_old}")
+                    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}")
+                    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
@@ -176,12 +189,16 @@ class BenchMarkRunner:
                     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}")
+                    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}")
+                    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
@@ -200,10 +217,14 @@ class BenchMarkRunner:
 
             send_conn.send((SIGNAL.SUCCESS, None))
             send_conn.close()
-            log.info(f"Success to finish task: label={running_task.task_label}, run_id={running_task.run_id}")
+            log.info(
+                f"Success 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}"
+            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))
@@ -226,16 +247,26 @@ class BenchMarkRunner:
             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)}")
+        log.info(
+            f"task submitted: id={self.running_task.run_id}, {self.running_task.task_label}, ",
+            f"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"))
+        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):
+    def kill_proc_tree(
+        self,
+        sig: int = signal.SIGTERM,
+        timeout: float | None = None,
+        on_terminate: Callable | None = 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
@@ -248,12 +279,11 @@ class BenchMarkRunner:
                 p.send_signal(sig)
             except psutil.NoSuchProcess:
                 pass
-        gone, alive = psutil.wait_procs(children, timeout=timeout,
-                                        callback=on_terminate)
+        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()
+benchmark_runner = BenchMarkRunner()

vectordb_bench/log_util.py

@@ -1,102 +1,97 @@
 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)',
+
+def init(log_level: str):
+    log_config = {
+        "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,
+            "colorful_console": {
+                "format": "%(asctime)s | %(levelname)s: %(message)s (%(filename)s:%(lineno)s) (%(process)s)",
+                "()": ColorfulFormatter,
             },
         },
-        'handlers': {
-            'console': {
-                'class': 'logging.StreamHandler',
-                'formatter': 'colorful_console',
+        "handlers": {
+            "console": {
+                "class": "logging.StreamHandler",
+                "formatter": "colorful_console",
             },
-            'no_color_console': {
-                'class': 'logging.StreamHandler',
-                'formatter': 'default',
+            "no_color_console": {
+                "class": "logging.StreamHandler",
+                "formatter": "default",
             },
         },
-        'loggers': {
-            'vectordb_bench': {
-                'handlers': ['console'],
-                'level': log_level,
-                'propagate': False
+        "loggers": {
+            "vectordb_bench": {
+                "handlers": ["console"],
+                "level": log_level,
+                "propagate": False,
             },
-            'no_color': {
-                'handlers': ['no_color_console'],
-                'level': log_level,
-                'propagate': False
+            "no_color": {
+                "handlers": ["no_color_console"],
+                "level": log_level,
+                "propagate": False,
             },
         },
-        'propagate': False,
+        "propagate": False,
     }
 
-    config.dictConfig(LOGGING)
+    config.dictConfig(log_config)
 
-class colors:
-    HEADER= '\033[95m'
-    INFO= '\033[92m'
-    DEBUG= '\033[94m'
-    WARNING= '\033[93m'
-    ERROR= '\033[95m'
-    CRITICAL= '\033[91m'
-    ENDC= '\033[0m'
 
+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,
+    "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
+    def format_col(self, message: str, level_name: str):
+        if level_name in COLORS:
+            message = COLORS[level_name] + message + COLORS["ENDC"]
+        return message
 
 
 class ColorfulLogRecordProxy(logging.LogRecord):
-    def __init__(self, record):
+    def __init__(self, record: any):
         self._record = record
-        msg_level = record.levelname + 'M'
+        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.lineno = f"{record.lineno}"
+        self.process = f"{record.process}"
         self.levelname = f"{COLORS[record.levelname]}{record.levelname}{COLORS['ENDC']}"
 
-    def __getattr__(self, attr):
+    def __getattr__(self, attr: any):
         if attr not in self.__dict__:
             return getattr(self._record, attr)
         return getattr(self, attr)
 
 
 class ColorfulFormatter(ColorFulFormatColMixin, logging.Formatter):
-    def format(self, record):
+    def format(self, record: any):
         proxy = ColorfulLogRecordProxy(record)
-        message_str = super().format(proxy)
-
-        return message_str
+        return super().format(proxy)

vectordb_bench/metric.py

@@ -1,8 +1,7 @@
 import logging
-import numpy as np
-
 from dataclasses import dataclass, field
 
+import numpy as np
 
 log = logging.getLogger(__name__)
 
@@ -33,19 +32,19 @@ MAX_LOAD_COUNT_METRIC = "max_load_count"
 QPS_METRIC = "qps"
 RECALL_METRIC = "recall"
 
-metricUnitMap = {
+metric_unit_map = {
     LOAD_DURATION_METRIC: "s",
     SERIAL_LATENCY_P99_METRIC: "ms",
     MAX_LOAD_COUNT_METRIC: "K",
     QURIES_PER_DOLLAR_METRIC: "K",
 }
 
-lowerIsBetterMetricList = [
+lower_is_better_metrics = [
     LOAD_DURATION_METRIC,
     SERIAL_LATENCY_P99_METRIC,
 ]
 
-metricOrder = [
+metric_order = [
     QPS_METRIC,
     RECALL_METRIC,
     LOAD_DURATION_METRIC,
@@ -55,7 +54,7 @@ metricOrder = [
 
 
 def isLowerIsBetterMetric(metric: str) -> bool:
-    return metric in lowerIsBetterMetricList
+    return metric in lower_is_better_metrics
 
 
 def calc_recall(count: int, ground_truth: list[int], got: list[int]) -> float:
@@ -70,7 +69,7 @@ def calc_recall(count: int, ground_truth: list[int], got: list[int]) -> float:
 def get_ideal_dcg(k: int):
     ideal_dcg = 0
     for i in range(k):
-        ideal_dcg += 1 / np.log2(i+2)
+        ideal_dcg += 1 / np.log2(i + 2)
 
     return ideal_dcg
 
@@ -78,8 +77,8 @@ def get_ideal_dcg(k: int):
 def calc_ndcg(ground_truth: list[int], got: list[int], ideal_dcg: float) -> float:
     dcg = 0
     ground_truth = list(ground_truth)
-    for id in set(got):
-        if id in ground_truth:
-            idx = ground_truth.index(id)
-            dcg += 1 / np.log2(idx+2)
+    for got_id in set(got):
+        if got_id in ground_truth:
+            idx = ground_truth.index(got_id)
+            dcg += 1 / np.log2(idx + 2)
     return dcg / ideal_dcg

vectordb_bench/models.py

@@ -2,29 +2,31 @@ import logging
 import pathlib
 from datetime import date, datetime
 from enum import Enum, StrEnum, auto
-from typing import List, Self
+from typing import Self
 
 import ujson
 
+from . import config
+from .backend.cases import CaseType
 from .backend.clients import (
     DB,
-    DBConfig,
     DBCaseConfig,
+    DBConfig,
 )
-from .backend.cases import CaseType
 from .base import BaseModel
-from . import config
 from .metric import Metric
 
 log = logging.getLogger(__name__)
 
 
 class LoadTimeoutError(TimeoutError):
-    pass
+    def __init__(self, duration: int):
+        super().__init__(f"capacity case load timeout in {duration}s")
 
 
 class PerformanceTimeoutError(TimeoutError):
-    pass
+    def __init__(self):
+        super().__init__("Performance case optimize timeout")
 
 
 class CaseConfigParamType(Enum):
@@ -92,7 +94,7 @@ class CustomizedCase(BaseModel):
 
 
 class ConcurrencySearchConfig(BaseModel):
-    num_concurrency: List[int] = config.NUM_CONCURRENCY
+    num_concurrency: list[int] = config.NUM_CONCURRENCY
     concurrency_duration: int = config.CONCURRENCY_DURATION
 
 
@@ -146,7 +148,7 @@ class TaskConfig(BaseModel):
     db_config: DBConfig
     db_case_config: DBCaseConfig
     case_config: CaseConfig
-    stages: List[TaskStage] = ALL_TASK_STAGES
+    stages: list[TaskStage] = ALL_TASK_STAGES
 
     @property
     def db_name(self):
@@ -210,26 +212,23 @@ class TestResult(BaseModel):
             log.info(f"local result directory not exist, creating it: {result_dir}")
             result_dir.mkdir(parents=True)
 
-        file_name = self.file_fmt.format(
-            date.today().strftime("%Y%m%d"), partial.task_label, db
-        )
+        file_name = self.file_fmt.format(date.today().strftime("%Y%m%d"), partial.task_label, db)
         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.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:
+        with pathlib.Path(result_file).open("w") as f:
             b = partial.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}")
+            msg = f"No such file: {full_path}"
+            raise ValueError(msg)
 
-        with open(full_path) as f:
+        with pathlib.Path(full_path).open("r") as f:
             test_result = ujson.loads(f.read())
             if "task_label" not in test_result:
                 test_result["task_label"] = test_result["run_id"]
@@ -248,19 +247,16 @@ class TestResult(BaseModel):
                 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_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
+            return TestResult.validate(test_result)
 
+    # ruff: noqa
     def display(self, dbs: list[DB] | None = None):
         filter_list = dbs if dbs and isinstance(dbs, list) else None
         sorted_results = sorted(
@@ -273,31 +269,18 @@ class TestResult(BaseModel):
             reverse=True,
         )
 
-        filtered_results = [
-            r
-            for r in sorted_results
-            if not filter_list or r.task_config.db not in filter_list
-        ]
+        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):
+        def append_return(x: any, y: any):
             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_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_recall = max(map(len, [str(f.metrics.recall) for f in filtered_results])) + 3
 
         max_db_labels = 8 if max_db_labels < 8 else max_db_labels
         max_load_dur = 11 if max_load_dur < 11 else max_load_dur
@@ -356,7 +339,7 @@ class TestResult(BaseModel):
                     f.metrics.recall,
                     f.metrics.max_load_count,
                     f.label.value,
-                )
+                ),
             )
 
         tmp_logger = logging.getLogger("no_color")