vectordb-bench 0.0.9__py3-none-any.whl → 0.0.11__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (40) hide show
  1. vectordb_bench/__init__.py +22 -9
  2. vectordb_bench/backend/cases.py +32 -12
  3. vectordb_bench/backend/clients/__init__.py +1 -0
  4. vectordb_bench/backend/clients/api.py +1 -1
  5. vectordb_bench/backend/clients/milvus/cli.py +291 -0
  6. vectordb_bench/backend/clients/milvus/config.py +6 -2
  7. vectordb_bench/backend/clients/milvus/milvus.py +16 -5
  8. vectordb_bench/backend/clients/pgvector/cli.py +116 -0
  9. vectordb_bench/backend/clients/pgvector/config.py +1 -1
  10. vectordb_bench/backend/clients/pgvector/pgvector.py +7 -4
  11. vectordb_bench/backend/clients/redis/cli.py +74 -0
  12. vectordb_bench/backend/clients/test/cli.py +25 -0
  13. vectordb_bench/backend/clients/test/config.py +18 -0
  14. vectordb_bench/backend/clients/test/test.py +62 -0
  15. vectordb_bench/backend/clients/weaviate_cloud/cli.py +41 -0
  16. vectordb_bench/backend/clients/zilliz_cloud/cli.py +55 -0
  17. vectordb_bench/backend/runner/mp_runner.py +14 -3
  18. vectordb_bench/backend/runner/serial_runner.py +7 -3
  19. vectordb_bench/backend/task_runner.py +76 -26
  20. vectordb_bench/cli/__init__.py +0 -0
  21. vectordb_bench/cli/cli.py +362 -0
  22. vectordb_bench/cli/vectordbbench.py +20 -0
  23. vectordb_bench/config-files/sample_config.yml +17 -0
  24. vectordb_bench/frontend/components/check_results/data.py +11 -8
  25. vectordb_bench/frontend/components/concurrent/charts.py +82 -0
  26. vectordb_bench/frontend/components/run_test/dbSelector.py +7 -1
  27. vectordb_bench/frontend/components/run_test/submitTask.py +12 -4
  28. vectordb_bench/frontend/components/tables/data.py +44 -0
  29. vectordb_bench/frontend/const/dbCaseConfigs.py +2 -1
  30. vectordb_bench/frontend/pages/concurrent.py +72 -0
  31. vectordb_bench/frontend/pages/tables.py +24 -0
  32. vectordb_bench/interface.py +21 -25
  33. vectordb_bench/metric.py +23 -1
  34. vectordb_bench/models.py +45 -5
  35. {vectordb_bench-0.0.9.dist-info → vectordb_bench-0.0.11.dist-info}/METADATA +193 -2
  36. {vectordb_bench-0.0.9.dist-info → vectordb_bench-0.0.11.dist-info}/RECORD +40 -24
  37. {vectordb_bench-0.0.9.dist-info → vectordb_bench-0.0.11.dist-info}/WHEEL +1 -1
  38. {vectordb_bench-0.0.9.dist-info → vectordb_bench-0.0.11.dist-info}/entry_points.txt +1 -0
  39. {vectordb_bench-0.0.9.dist-info → vectordb_bench-0.0.11.dist-info}/LICENSE +0 -0
  40. {vectordb_bench-0.0.9.dist-info → vectordb_bench-0.0.11.dist-info}/top_level.txt +0 -0
@@ -1,11 +1,13 @@
1
- import environs
2
1
  import inspect
3
2
  import pathlib
4
- from . import log_util
5
3
 
4
+ import environs
5
+
6
+ from . import log_util
6
7
 
7
8
  env = environs.Env()
8
- env.read_env(".env")
9
+ env.read_env(".env", False)
10
+
9
11
 
10
12
  class config:
11
13
  ALIYUN_OSS_URL = "assets.zilliz.com.cn/benchmark/"
@@ -19,9 +21,20 @@ class config:
19
21
 
20
22
  DROP_OLD = env.bool("DROP_OLD", True)
21
23
  USE_SHUFFLED_DATA = env.bool("USE_SHUFFLED_DATA", True)
