vectordb-bench 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

vectordb_bench/backend/task_runner.py

@@ -1,4 +1,5 @@
 import logging
+import psutil
 import traceback
 import concurrent
 import numpy as np
@@ -7,7 +8,7 @@ from enum import Enum, auto
 from . import utils
 from .cases import Case, CaseLabel
 from ..base import BaseModel
-from ..models import TaskConfig
+from ..models import TaskConfig, PerformanceTimeoutError
 
 from .clients import (
     api,
@@ -92,80 +93,70 @@ class CaseRunner(BaseModel):
         self._pre_run(drop_old)
 
         if self.ca.label == CaseLabel.Load:
-            return self._run_load_case()
+            return self._run_capacity_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}")
+            msg = f"unknown case type: {self.ca.label}"
+            log.warning(msg)
+            raise ValueError(msg)
 
-
-    def _run_load_case(self) -> Metric:
-        """ run load cases
+    def _run_capacity_case(self) -> Metric:
+        """ run capacity 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:
+            runner = SerialInsertRunner(self.db, self.ca.dataset, self.normalize, self.ca.load_timeout)
             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}")
+            log.warning(f"Failed to run capacity case, reason = {e}")
             raise e from None
-        log.info("End capacity case")
-
+        else:
+            log.info(f"Capacity case loading dataset reaches VectorDB's limit: max capacity = {count}")
+            return Metric(max_load_count=count)
 
     def _run_perf_case(self, drop_old: bool = True) -> Metric:
+        """ run performance cases
+
+        Returns:
+            Metric: load_duration, recall, serial_latency_p99, and, qps
+        """
         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)
+                log.info(
+                    f"Finish loading the entire dataset into VectorDB,"
+                    f" insert_duration={load_dur}, optimize_duration={build_dur}"
+                    f" load_duration(insert + optimize) = {m.load_duration}"
+                )
 
             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}")
+            log.warning(f"Failed to run performance case, reason = {e}")
             traceback.print_exc()
-            raise e
+            raise e from None
+        else:
+            log.info(f"Performance case got result: {m}")
+            return m
 
     @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
-
+        try:
+            runner = SerialInsertRunner(self.db, self.ca.dataset, self.normalize, self.ca.load_timeout)
+            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,
@@ -198,17 +189,21 @@ class CaseRunner(BaseModel):
 
     @utils.time_it
     def _task(self) -> None:
-        """"""
         with self.db.init():
-            self.db.ready_to_search()
+            self.db.optimize()
 
     def _optimize(self) -> float:
         with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
             future = executor.submit(self._task)
             try:
-                return future.result()[1]
+                return future.result(timeout=self.ca.optimize_timeout)[1]
+            except TimeoutError as e:
+                log.warning(f"VectorDB optimize timeout in {self.ca.optimize_timeout}")
+                for pid, _ in executor._processes.items():
+                    psutil.Process(pid).kill()
+                raise PerformanceTimeoutError("Performance case optimize timeout") from e
             except Exception as e:
-                log.warning(f"VectorDB ready_to_search error: {e}")
+                log.warning(f"VectorDB optimize error: {e}")
                 raise e from None
 
     def _init_search_runner(self):

vectordb_bench/frontend/components/check_results/charts.py

@@ -1,30 +1,19 @@
+from vectordb_bench.backend.cases import Case
+from vectordb_bench.frontend.components.check_results.expanderStyle import initMainExpanderStyle
 from vectordb_bench.metric import metricOrder, isLowerIsBetterMetric, metricUnitMap
-from vectordb_bench.frontend.const import *
+from vectordb_bench.frontend.const.styles import *
 from vectordb_bench.models import ResultLabel
 import plotly.express as px
 
 
-def drawCharts(st, allData, failedTasks, cases):
-    st.markdown(
-        "<style> .main .streamlit-expanderHeader p {font-size: 20px; font-weight: 600;} </style>",
-        unsafe_allow_html=True,
-    )
-    st.markdown(
-        """<style>
-            .main div[data-testid='stExpander'] {
-                background-color: #F6F8FA;
-                border: 1px solid #A9BDD140;
-                border-radius: 8px;
-            }
-        </style>""",
-        unsafe_allow_html=True,
-    )
+def drawCharts(st, allData, failedTasks, cases: list[Case]):
+    initMainExpanderStyle(st)
     for case in cases:
