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

vectordb_bench/frontend/pages/custom.py

@@ -1,13 +1,21 @@
+from functools import partial
 import streamlit as st
 from vectordb_bench.frontend.components.check_results.headerIcon import drawHeaderIcon
-from vectordb_bench.frontend.components.custom.displayCustomCase import displayCustomCase
+from vectordb_bench.frontend.components.custom.displayCustomCase import (
+    displayCustomCase,
+)
 from vectordb_bench.frontend.components.custom.displaypPrams import displayParams
-from vectordb_bench.frontend.components.custom.getCustomConfig import CustomCaseConfig, generate_custom_case, get_custom_configs, save_custom_configs
+from vectordb_bench.frontend.components.custom.getCustomConfig import (
+    CustomCaseConfig,
+    generate_custom_case,
+    get_custom_configs,
+    save_custom_configs,
+)
 from vectordb_bench.frontend.components.custom.initStyle import initStyle
 from vectordb_bench.frontend.config.styles import FAVICON, PAGE_TITLE
 
 
-class CustomCaseManager():
+class CustomCaseManager:
     customCaseItems: list[CustomCaseConfig]
 
     def __init__(self):
@@ -52,12 +60,25 @@ def main():
 
         columns = expander.columns(8)
         columns[0].button(
-            "Save", key=f"{key}_", type="secondary", on_click=lambda: customCaseManager.save())
-        columns[1].button(":red[Delete]", key=f"{key}_delete", type="secondary",
-                          on_click=lambda: customCaseManager.deleteCase(idx))
-
-    st.button("\+ New Dataset", key=f"add_custom_configs",
-              type="primary", on_click=lambda: customCaseManager.addCase())
+            "Save",
+            key=f"{key}_",
+            type="secondary",
+            on_click=lambda: customCaseManager.save(),
+        )
+        columns[1].button(
+            ":red[Delete]",
+            key=f"{key}_delete",
+            type="secondary",
+            # B023
+            on_click=partial(lambda idx: customCaseManager.deleteCase(idx), idx=idx),
+        )
+
+    st.button(
+        "\+ New Dataset",
+        key="add_custom_configs",
+        type="primary",
+        on_click=lambda: customCaseManager.addCase(),
+    )
 
 
 if __name__ == "__main__":

vectordb_bench/frontend/pages/quries_per_dollar.py

@@ -15,8 +15,8 @@ from vectordb_bench.frontend.components.check_results.nav import (
 from vectordb_bench.frontend.components.check_results.charts import drawMetricChart
 from vectordb_bench.frontend.components.check_results.filters import getshownData
 from vectordb_bench.frontend.components.get_results.saveAsImage import getResults
-from vectordb_bench.frontend.config.styles import *
-from vectordb_bench.interface import benchMarkRunner
+
+from vectordb_bench.interface import benchmark_runner
 from vectordb_bench.metric import QURIES_PER_DOLLAR_METRIC
 
 
@@ -27,7 +27,7 @@ def main():
     # header
     drawHeaderIcon(st)
 
-    allResults = benchMarkRunner.get_results()
+    allResults = benchmark_runner.get_results()
 
     st.title("Vector DB Benchmark (QP$)")
 

vectordb_bench/frontend/pages/run_test.py

@@ -15,10 +15,10 @@ from vectordb_bench.frontend.components.check_results.stPageConfig import initRu
 def main():
     # set page config
     initRunTestPageConfig(st)
-    
+
     # init style
     initStyle(st)
-    
+
     # header
     drawHeaderIcon(st)
 
@@ -48,11 +48,7 @@ def main():
     activedCaseList, allCaseConfigs = caseSelector(caseSelectorContainer, activedDbList)
 
     # generate tasks
-    tasks = (
-        generate_tasks(activedDbList, dbConfigs, activedCaseList, allCaseConfigs)
-        if isAllValid
-        else []
-    )
+    tasks = generate_tasks(activedDbList, dbConfigs, activedCaseList, allCaseConfigs) if isAllValid else []
 
     # submit
     submitContainer = st.container()

vectordb_bench/frontend/utils.py

@@ -18,5 +18,5 @@ def addHorizontalLine(st):
 
 def generate_random_string(length):
     letters = string.ascii_letters + string.digits
-    result = ''.join(random.choice(letters) for _ in range(length))
+    result = "".join(random.choice(letters) for _ in range(length))
     return result

vectordb_bench/frontend/vdb_benchmark.py

@@ -11,8 +11,8 @@ from vectordb_bench.frontend.components.check_results.nav import (
 from vectordb_bench.frontend.components.check_results.charts import drawCharts
 from vectordb_bench.frontend.components.check_results.filters import getshownData
 from vectordb_bench.frontend.components.get_results.saveAsImage import getResults
-from vectordb_bench.frontend.config.styles import *
-from vectordb_bench.interface import benchMarkRunner
+
+from vectordb_bench.interface import benchmark_runner
 
 
 def main():
@@ -22,7 +22,7 @@ def main():
     # header
     drawHeaderIcon(st)
 
-    allResults = benchMarkRunner.get_results()
+    allResults = benchmark_runner.get_results()
 
     st.title("Vector Database Benchmark")
     st.caption(
@@ -32,9 +32,7 @@ def main():
 
     # results selector and filter
     resultSelectorContainer = st.sidebar.container()
-    shownData, failedTasks, showCaseNames = getshownData(
-        allResults, resultSelectorContainer
-    )
+    shownData, failedTasks, showCaseNames = getshownData(allResults, resultSelectorContainer)
 
     resultSelectorContainer.divider()
 

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