22
- NUM_CONCURRENCY = [1, 5, 10, 15, 20, 25, 30, 35]
23
24
 
24
- RESULTS_LOCAL_DIR = pathlib.Path(__file__).parent.joinpath("results")
25
+ NUM_CONCURRENCY = env.list("NUM_CONCURRENCY", [1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75, 80, 85, 90, 95, 100], subcast=int )
26
+
27
+ CONCURRENCY_DURATION = 30
28
+
29
+ RESULTS_LOCAL_DIR = env.path(
30
+ "RESULTS_LOCAL_DIR", pathlib.Path(__file__).parent.joinpath("results")
31
+ )
32
+ CONFIG_LOCAL_DIR = env.path(
33
+ "CONFIG_LOCAL_DIR", pathlib.Path(__file__).parent.joinpath("config-files")
34
+ )
35
+
36
+
37
+ K_DEFAULT = 100 # default return top k nearest neighbors during search
25
38
 
26
39
  CAPACITY_TIMEOUT_IN_SECONDS = 24 * 3600 # 24h
27
40
  LOAD_TIMEOUT_DEFAULT = 2.5 * 3600 # 2.5h
@@ -32,10 +45,10 @@ class config:
32
45
  LOAD_TIMEOUT_1536D_500K = 2.5 * 3600 # 2.5h
33
46
  LOAD_TIMEOUT_1536D_5M = 25 * 3600 # 25h
34
47
 
35
- OPTIMIZE_TIMEOUT_DEFAULT = 15 * 60 # 15min
36
- OPTIMIZE_TIMEOUT_768D_1M = 15 * 60 # 15min
37
- OPTIMIZE_TIMEOUT_768D_10M = 2.5 * 3600 # 2.5h
38
- OPTIMIZE_TIMEOUT_768D_100M = 25 * 3600 # 1.04d
48
+ OPTIMIZE_TIMEOUT_DEFAULT = 30 * 60 # 30min
49
+ OPTIMIZE_TIMEOUT_768D_1M = 30 * 60 # 30min
50
+ OPTIMIZE_TIMEOUT_768D_10M = 5 * 3600 # 5h
51
+ OPTIMIZE_TIMEOUT_768D_100M = 50 * 3600 # 50h
39
52
 
40
53
 
41
54
  OPTIMIZE_TIMEOUT_1536D_500K = 15 * 60 # 15min
@@ -1,6 +1,7 @@
1
1
  import typing
2
2
  import logging
3
3
  from enum import Enum, auto
4
+ from typing import Type
4
5
 
5
6
  from vectordb_bench import config
6
7
  from vectordb_bench.base import BaseModel
@@ -10,8 +11,6 @@ from .dataset import Dataset, DatasetManager
10
11
 
11
12
  log = logging.getLogger(__name__)
12
13
 
13
- Case = typing.TypeVar("Case")
14
-
15
14
 
16
15
  class CaseType(Enum):