-        chartContainer = st.expander(case, True)
-        data = [data for data in allData if data["case"] == case]
+        chartContainer = st.expander(case.name, True)
+        data = [data for data in allData if data["case_name"] == case.name]
         drawChart(data, chartContainer)
 
-        errorDBs = failedTasks[case]
+        errorDBs = failedTasks[case.name]
         showFailedDBs(chartContainer, errorDBs)
 
 
@@ -102,7 +91,7 @@ def drawMetricChart(data, metric, st):
     xmin = 0
     xmax = max([d.get(metric, 0) for d in dataWithMetric])
     xpadding = (xmax - xmin) / 16
-    xpadding_multiplier = 1.6
+    xpadding_multiplier = 1.8
     xrange = [xmin, xmax + xpadding * xpadding_multiplier]
     unit = metricUnitMap.get(metric, "")
     labelToShapeMap = getLabelToShapeMap(dataWithMetric)
@@ -136,7 +125,7 @@ def drawMetricChart(data, metric, st):
             font=dict(
                 size=1,
             ),
-            # text="",
+            text="",
         )
     )
     fig.update_traces(

vectordb_bench/frontend/components/check_results/data.py

@@ -1,63 +1,78 @@
 from collections import defaultdict
 from dataclasses import asdict
+from vectordb_bench.backend.cases import Case
 from vectordb_bench.metric import isLowerIsBetterMetric
-from vectordb_bench.models import ResultLabel
+from vectordb_bench.models import CaseResult, ResultLabel
 
 
-def getChartData(tasks, dbNames, cases):
+def getChartData(
+    tasks: list[CaseResult],
+    dbNames: list[str],
+    cases: list[Case],
+):
     filterTasks = getFilterTasks(tasks, dbNames, cases)
     mergedTasks, failedTasks = mergeTasks(filterTasks)
     return mergedTasks, failedTasks
 
 
-def getFilterTasks(tasks, dbNames, cases):
+def getFilterTasks(
+    tasks: list[CaseResult],
+    dbNames: list[str],
+    cases: list[Case],
+) -> list[CaseResult]:
+    case_ids = [case.case_id for case in cases]
     filterTasks = [
         task
         for task in tasks
         if task.task_config.db_name in dbNames
-        and task.task_config.case_config.case_id.value in cases
+        and task.task_config.case_config.case_id in case_ids
     ]
     return filterTasks
 
 
-def mergeTasks(tasks):
+def mergeTasks(tasks: list[CaseResult]):
     dbCaseMetricsMap = defaultdict(lambda: defaultdict(dict))
     for task in tasks:
         db_name = task.task_config.db_name
         db = task.task_config.db.value
         db_label = task.task_config.db_config.db_label or ""
-        case = task.task_config.case_config.case_id.value
-        dbCaseMetricsMap[db_name][case] = {
+        case_id = task.task_config.case_config.case_id
+        dbCaseMetricsMap[db_name][case_id] = {
             "db": db,
             "db_label": db_label,
             "metrics": mergeMetrics(
-                dbCaseMetricsMap[db_name][case].get("metrics", {}), asdict(task.metrics)
+                dbCaseMetricsMap[db_name][case_id].get("metrics", {}),
+                asdict(task.metrics),
+            ),
+            "label": getBetterLabel(
+                dbCaseMetricsMap[db_name][case_id].get("label", ResultLabel.FAILED),
+                task.label,
             ),
-            "label": getBetterLabel(dbCaseMetricsMap[db_name][case].get("label", ResultLabel.FAILED), task.label)
         }
 
     mergedTasks = []
     failedTasks = defaultdict(lambda: defaultdict(str))
     for db_name, caseMetricsMap in dbCaseMetricsMap.items():
-        for case, metricInfo in caseMetricsMap.items():
+        for case_id, metricInfo in caseMetricsMap.items():
             metrics = metricInfo["metrics"]
             db = metricInfo["db"]
             db_label = metricInfo["db_label"]
             label = metricInfo["label"]
+            case_name = case_id.case_name
             if label == ResultLabel.NORMAL:
                 mergedTasks.append(
                     {
                         "db_name": db_name,
                         "db": db,
                         "db_label": db_label,
-                        "case": case,
+                        "case_name": case_name,
                         "metricsSet": set(metrics.keys()),
                         **metrics,
                     }
                 )
-            else: 
-                failedTasks[case][db_name] = label
-                
+            else:
+                failedTasks[case_name][db_name] = label
+
     return mergedTasks, failedTasks
 
 
@@ -81,6 +96,7 @@ def getBetterMetric(metric, value_1, value_2):
         if isLowerIsBetterMetric(metric)
         else max(value_1, value_2)
     )
-    
+
+
 def getBetterLabel(label_1: ResultLabel, label_2: ResultLabel):
     return label_2 if label_1 != ResultLabel.NORMAL else label_1

vectordb_bench/frontend/components/check_results/expanderStyle.py

@@ -0,0 +1,37 @@
+def initMainExpanderStyle(st):
+    st.markdown(
+        """<style>
+            .main .streamlit-expanderHeader p {font-size: 20px; font-weight: 600;}
+            .main div[data-testid='stExpander'] {
+                background-color: #F6F8FA;
+                border: 1px solid #A9BDD140;
+                border-radius: 8px;
+            }
+        </style>""",
+        unsafe_allow_html=True,
+    )
+
+
+def initSidebarExanderStyle(st):
+    st.markdown(
+        """<style>
+            section[data-testid='stSidebar']
+                div[data-testid='stExpander']
+                    div[data-testid='stVerticalBlock']
+                        { gap: 0.2rem; }
+            div[data-testid='stExpander']
+                { background-color: #ffffff; }
+            section[data-testid='stSidebar'] 
+                .streamlit-expanderHeader 
+                    p { font-size: 16px; font-weight: 600; }
+            section[data-testid='stSidebar']
+                div[data-testid='stExpander']
+                    div[data-testid='stVerticalBlock'] 
+                        button {
+                            padding: 0 0.5rem;
+                            margin-bottom: 8px;
+                            float: right;
+                        }
+        <style>""",
+        unsafe_allow_html=True,
+    )

vectordb_bench/frontend/components/check_results/filters.py

@@ -1,8 +1,14 @@
+from vectordb_bench.backend.cases import Case
 from vectordb_bench.frontend.components.check_results.data import getChartData
-from vectordb_bench.frontend.const import *
+from vectordb_bench.frontend.components.check_results.expanderStyle import initSidebarExanderStyle
+from vectordb_bench.frontend.const.dbCaseConfigs import CASE_LIST
+from vectordb_bench.frontend.const.styles import *
+import streamlit as st
 
+from vectordb_bench.models import CaseResult, TestResult
 
-def getshownData(results, st):
+
+def getshownData(results: list[TestResult], st):
     # hide the nav
     st.markdown(
         "<style> div[data-testid='stSidebarNav'] {display: none;} </style>",
@@ -19,7 +25,7 @@ def getshownData(results, st):
     return shownData, failedTasks, showCases
 
 
-def getshownResults(results, st):
+def getshownResults(results: list[TestResult], st) -> list[CaseResult]:
     resultSelectOptions = [
         result.task_label
         if result.task_label != result.run_id
@@ -38,7 +44,7 @@ def getshownResults(results, st):
         # label_visibility="hidden",
         default=resultSelectOptions,
     )
-    selectedResult = []
+    selectedResult: list[CaseResult] = []
     for option in selectedResultSelectedOptions:
         result = results[resultSelectOptions.index(option)].results
         selectedResult += result
@@ -46,52 +52,74 @@ def getshownResults(results, st):
     return selectedResult
 
 
-def getShowDbsAndCases(result, st):
-    # expanderStyles
-    st.markdown("<style> section[data-testid='stSidebar'] div[data-testid='stExpander'] div[data-testid='stVerticalBlock'] { gap: 0.2rem; }  </style>", unsafe_allow_html=True,)
-    st.markdown(
-        "<style> div[data-testid='stExpander'] {background-color: #ffffff;} </style>",
-        unsafe_allow_html=True,
-    )
-    st.markdown(
-        "<style> section[data-testid='stSidebar'] .streamlit-expanderHeader p {font-size: 16px; font-weight: 600;} </style>",
-        unsafe_allow_html=True,
-    )
-
+def getShowDbsAndCases(result: list[CaseResult], st) -> tuple[list[str], list[Case]]:
+    initSidebarExanderStyle(st)
     allDbNames = list(set({res.task_config.db_name for res in result}))
     allDbNames.sort()
     allCasesSet = set({res.task_config.case_config.case_id for res in result})
-    allCases = [case["name"].value for case in CASE_LIST if case["name"] in allCasesSet]
-
-    # dbFilterContainer = st.container()
-    # dbFilterContainer.subheader("DB Filter")
-    dbFilterContainer = st.expander("DB Filter", True)
-    showDBNames = filterView(allDbNames, dbFilterContainer, col=1)
+    allCases: list[Case] = [case.case_cls() for case in CASE_LIST if case in allCasesSet]
+
+    # DB Filter
+    dbFilterContainer = st.container()
+    showDBNames = filterView(
+        dbFilterContainer,
+        "DB Filter",
+        allDbNames,
+        col=1,
+    )
 
-    # caseFilterContainer = st.container()
-    # caseFilterContainer.subheader("Case Filter")
-    caseFilterContainer = st.expander("Case Filter", True)
+    # Case Filter
+    caseFilterContainer = st.container()
     showCases = filterView(
-        allCases,
         caseFilterContainer,
+        "Case Filter",
+        [case for case in allCases],
         col=1,
-        optionLables=[case for case in allCases],
+        optionLables=[case.name for case in allCases],
     )
 
     return showDBNames, showCases
 
 
-def filterView(options, st, col, optionLables=None):
-    columns = st.columns(
+def filterView(container, header, options, col, optionLables=None):
+    selectAllState = f"{header}-select-all-state"
+    if selectAllState not in st.session_state:
+        st.session_state[selectAllState] = True
+
+    countKeyState = f"{header}-select-all-count-key"
+    if countKeyState not in st.session_state:
+        st.session_state[countKeyState] = 0
+
+    expander = container.expander(header, True)
+    selectAllColumns = expander.columns(SIDEBAR_CONTROL_COLUMNS, gap="small")
+    selectAllButton = selectAllColumns[SIDEBAR_CONTROL_COLUMNS - 2].button(
+        "select all",
+        key=f"{header}-select-all-button",
+        # type="primary",
+    )
+    clearAllButton = selectAllColumns[SIDEBAR_CONTROL_COLUMNS - 1].button(
+        "clear all",
+        key=f"{header}-clear-all-button",
+        # type="primary",
+    )
+    if selectAllButton:
+        st.session_state[selectAllState] = True
+        st.session_state[countKeyState] += 1
+    if clearAllButton:
+        st.session_state[selectAllState] = False
+        st.session_state[countKeyState] += 1
+    columns = expander.columns(
         col,
         gap="small",
     )
-    isActive = {option: True for option in options}
     if optionLables is None:
         optionLables = options
-    for i, option in enumerate(options):
+    isActive = {option: st.session_state[selectAllState] for option in optionLables}
+    for i, option in enumerate(optionLables):
         isActive[option] = columns[i % col].checkbox(
-            optionLables[i], value=isActive[option]
+            optionLables[i],
+            value=isActive[option],
+            key=f"{optionLables[i]}-{st.session_state[countKeyState]}",
         )
 
-    return [option for option in options if isActive[option]]
+    return [options[i] for i, option in enumerate(optionLables) if isActive[option]]

vectordb_bench/frontend/components/check_results/footer.py

@@ -0,0 +1,8 @@
+def footer(st):
+    text = "* All test results are from community contributors. If there is any ambiguity, feel free to raise an issue or make amendments on our <a href='https://github.com/zilliztech/VectorDBBench'>GitHub page</a>."
+    st.markdown(
+        f"""
+        <div style="margin-top: 16px; color: #aaa; font-size: 14px;">{text}</div
+        """,
+        unsafe_allow_html=True,
+    )

vectordb_bench/frontend/components/check_results/headerIcon.py

@@ -1,17 +1,21 @@
+from vectordb_bench.frontend.const.styles import HEADER_ICON
+
+
 def drawHeaderIcon(st):
-    st.markdown("""
+    st.markdown(
+        f"""
 <div class="headerIconContainer"></div>
 
 <style>
-.headerIconContainer {
+.headerIconContainer {{
     position: absolute;
     top: -50px;
     height: 50px;
     width: 100%;
     border-bottom: 2px solid #E8EAEE;
-    background-image: url(https://assets.zilliz.com/vdb_benchmark_db790b5387.png);
+    background-image: url({HEADER_ICON});
     background-repeat: no-repeat;
-}
+}}
 </style
 """,
         unsafe_allow_html=True,

vectordb_bench/frontend/components/check_results/nav.py

@@ -2,20 +2,21 @@ from streamlit_extras.switch_page_button import switch_page
 
 
 def NavToRunTest(st):
-    st.header("Run your test")
+    st.subheader("Run your test")
     st.write("You can set the configs and run your own test.")
     navClick = st.button("Run Your Test   >")
     if navClick:
         switch_page("run test")
         
         
-def NavToQPSWithPrice(st):
-    navClick = st.button("QPS with Price   >")
+def NavToQuriesPerDollar(st):
+    st.subheader("Compare qps with price.")
+    navClick = st.button("QP$ (Quries per Dollar)   >")
     if navClick:
-        switch_page("qps with price")
+        switch_page("quries_per_dollar")
         
         
-def NavToResults(st):
-    navClick = st.button("< &nbsp;&nbsp;Back to Results")
+def NavToResults(st, key="nav-to-results"):
+    navClick = st.button("< &nbsp;&nbsp;Back to Results", key=key)
     if navClick:
         switch_page("vdb benchmark")

vectordb_bench/frontend/components/check_results/priceTable.py

@@ -1,9 +1,10 @@
 from vectordb_bench.backend.clients import DB
-from vectordb_bench.frontend.const import DB_DBLABEL_TO_PRICE
 import pandas as pd
 from collections import defaultdict
 import streamlit as st
 
+from vectordb_bench.frontend.const.dbPrices import DB_DBLABEL_TO_PRICE
+
 
 def priceTable(container, data):
     dbAndLabelSet = {
@@ -25,7 +26,7 @@ def priceTable(container, data):
     )
     height = len(table) * 35 + 38
 
-    expander = container.expander("You can edit the price.")
+    expander = container.expander("Price List (Editable).")
     editTable = expander.data_editor(
         table,
         use_container_width=True,

vectordb_bench/frontend/components/check_results/stPageConfig.py

@@ -0,0 +1,18 @@
+from vectordb_bench.frontend.const.styles import *
+
+
+def initResultsPageConfig(st):
+    st.set_page_config(
+        page_title=PAGE_TITLE,
+        page_icon=FAVICON,
+        # layout="wide",
+        # initial_sidebar_state="collapsed",
+    )
+
+def initRunTestPageConfig(st):
+    st.set_page_config(
+        page_title=PAGE_TITLE,
+        page_icon=FAVICON,
+        # layout="wide",
+        initial_sidebar_state="collapsed",
+    )

vectordb_bench/frontend/components/get_results/saveAsImage.py

@@ -0,0 +1,50 @@
+import requests
+import streamlit as st
+import streamlit.components.v1 as components
+
+HTML_2_CANVAS_URL = "https://unpkg.com/html2canvas@1.4.1/dist/html2canvas.js"
+
+
+@st.cache_data
+def load_unpkg(src: str) -> str:
+    return requests.get(src).text
+
+def getResults(container, pageName="vectordb_bench"):
+    container.subheader("Get results")
+    saveAsImage(container, pageName)
+
+def saveAsImage(container, pageName):
+    html2canvasJS = load_unpkg(HTML_2_CANVAS_URL)
+    container.write()
+    buttonText = "Save as Image"
+    savePDFButton = container.button(buttonText)
+    if savePDFButton:
+        components.html(
+            f"""
+<script>{html2canvasJS}</script>
+
+<script>
+const html2canvas = window.html2canvas
+
+const streamlitDoc = window.parent.document;
+const stApp = streamlitDoc.querySelector('.main > .block-container');
+
+const buttons = Array.from(streamlitDoc.querySelectorAll('.stButton > button'));
+const imgButton = buttons.find(el => el.innerText === '{buttonText}');
+
+if (imgButton)
+    imgButton.innerText = 'Creating Image...';
+
+html2canvas(stApp, {{ allowTaint: false, useCORS: true }}).then(function (canvas) {{
+    a = document.createElement('a');
+    a.href = canvas.toDataURL("image/jpeg", 1.0).replace("image/jpeg", "image/octet-stream");
+    a.download = '{pageName}.png';
+    a.click();
+    
+    if (imgButton)
+        imgButton.innerText = '{buttonText}';
+}})
+</script>""",
+            height=0,
+            width=0,
+        )

vectordb_bench/frontend/components/run_test/autoRefresh.py

@@ -1,5 +1,5 @@
 from streamlit_autorefresh import st_autorefresh
-from vectordb_bench.frontend.const import *
+from vectordb_bench.frontend.const.styles import *
 
 
 def autoRefresh():