17
16
  """
@@ -42,11 +41,15 @@ class CaseType(Enum):
42
41
  Performance1536D500K99P = 14
43
42
  Performance1536D5M99P = 15
44
43
 
44
+ Performance1536D50K = 50
45
+
45
46
  Custom = 100
46
47
 
47
48
  @property
48
- def case_cls(self, custom_configs: dict | None = None) -> Case:
49
- return type2case.get(self)
49
+ def case_cls(self, custom_configs: dict | None = None) -> Type["Case"]:
50
+ if self not in type2case:
51
+ raise NotImplementedError(f"Case {self} has not implemented. You can add it manually to vectordb_bench.backend.cases.type2case or define a custom_configs['custom_cls']")
52
+ return type2case[self]
50
53
 
51
54
  @property
52
55
  def case_name(self) -> str:
@@ -69,7 +72,7 @@ class CaseLabel(Enum):
69
72
 
70
73
 
71
74
  class Case(BaseModel):
72
- """Undifined case
75
+ """Undefined case
73
76
 
74
77
  Fields:
75
78
  case_id(CaseType): default 9 case type plus one custom cases.
@@ -86,9 +89,9 @@ class Case(BaseModel):
86
89
  dataset: DatasetManager
87
90
 
88
91
  load_timeout: float | int
89
- optimize_timeout: float | int | None
92
+ optimize_timeout: float | int | None = None
90
93
 
91
- filter_rate: float | None
94
+ filter_rate: float | None = None
92
95
 
93
96
  @property
94
97
  def filters(self) -> dict | None:
@@ -115,20 +118,23 @@ class PerformanceCase(Case, BaseModel):
115
118
  load_timeout: float | int = config.LOAD_TIMEOUT_DEFAULT
116
119
  optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_DEFAULT
117
120
 
121
+
118
122
  class CapacityDim960(CapacityCase):
119
123
  case_id: CaseType = CaseType.CapacityDim960
120
124
  dataset: DatasetManager = Dataset.GIST.manager(100_000)
121
125
  name: str = "Capacity Test (960 Dim Repeated)"
122
- description: str = """This case tests the vector database's loading capacity by repeatedly inserting large-dimension vectors (GIST 100K vectors, <b>960 dimensions</b>) until it is fully loaded.
123
- Number of inserted vectors will be reported."""
126
+ description: str = """This case tests the vector database's loading capacity by repeatedly inserting large-dimension
127
+ vectors (GIST 100K vectors, <b>960 dimensions</b>) until it is fully loaded. Number of inserted vectors will be
128
+ reported."""
124
129
 
125
130
 
126
131
  class CapacityDim128(CapacityCase):
127
132
  case_id: CaseType = CaseType.CapacityDim128
128
133
  dataset: DatasetManager = Dataset.SIFT.manager(500_000)
129
134
  name: str = "Capacity Test (128 Dim Repeated)"
130
- description: str = """This case tests the vector database's loading capacity by repeatedly inserting small-dimension vectors (SIFT 100K vectors, <b>128 dimensions</b>) until it is fully loaded.
131
- Number of inserted vectors will be reported."""
135
+ description: str = """This case tests the vector database's loading capacity by repeatedly inserting small-dimension
136
+ vectors (SIFT 100K vectors, <b>128 dimensions</b>) until it is fully loaded. Number of inserted vectors will be
137
+ reported."""
132
138
 
133
139
 
134
140
  class Performance768D10M(PerformanceCase):
@@ -238,6 +244,7 @@ Results will show index building time, recall, and maximum QPS."""
238
244
  load_timeout: float | int = config.LOAD_TIMEOUT_1536D_500K
239
245
  optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_500K
240
246
 
247
+
241
248
  class Performance1536D5M1P(PerformanceCase):
242
249
  case_id: CaseType = CaseType.Performance1536D5M1P
243
250
  filter_rate: float | int | None = 0.01
@@ -248,6 +255,7 @@ Results will show index building time, recall, and maximum QPS."""
248
255
  load_timeout: float | int = config.LOAD_TIMEOUT_1536D_5M
249
256
  optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_5M
250
257
 
258
+
251
259
  class Performance1536D500K99P(PerformanceCase):
252
260
  case_id: CaseType = CaseType.Performance1536D500K99P
253
261
  filter_rate: float | int | None = 0.99
@@ -258,6 +266,7 @@ Results will show index building time, recall, and maximum QPS."""
258
266
  load_timeout: float | int = config.LOAD_TIMEOUT_1536D_500K
259
267
  optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_500K
260
268
 
269
+
261
270
  class Performance1536D5M99P(PerformanceCase):
262
271
  case_id: CaseType = CaseType.Performance1536D5M99P
263
272
  filter_rate: float | int | None = 0.99
@@ -269,6 +278,17 @@ Results will show index building time, recall, and maximum QPS."""
269
278
  optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_5M
270
279
 
271
280
 
281
+ class Performance1536D50K(PerformanceCase):
282
+ case_id: CaseType = CaseType.Performance1536D50K
283
+ filter_rate: float | int | None = None
284
+ dataset: DatasetManager = Dataset.OPENAI.manager(50_000)
285
+ name: str = "Search Performance Test (50K Dataset, 1536 Dim)"
286
+ description: str = """This case tests the search performance of a vector database with a medium 50K dataset (<b>OpenAI 50K vectors</b>, 1536 dimensions), at varying parallel levels.
287
+ Results will show index building time, recall, and maximum QPS."""
288
+ load_timeout: float | int = 3600
289
+ optimize_timeout: float | int | None = 15 * 60
290
+
291
+
272
292
  type2case = {
273
293
  CaseType.CapacityDim960: CapacityDim960,
274
294
  CaseType.CapacityDim128: CapacityDim128,
@@ -290,5 +310,5 @@ type2case = {
290
310
 
291
311
  CaseType.Performance1536D500K99P: Performance1536D500K99P,
292
312
  CaseType.Performance1536D5M99P: Performance1536D5M99P,
293
-
313
+ CaseType.Performance1536D50K: Performance1536D50K,
294
314
  }
@@ -32,6 +32,7 @@ class DB(Enum):
32
32
  PgVectoRS = "PgVectoRS"
33
33
  Redis = "Redis"
34
34
  Chroma = "Chroma"
35
+ Test = "test"
35
36
 
36
37
 
37
38
  @property
@@ -47,7 +47,7 @@ class DBConfig(ABC, BaseModel):
47
47
  def not_empty_field(cls, v, field):
48
48
  if field.name == "db_label":
49
49
  return v
50
- if isinstance(v, (str, SecretStr)) and len(v) == 0:
50
+ if not v and isinstance(v, (str, SecretStr)):
51
51
  raise ValueError("Empty string!")
52
52
  return v
53
53
 
@@ -0,0 +1,291 @@
1
+ from typing import Annotated, TypedDict, Unpack
2
+
3
+ import click
4
+ from pydantic import SecretStr
5
+
6
+ from vectordb_bench.cli.cli import (
7
+ CommonTypedDict,
8
+ HNSWFlavor3,
9
+ IVFFlatTypedDictN,
10
+ cli,
11
+ click_parameter_decorators_from_typed_dict,
12
+ run,
13
+
14
+ )
15
+ from vectordb_bench.backend.clients import DB
16
+
17
+ DBTYPE = DB.Milvus
18
+
19
+
20
+ class MilvusTypedDict(TypedDict):
21
+ uri: Annotated[
22
+ str, click.option("--uri", type=str, help="uri connection string", required=True)
23
+ ]
24
+
25
+
26
+ class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict):
27
+ ...
28
+
29
+
30
+ @cli.command()
31
+ @click_parameter_decorators_from_typed_dict(MilvusAutoIndexTypedDict)
32
+ def MilvusAutoIndex(**parameters: Unpack[MilvusAutoIndexTypedDict]):
33
+ from .config import MilvusConfig, AutoIndexConfig
34
+
35
+ run(
36
+ db=DBTYPE,
37
+ db_config=MilvusConfig(
38
+ db_label=parameters["db_label"],
39
+ uri=SecretStr(parameters["uri"]),
40
+ ),
41
+ db_case_config=AutoIndexConfig(),
42
+ **parameters,
43
+ )
44
+
45
+
46
+ @cli.command()
47
+ @click_parameter_decorators_from_typed_dict(MilvusAutoIndexTypedDict)
48
+ def MilvusFlat(**parameters: Unpack[MilvusAutoIndexTypedDict]):
49
+ from .config import MilvusConfig, FLATConfig
50
+
51
+ run(
52
+ db=DBTYPE,
53
+ db_config=MilvusConfig(
54
+ db_label=parameters["db_label"],
55
+ uri=SecretStr(parameters["uri"]),
56
+ ),
57
+ db_case_config=FLATConfig(),
58
+ **parameters,
59
+ )
60
+
61
+
62
+ class MilvusHNSWTypedDict(CommonTypedDict, MilvusTypedDict, HNSWFlavor3):
63
+ ...
64
+
65
+
66
+ @cli.command()
67
+ @click_parameter_decorators_from_typed_dict(MilvusHNSWTypedDict)
68
+ def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]):
69
+ from .config import MilvusConfig, HNSWConfig
70
+
71
+ run(
72
+ db=DBTYPE,
73
+ db_config=MilvusConfig(
74
+ db_label=parameters["db_label"],
75
+ uri=SecretStr(parameters["uri"]),
76
+ ),
77
+ db_case_config=HNSWConfig(
78
+ M=parameters["m"],
79
+ efConstruction=parameters["ef_construction"],
80
+ ef=parameters["ef_search"],
81
+ ),
82
+ **parameters,
83
+ )
84
+
85
+
86
+ class MilvusIVFFlatTypedDict(CommonTypedDict, MilvusTypedDict, IVFFlatTypedDictN):
87
+ ...
88
+
89
+
90
+ @cli.command()
91
+ @click_parameter_decorators_from_typed_dict(MilvusIVFFlatTypedDict)
92
+ def MilvusIVFFlat(**parameters: Unpack[MilvusIVFFlatTypedDict]):
93
+ from .config import MilvusConfig, IVFFlatConfig
94
+
95
+ run(
96
+ db=DBTYPE,
97
+ db_config=MilvusConfig(
98
+ db_label=parameters["db_label"],
99
+ uri=SecretStr(parameters["uri"]),
100
+ ),
101
+ db_case_config=IVFFlatConfig(
102
+ nlist=parameters["nlist"],
103
+ nprobe=parameters["nprobe"],
104
+ ),
105
+ **parameters,
106
+ )
107
+
108
+
109
+ @cli.command()
110
+ @click_parameter_decorators_from_typed_dict(MilvusIVFFlatTypedDict)
111
+ def MilvusIVFSQ8(**parameters: Unpack[MilvusIVFFlatTypedDict]):
112
+ from .config import MilvusConfig, IVFSQ8Config
113
+
114
+ run(
115
+ db=DBTYPE,
116
+ db_config=MilvusConfig(
117
+ db_label=parameters["db_label"],
118
+ uri=SecretStr(parameters["uri"]),
119
+ ),
120
+ db_case_config=IVFSQ8Config(
121
+ nlist=parameters["nlist"],
122
+ nprobe=parameters["nprobe"],
123
+ ),
124
+ **parameters,
125
+ )
126
+
127
+
128
+ class MilvusDISKANNTypedDict(CommonTypedDict, MilvusTypedDict):
129
+ search_list: Annotated[
130
+ str, click.option("--search-list",
131
+ type=int,
132
+ required=True)
133
+ ]
134
+
135
+
136
+ @cli.command()
137
+ @click_parameter_decorators_from_typed_dict(MilvusDISKANNTypedDict)
138
+ def MilvusDISKANN(**parameters: Unpack[MilvusDISKANNTypedDict]):
139
+ from .config import MilvusConfig, DISKANNConfig
140
+
141
+ run(
142
+ db=DBTYPE,
143
+ db_config=MilvusConfig(
144
+ db_label=parameters["db_label"],
145
+ uri=SecretStr(parameters["uri"]),
146
+ ),
147
+ db_case_config=DISKANNConfig(
148
+ search_list=parameters["search_list"],
149
+ ),
150
+ **parameters,
151
+ )
152
+
153
+
154
+ class MilvusGPUIVFTypedDict(CommonTypedDict, MilvusTypedDict, MilvusIVFFlatTypedDict):
155
+ cache_dataset_on_device: Annotated[
156
+ str, click.option("--cache-dataset-on-device",
157
+ type=str,
158
+ required=True)
159
+ ]
160
+ refine_ratio: Annotated[
161
+ str, click.option("--refine-ratio",
162
+ type=float,
163
+ required=True)
164
+ ]
165
+
166
+
167
+ @cli.command()
168
+ @click_parameter_decorators_from_typed_dict(MilvusGPUIVFTypedDict)
169
+ def MilvusGPUIVFFlat(**parameters: Unpack[MilvusGPUIVFTypedDict]):
170
+ from .config import MilvusConfig, GPUIVFFlatConfig
171
+
172
+ run(
173
+ db=DBTYPE,
174
+ db_config=MilvusConfig(
175
+ db_label=parameters["db_label"],
176
+ uri=SecretStr(parameters["uri"]),
177
+ ),
178
+ db_case_config=GPUIVFFlatConfig(
179
+ nlist=parameters["nlist"],
180
+ nprobe=parameters["nprobe"],
181
+ cache_dataset_on_device=parameters["cache_dataset_on_device"],
182
+ refine_ratio=parameters.get("refine_ratio"),
183
+ ),
184
+ **parameters,
185
+ )
186
+
187
+
188
+ class MilvusGPUIVFPQTypedDict(CommonTypedDict, MilvusTypedDict, MilvusIVFFlatTypedDict, MilvusGPUIVFTypedDict):
189
+ m: Annotated[
190
+ str, click.option("--m",
191
+ type=int, help="hnsw m",
192
+ required=True)
193
+ ]
194
+ nbits: Annotated[
195
+ str, click.option("--nbits",
196
+ type=int,
197
+ required=True)
198
+ ]
199
+
200
+
201
+ @cli.command()
202
+ @click_parameter_decorators_from_typed_dict(MilvusGPUIVFPQTypedDict)
203
+ def MilvusGPUIVFPQ(**parameters: Unpack[MilvusGPUIVFPQTypedDict]):
204
+ from .config import MilvusConfig, GPUIVFPQConfig
205
+
206
+ run(
207
+ db=DBTYPE,
208
+ db_config=MilvusConfig(
209
+ db_label=parameters["db_label"],
210
+ uri=SecretStr(parameters["uri"]),
211
+ ),
212
+ db_case_config=GPUIVFPQConfig(
213
+ nlist=parameters["nlist"],
214
+ nprobe=parameters["nprobe"],
215
+ m=parameters["m"],
216
+ nbits=parameters["nbits"],
217
+ cache_dataset_on_device=parameters["cache_dataset_on_device"],
218
+ refine_ratio=parameters["refine_ratio"],
219
+ ),
220
+ **parameters,
221
+ )
222
+
223
+
224
+ class MilvusGPUCAGRATypedDict(CommonTypedDict, MilvusTypedDict, MilvusGPUIVFTypedDict):
225
+ intermediate_graph_degree: Annotated[
226
+ str, click.option("--intermediate-graph-degree",
227
+ type=int,
228
+ required=True)
229
+ ]
230
+ graph_degree: Annotated[
231
+ str, click.option("--graph-degree",
232
+ type=int,
233
+ required=True)
234
+ ]
235
+ build_algo: Annotated[
236
+ str, click.option("--build_algo",
237
+ type=str,
238
+ required=True)
239
+ ]
240
+ team_size: Annotated[
241
+ str, click.option("--team-size",
242
+ type=int,
243
+ required=True)
244
+ ]
245
+ search_width: Annotated[
246
+ str, click.option("--search-width",
247
+ type=int,
248
+ required=True)
249
+ ]
250
+ itopk_size: Annotated[
251
+ str, click.option("--itopk-size",
252
+ type=int,
253
+ required=True)
254
+ ]
255
+ min_iterations: Annotated[
256
+ str, click.option("--min-iterations",
257
+ type=int,
258
+ required=True)
259
+ ]
260
+ max_iterations: Annotated[
261
+ str, click.option("--max-iterations",
262
+ type=int,
263
+ required=True)
264
+ ]
265
+
266
+
267
+ @cli.command()
268
+ @click_parameter_decorators_from_typed_dict(MilvusGPUCAGRATypedDict)
269
+ def MilvusGPUCAGRA(**parameters: Unpack[MilvusGPUCAGRATypedDict]):
270
+ from .config import MilvusConfig, GPUCAGRAConfig
271
+
272
+ run(
273
+ db=DBTYPE,
274
+ db_config=MilvusConfig(
275
+ db_label=parameters["db_label"],
276
+ uri=SecretStr(parameters["uri"]),
277
+ ),
278
+ db_case_config=GPUCAGRAConfig(
279
+ intermediate_graph_degree=parameters["intermediate_graph_degree"],
280
+ graph_degree=parameters["graph_degree"],
281
+ itopk_size=parameters["itopk_size"],
282
+ team_size=parameters["team_size"],
283
+ search_width=parameters["search_width"],
284
+ min_iterations=parameters["min_iterations"],
285
+ max_iterations=parameters["max_iterations"],
286
+ build_algo=parameters["build_algo"],
287
+ cache_dataset_on_device=parameters["cache_dataset_on_device"],
288
+ refine_ratio=parameters["refine_ratio"],
289
+ ),
290
+ **parameters,
291
+ )
@@ -14,13 +14,17 @@ class MilvusIndexConfig(BaseModel):
14
14
 
15
15
  index: IndexType
16
16
  metric_type: MetricType | None = None
17
+
18
+ @property
19
+ def is_gpu_index(self) -> bool:
20
+ return self.index in [IndexType.GPU_CAGRA, IndexType.GPU_IVF_FLAT, IndexType.GPU_IVF_PQ]
17
21
 
18
22
  def parse_metric(self) -> str:
19
23
  if not self.metric_type:
20
24
  return ""
21
25
 
22
- # if self.metric_type == MetricType.COSINE:
23
- # return MetricType.L2.value
26
+ if self.is_gpu_index and self.metric_type == MetricType.COSINE:
27
+ return MetricType.L2.value
24
28
  return self.metric_type.value
25
29
 
26
30
 
@@ -119,13 +119,21 @@ class Milvus(VectorDB):
119
119
  wait_index()
120
120
 
121
121
  # Skip compaction if use GPU indexType
122
- if self.case_config.index in [IndexType.GPU_CAGRA, IndexType.GPU_IVF_FLAT, IndexType.GPU_IVF_PQ]:
122
+ if self.case_config.is_gpu_index:
123
123
  log.debug("skip compaction for gpu index type.")
124
124
  else :
125
- self.col.compact()
126
- self.col.wait_for_compaction_completed()
125
+ try:
126
+ self.col.compact()
127
+ self.col.wait_for_compaction_completed()
128
+ except Exception as e:
129
+ log.warning(f"{self.name} compact error: {e}")
130
+ if hasattr(e, 'code'):
131
+ if e.code().name == 'PERMISSION_DENIED':
132
+ log.warning(f"Skip compact due to permission denied.")
133
+ pass
134
+ else:
135
+ raise e
127
136
  wait_index()
128
-
129
137
  except Exception as e:
130
138
  log.warning(f"{self.name} optimize error: {e}")
131
139
  raise e from None
@@ -143,7 +151,6 @@ class Milvus(VectorDB):
143
151
  self.case_config.index_param(),
144
152
  index_name=self._index_name,
145
153
  )
146
-
147
154
  coll.load()
148
155
  log.info(f"{self.name} load")
149
156
  except Exception as e:
@@ -157,6 +164,10 @@ class Milvus(VectorDB):
157
164
 
158
165
  def need_normalize_cosine(self) -> bool:
159
166
  """Wheather this database need to normalize dataset to support COSINE"""
167
+ if self.case_config.is_gpu_index:
168
+ log.info(f"current gpu_index only supports IP / L2, cosine dataset need normalize.")
169
+ return True
170
+
160
171
  return False
161
172
 
162
173
  def insert_embeddings(
@@ -0,0 +1,116 @@
1
+ from typing import Annotated, Optional, TypedDict, Unpack
2
+
3
+ import click
4
+ import os
5
+ from pydantic import SecretStr
6
+
7
+ from ....cli.cli import (
8
+ CommonTypedDict,
9
+ HNSWFlavor1,
10
+ IVFFlatTypedDict,
11
+ cli,
12
+ click_parameter_decorators_from_typed_dict,
13
+ run,
14
+ )
15
+ from vectordb_bench.backend.clients import DB
16
+
17
+
18
+ class PgVectorTypedDict(CommonTypedDict):
19
+ user_name: Annotated[
20
+ str, click.option("--user-name", type=str, help="Db username", required=True)
21
+ ]
22
+ password: Annotated[
23
+ str,
24
+ click.option("--password",
25
+ type=str,
26
+ help="Postgres database password",
27
+ default=lambda: os.environ.get("POSTGRES_PASSWORD", ""),
28
+ show_default="$POSTGRES_PASSWORD",
29
+ ),
30
+ ]
31
+
32
+ host: Annotated[
33
+ str, click.option("--host", type=str, help="Db host", required=True)
34
+ ]
35
+ db_name: Annotated[
36
+ str, click.option("--db-name", type=str, help="Db name", required=True)
37
+ ]
38
+ maintenance_work_mem: Annotated[
39
+ Optional[str],
40
+ click.option(
41
+ "--maintenance-work-mem",
42
+ type=str,
43
+ help="Sets the maximum memory to be used for maintenance operations (index creation). "
44
+ "Can be entered as string with unit like '64GB' or as an integer number of KB."
45
+ "This will set the parameters: max_parallel_maintenance_workers,"
46
+ " max_parallel_workers & table(parallel_workers)",
47
+ required=False,
48
+ ),
49
+ ]
50
+ max_parallel_workers: Annotated[
51
+ Optional[int],
52
+ click.option(
53
+ "--max-parallel-workers",
54
+ type=int,
55
+ help="Sets the maximum number of parallel processes per maintenance operation (index creation)",
56
+ required=False,
57
+ ),
58
+ ]
59
+
60
+
61
+ class PgVectorIVFFlatTypedDict(PgVectorTypedDict, IVFFlatTypedDict):
62
+ ...
63
+
64
+
65
+ @cli.command()
66
+ @click_parameter_decorators_from_typed_dict(PgVectorIVFFlatTypedDict)
67
+ def PgVectorIVFFlat(
68
+ **parameters: Unpack[PgVectorIVFFlatTypedDict],
69
+ ):
70
+ from .config import PgVectorConfig, PgVectorIVFFlatConfig
71
+
72
+ run(
73
+ db=DB.PgVector,
74
+ db_config=PgVectorConfig(
75
+ db_label=parameters["db_label"],
76
+ user_name=SecretStr(parameters["user_name"]),
77
+ password=SecretStr(parameters["password"]),
78
+ host=parameters["host"],
79
+ db_name=parameters["db_name"],
80
+ ),
81
+ db_case_config=PgVectorIVFFlatConfig(
82
+ metric_type=None, lists=parameters["lists"], probes=parameters["probes"]
83
+ ),
84
+ **parameters,
85
+ )
86
+
87
+
88
+ class PgVectorHNSWTypedDict(PgVectorTypedDict, HNSWFlavor1):
89
+ ...
90
+
91
+
92
+ @cli.command()
93
+ @click_parameter_decorators_from_typed_dict(PgVectorHNSWTypedDict)
94
+ def PgVectorHNSW(
95
+ **parameters: Unpack[PgVectorHNSWTypedDict],
96
+ ):
97
+ from .config import PgVectorConfig, PgVectorHNSWConfig
98
+
99
+ run(
100
+ db=DB.PgVector,
101
+ db_config=PgVectorConfig(
102
+ db_label=parameters["db_label"],
103
+ user_name=SecretStr(parameters["user_name"]),
104
+ password=SecretStr(parameters["password"]),
105
+ host=parameters["host"],
106
+ db_name=parameters["db_name"],
107
+ ),
108
+ db_case_config=PgVectorHNSWConfig(
109
+ m=parameters["m"],
110
+ ef_construction=parameters["ef_construction"],
111
+ ef_search=parameters["ef_search"],
112
+ maintenance_work_mem=parameters["maintenance_work_mem"],
113
+ max_parallel_workers=parameters["max_parallel_workers"],
114
+ ),
115
+ **parameters,
116
+ )
@@ -109,7 +109,7 @@ class PgVectorIndexConfig(BaseModel, DBCaseConfig):
109
109
  def _optionally_build_set_options(
110
110
  set_mapping: Mapping[str, Any]
111
111
  ) -> Sequence[dict[str, Any]]:
112
- """Walk through options, creating 'SET 'key1 = "value1";' commands"""
112
+ """Walk through options, creating 'SET 'key1 = "value1";' list"""
113
113
  session_options = []
114
114
  for setting_name, value in set_mapping.items():
115
115
  if value: