From 40df7e3424e2dd535e6e3f5c19b287f06af41516 Mon Sep 17 00:00:00 2001 From: YangYanbin Date: Tue, 15 Sep 2026 19:53:21 +0800 Subject: [PATCH] feat: support configurable NQ for concurrent Milvus searches Signed-off-by: YangYanbin --- README.md | 26 ++++++++++ tests/test_case_runner_reuse.py | 24 +++++++++ tests/test_large_topk_frontend.py | 13 +++++ tests/test_milvus.py | 43 +++++++++++++++ tests/test_milvus_zilliz_cli.py | 32 ++++++++++++ vectordb_bench/backend/clients/api.py | 17 ++++++ .../backend/clients/milvus/milvus.py | 16 +++++- vectordb_bench/backend/runner/mp_runner.py | 52 +++++++++++++------ vectordb_bench/backend/task_runner.py | 11 ++++ vectordb_bench/cli/cli.py | 11 ++++ .../frontend/components/check_results/data.py | 2 + vectordb_bench/models.py | 1 + 12 files changed, 229 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index 3fd71fae5..0ffa98f0e 100644 --- a/README.md +++ b/README.md @@ -138,6 +138,32 @@ vectordbbench milvusautoindex \ ``` +Use `--nq` (default: `1`) to set the number of query vectors per concurrent search request. +Milvus and Zilliz Cloud support batching for single-tenant vector performance cases, +including filtered searches. This example reuses a collection containing the selected +dataset and a matching HNSW index: + +```shell +vectordbbench milvushnsw \ + --uri http://localhost:19530 --collection-name VDBBench \ + --case-type Performance768D1M --m 16 --ef-construction 128 --ef-search 128 \ + --k 10 --nq 2 \ + --skip-drop-old --skip-load +``` + +YAML uses `nq: 2`. The value is saved in `task_config.case_config.nq` in result JSON. +Result views distinguish non-default NQ values in case names so different batch sizes stay separate. +Omitting `--num-concurrency` uses the configured concurrency list. Each request contains +exactly `nq` vectors, wrapping through the query dataset as needed. + +- `qps` and `conc_qps_list` count successful query vectors per second. Successful request throughput + is QPS divided by `nq`. Failed batches or responses missing a query's result list contribute zero. +- `conc_latency_*` measures the whole successful batch request. +- `serial_latency_*` and recall use NQ=1 and concurrency=1. + +Other backends and streaming, full-text, cold-latency and multi-tenant cases reject +`nq > 1` before data loading. + Use `--note` or `--note-file` to preserve deployment, resource, client, network, and constraint context in each result JSON under `task_config.db_config.note`. The options are mutually exclusive. Prefer `--note-file` for structured or multiline context, and never include credentials, tokens, or sensitive connection details. ```shell diff --git a/tests/test_case_runner_reuse.py b/tests/test_case_runner_reuse.py index 44315c635..4ea05f521 100644 --- a/tests/test_case_runner_reuse.py +++ b/tests/test_case_runner_reuse.py @@ -1,3 +1,6 @@ +from unittest.mock import Mock + +import pytest from pydantic import SecretStr from vectordb_bench import config @@ -66,6 +69,27 @@ def make_runner( ) +def test_nq_rejects_unsupported_cases_before_data_loading(monkeypatch): + runner = make_runner(db=DB.Test) + runner.config.case_config.nq = 2 + dataset = Mock() + init_db = Mock() + monkeypatch.setattr(runner.ca, "dataset", dataset) + monkeypatch.setattr(CaseRunner, "init_db", init_db) + + with pytest.raises(NotImplementedError, match="nq > 1"): + runner._pre_run() + + assert dataset.mock_calls == [] + init_db.assert_not_called() + + for case_id in (CaseType.FTSBm25Performance, CaseType.CloudMultiTenantSearchCase): + runner = make_runner(db=DB.ZillizCloud, case_id=case_id) + runner.config.case_config.nq = 2 + with pytest.raises(ValueError, match="nq > 1"): + runner._validate_nq() + + def assert_not_reusable(left: CaseRunner, right: CaseRunner) -> None: assert left != right assert right != left diff --git a/tests/test_large_topk_frontend.py b/tests/test_large_topk_frontend.py index 7afd3cba2..7b0adc818 100644 --- a/tests/test_large_topk_frontend.py +++ b/tests/test_large_topk_frontend.py @@ -152,6 +152,17 @@ def test_merge_tasks_keeps_results_with_different_k_separate(): assert len({item["case_name"] for item in merged}) == 2 +def test_results_keep_different_nq_separate(): + tasks = [_case_result(k=100, qps=10), _case_result(k=100, qps=20, nq=2)] + base_name = tasks[0].task_config.case_config.case.name + names = [base_name, f"{base_name} (NQ=2)"] + + assert [data.getCaseResultName(task) for task in tasks] == names + merged, failed = data.mergeTasks(tasks) + assert failed == {} + assert {row["case_name"]: row["qps"] for row in merged} == {names[0]: 10, names[1]: 20} + + def test_merge_tasks_keeps_payload_profiles_separate_for_same_k(): merged, failed = data.mergeTasks( [ @@ -215,6 +226,7 @@ def _case_result( *, k: int, qps: float, + nq: int = 1, payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY, ) -> CaseResult: return CaseResult( @@ -225,6 +237,7 @@ def _case_result( case_config=CaseConfig( case_id=CaseType.Performance768D100M, k=k, + nq=nq, payload_profile=payload_profile, ), ), diff --git a/tests/test_milvus.py b/tests/test_milvus.py index d08faf0a2..83fabf957 100644 --- a/tests/test_milvus.py +++ b/tests/test_milvus.py @@ -4,6 +4,7 @@ """ import logging +from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import MagicMock, call @@ -16,6 +17,7 @@ from vectordb_bench.backend.clients.milvus.config import MilvusConfig, MilvusFtsConfig from vectordb_bench.backend.clients.milvus.milvus import MILVUS_FORCE_MERGE_TARGET_SIZE_MB, Milvus from vectordb_bench.backend.payload import PayloadProfile +from vectordb_bench.backend.runner.mp_runner import MultiProcessingSearchRunner from vectordb_bench.interface import BenchMarkRunner from vectordb_bench.models import CaseConfig, TaskConfig @@ -44,6 +46,47 @@ def search(**kwargs): assert captured["output_fields"] == ["vector"] +def test_milvus_batch_search_counts_successful_queries(monkeypatch): + queries = [[0.0], [1.0], [2.0]] + # Complete zero-hit batch, incomplete batch, SDK failure. + responses = [[[], []], [[]], RuntimeError("search failed")] + db = object.__new__(Milvus) + db.client = SimpleNamespace(search=MagicMock(side_effect=responses)) + db.collection_name = "test_collection" + db._vector_field = "vector" + db._primary_field = "pk" + db.case_config = SimpleNamespace(search_param=lambda: {"metric_type": "COSINE"}) + db.expr = "" + db.init = lambda: nullcontext() + db.prepare_filter = MagicMock() + runner = MultiProcessingSearchRunner(db, queries, nq=2, duration=3) + monkeypatch.setattr("vectordb_bench.backend.runner.mp_runner.random.randint", lambda _a, _b: 0) + # Each SDK request takes one simulated second, regardless of clock reads. + clock = MagicMock() + + def perf_counter(): + assert clock.call_count < 100, "Search loop did not finish" + return float(db.client.search.call_count) + + clock.side_effect = perf_counter + monkeypatch.setattr("vectordb_bench.backend.runner.mp_runner.time", SimpleNamespace(perf_counter=clock)) + + count, _, latencies = runner.search(queries, MagicMock(), MagicMock()) + assert count == 2 + assert latencies == [1.0] + assert db.client.search.call_count == 3 + expected_batches = [queries[:2], [queries[2], queries[0]], queries[1:]] + assert [c.kwargs["data"] for c in db.client.search.call_args_list] == expected_batches + + db.client.search.reset_mock() + db.client.search.side_effect = responses + count, failed_count, stats = runner.search_by_dur(3, queries, MagicMock(), MagicMock()) + assert (count, failed_count) == (2, 4) + assert stats["count"] == 1 + assert stats["avg"] == pytest.approx(1.0, rel=0.01) + assert [c.kwargs["data"] for c in db.client.search.call_args_list] == expected_batches + + def _fake_milvus_client(monkeypatch, *, collection_exists=False, properties=None): client = MagicMock() client.has_collection.return_value = collection_exists diff --git a/tests/test_milvus_zilliz_cli.py b/tests/test_milvus_zilliz_cli.py index 967e5f109..6db46e9c7 100644 --- a/tests/test_milvus_zilliz_cli.py +++ b/tests/test_milvus_zilliz_cli.py @@ -4,6 +4,7 @@ from vectordb_bench.backend.clients.milvus import cli as milvus_cli from vectordb_bench.backend.clients.zilliz_cloud import cli as zilliz_cli +from vectordb_bench.cli import cli as common_cli def test_milvus_cli_builds_shared_connection_config() -> None: @@ -131,3 +132,34 @@ def fake_run(**kwargs): assert result.exit_code == 0, result.output assert captured["db_case_config"].use_partition_key is True + + +def test_milvus_autoindex_cli_nq(monkeypatch: MonkeyPatch) -> None: + captured = {} + + def fake_run(tasks, task_label): + captured["task"] = tasks[0] + + monkeypatch.setattr(common_cli.benchmark_runner, "run", fake_run) + monkeypatch.setattr(common_cli.benchmark_runner, "has_running", lambda: False) + + result = CliRunner().invoke( + milvus_cli.MilvusAutoIndex, + ["--uri", "http://localhost:19530"], + ) + assert result.exit_code == 0, result.output + assert captured["task"].case_config.nq == 1 + + result = CliRunner().invoke( + milvus_cli.MilvusAutoIndex, + ["--uri", "http://localhost:19530", "--nq", "2"], + ) + assert result.exit_code == 0, result.output + assert captured["task"].case_config.nq == 2 + + result = CliRunner().invoke( + milvus_cli.MilvusAutoIndex, + ["--uri", "http://localhost:19530", "--nq", "0", "--dry-run"], + ) + assert result.exit_code == 2, result.output + assert "Invalid value for '--nq'" in result.output diff --git a/vectordb_bench/backend/clients/api.py b/vectordb_bench/backend/clients/api.py index 2365f72aa..e1c74e9df 100644 --- a/vectordb_bench/backend/clients/api.py +++ b/vectordb_bench/backend/clients/api.py @@ -183,6 +183,7 @@ class VectorDB(ABC): "The filtering types supported by the VectorDB Client, default only non-filter" supported_filter_types: list[FilterOp] = [FilterOp.NonFilter] name: str = "" + supports_batch_search: bool = False # Whether the client can share a single connection across threads. # If False, concurrent runners will deep-copy the instance and call @@ -377,6 +378,22 @@ def search_embedding( """ raise NotImplementedError + def search_embeddings( + self, + queries: list[list[float]], + k: int = 100, + payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY, + tenant: str | None = None, + ) -> list[list[int]]: + """Search a batch in one request, returning one ordered result list per query. + + Implementations must raise on incomplete or failed batch responses. Backends + supporting this API must set supports_batch_search=True; sequential calls to + search_embedding are not a batch implementation. + """ + msg = f"{self.name} does not support batch vector search" + raise NotImplementedError(msg) + @abstractmethod def optimize(self, data_size: int | None = None): """optimize will be called between insertion and search in performance cases. diff --git a/vectordb_bench/backend/clients/milvus/milvus.py b/vectordb_bench/backend/clients/milvus/milvus.py index 7c639a977..dcb7e9a5d 100644 --- a/vectordb_bench/backend/clients/milvus/milvus.py +++ b/vectordb_bench/backend/clients/milvus/milvus.py @@ -22,6 +22,7 @@ class Milvus(VectorDB): + supports_batch_search = True supported_filter_types: list[FilterOp] = [ FilterOp.NonFilter, FilterOp.NumGE, @@ -532,6 +533,15 @@ def search_embedding( tenant: str | None = None, ) -> list[int]: """Perform a search on a query embedding and return results.""" + return self.search_embeddings([query], k, payload_profile=payload_profile, tenant=tenant)[0] + + def search_embeddings( + self, + queries: list[list[float]], + k: int = 100, + payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY, + tenant: str | None = None, + ) -> list[list[int]]: assert self.client is not None output_fields = None @@ -548,7 +558,7 @@ def search_embedding( search_kwargs = { "collection_name": self.collection_name, - "data": [query], + "data": queries, "anns_field": self._vector_field, "search_params": self.case_config.search_param(), "limit": k, @@ -556,8 +566,10 @@ def search_embedding( "output_fields": output_fields, } res = self.client.search(**search_kwargs) + if len(res) != len(queries): + raise RuntimeError("Milvus batch search returned an incomplete response") - return [result[self._primary_field] for result in res[0]] + return [[result[self._primary_field] for result in hits] for hits in res] def search_documents( self, diff --git a/vectordb_bench/backend/runner/mp_runner.py b/vectordb_bench/backend/runner/mp_runner.py index 51a5422c6..03e010151 100644 --- a/vectordb_bench/backend/runner/mp_runner.py +++ b/vectordb_bench/backend/runner/mp_runner.py @@ -49,8 +49,14 @@ def __init__( payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY, tenant_labels: list[str] | None = None, workload_kind: WorkloadKind = WorkloadKind.VECTOR, + nq: int = 1, ): self.db = db + if nq < 1: + raise ValueError("nq must be positive") + if nq > 1 and (workload_kind != WorkloadKind.VECTOR or not db.supports_batch_search): + raise NotImplementedError("nq > 1 requires a backend with batch vector search support") + self.nq = nq self.k = k self.filters = filters self.workload_kind = workload_kind @@ -93,16 +99,19 @@ def __setstate__(self, state: dict): msg = f"Unsupported search workload: {self.workload_kind}" raise NotImplementedError(msg) - def _search_embedding(self, emb: list[float], tenant: str | None = None) -> list[int]: + def _search_embedding( + self, emb: list[float] | list[list[float]], tenant: str | None = None + ) -> list[int] | list[list[int]]: + search = self.db.search_embedding if self.nq == 1 else self.db.search_embeddings if tenant is None: if self.payload_profile == PayloadProfile.IDS_ONLY: - return self.db.search_embedding(emb, self.k) - return self.db.search_embedding(emb, self.k, payload_profile=self.payload_profile) + return search(emb, self.k) + return search(emb, self.k, payload_profile=self.payload_profile) if self.payload_profile == PayloadProfile.IDS_ONLY: - return self.db.search_embedding(emb, self.k, tenant=tenant) - return self.db.search_embedding(emb, self.k, payload_profile=self.payload_profile, tenant=tenant) + return search(emb, self.k, tenant=tenant) + return search(emb, self.k, payload_profile=self.payload_profile, tenant=tenant) - def _search_once(self, query: list[float] | str, tenant_rng: random.Random | None = None): + def _search_once(self, query: list[float] | list[list[float]] | str, tenant_rng: random.Random | None = None): if self.workload_kind == WorkloadKind.FULL_TEXT: if self.payload_profile == PayloadProfile.IDS_ONLY: return self._search_func(query, self.k) @@ -112,7 +121,16 @@ def _search_once(self, query: list[float] | str, tenant_rng: random.Random | Non if tenant_rng is not None and self.tenant_labels else None ) - return self._search_func(query, tenant=tenant) + results = self._search_func(query, tenant=tenant) + if self.nq > 1 and len(results) != len(query): + raise RuntimeError("Batch search must return one result list per query") + return results + + def _query_at(self, test_data: list, idx: int): + if self.nq == 1: + return test_data[idx] + # Wrap around to keep every throughput request at exactly nq vectors. + return [test_data[(idx + offset) % len(test_data)] for offset in range(self.nq)] def search( self, @@ -141,14 +159,14 @@ def search( while time.perf_counter() < start_time + self.duration: s = time.perf_counter() try: - self._search_once(test_data[idx], tenant_rng=tenant_rng) - count += 1 + self._search_once(self._query_at(test_data, idx), tenant_rng=tenant_rng) + count += self.nq latencies.append(time.perf_counter() - s) except Exception as e: log.warning(f"VectorDB search_embedding error: {e}") # loop through the test data - idx = idx + 1 if idx < num - 1 else 0 + idx = (idx + self.nq) % num if count % 500 == 0: log.debug( @@ -397,8 +415,8 @@ def search_by_dur( ) -> tuple[int, int, dict]: """ Returns: - int: successful requests count - int: failed requests count + int: successful query vectors count (requests for full-text search) + int: failed query vectors count (requests for full-text search) dict: latency statistics with p99, p95, p50, avg, count (computed via HDR Histogram) """ # sync all process @@ -419,20 +437,20 @@ def search_by_dur( while time.perf_counter() < start_time + dur: s = time.perf_counter() try: - self._search_once(test_data[idx]) - success_count += 1 + self._search_once(self._query_at(test_data, idx)) + success_count += self.nq latency_us = int((time.perf_counter() - s) * US_TO_SECONDS) histogram.record_value(max(HDR_HISTOGRAM_MIN_US, min(latency_us, HDR_HISTOGRAM_MAX_US))) except Exception as e: - failed_cnt += 1 + failed_cnt += self.nq # reduce log - if failed_cnt <= 3: + if failed_cnt <= 3 * self.nq: log.warning(f"VectorDB search_embedding error: {e}") else: log.debug(f"VectorDB search_embedding error: {e}") # loop through the test data - idx = idx + 1 if idx < num - 1 else 0 + idx = (idx + self.nq) % num if success_count % 500 == 0: log.debug( diff --git a/vectordb_bench/backend/task_runner.py b/vectordb_bench/backend/task_runner.py index ea32a0953..6605b4857 100644 --- a/vectordb_bench/backend/task_runner.py +++ b/vectordb_bench/backend/task_runner.py @@ -251,6 +251,7 @@ def _validate_vector_payload_profile(self) -> None: def _pre_run(self, drop_old: bool = True): try: + self._validate_nq() self._validate_cloud_cold_latency_config(drop_old) requested_k = self.config.case_config.k or config.K_DEFAULT ground_truth_k = ( @@ -300,6 +301,15 @@ def _pre_run(self, drop_old: bool = True): log.warning(f"pre run case error: please install client for db: {self.config.db}, error={e}") raise e from None + def _validate_nq(self) -> None: + if self.config.case_config.nq == 1: + return + if self.ca.label != CaseLabel.Performance or self.ca.is_multitenant: + raise ValueError("nq > 1 is only supported for single-tenant vector performance cases") + if not self.config.db.init_cls.supports_batch_search: + msg = f"{self.config.db.value} does not support nq > 1" + raise NotImplementedError(msg) + def _validate_cloud_cold_latency_config(self, drop_old: bool) -> None: if getattr(self.ca, "label", None) != CaseLabel.CloudColdLatency: return @@ -695,6 +705,7 @@ def _init_search_runner(self): payload_profile=self.ca.payload_profile, tenant_labels=tenant_labels, workload_kind=WorkloadKind.VECTOR, + nq=self.config.case_config.nq, ) def _init_fts_search_runner(self): diff --git a/vectordb_bench/cli/cli.py b/vectordb_bench/cli/cli.py index 9f4b5735d..1d405be79 100644 --- a/vectordb_bench/cli/cli.py +++ b/vectordb_bench/cli/cli.py @@ -506,6 +506,16 @@ class CommonTypedDict(TypedDict): help="Number of nearest neighbors. LAION 100M selects tiered GT automatically up to 1,000,000.", ), ] + nq: Annotated[ + int, + click.option( + "--nq", + type=click.IntRange(min=1), + default=1, + show_default=True, + help="Query vectors per concurrent search request; serial latency and recall always use nq=1.", + ), + ] concurrency_duration: Annotated[ int, click.option( @@ -1007,6 +1017,7 @@ def run( case_id=CaseType[parameters["case_type"]], payload_profile=get_case_payload_profile(parameters), k=parameters["k"], + nq=parameters.get("nq", 1), concurrency_search_config=ConcurrencySearchConfig( concurrency_duration=parameters["concurrency_duration"], num_concurrency=[int(s) for s in parameters["num_concurrency"]], diff --git a/vectordb_bench/frontend/components/check_results/data.py b/vectordb_bench/frontend/components/check_results/data.py index 21fa38d76..c7af4d10a 100644 --- a/vectordb_bench/frontend/components/check_results/data.py +++ b/vectordb_bench/frontend/components/check_results/data.py @@ -14,6 +14,8 @@ def getCaseResultName(task: CaseResult) -> str: details = [] if case_config.k is not None and case_config.k != config.K_DEFAULT: details.append(f"K={case_config.k:,}") + if case_config.nq != 1: + details.append(f"NQ={case_config.nq:,}") if ( isinstance(case, PerformanceCase) and case.case_id != CaseType.CloudPayloadSearchCase diff --git a/vectordb_bench/models.py b/vectordb_bench/models.py index 4e471363a..21eea9e95 100644 --- a/vectordb_bench/models.py +++ b/vectordb_bench/models.py @@ -218,6 +218,7 @@ class CaseConfig(BaseModel): custom_case: dict | None = None payload_profile: PayloadProfile | None = None k: int | None = config.K_DEFAULT + nq: PositiveInt = 1 # Concurrent search only; serial latency and recall always use one query per request. concurrency_search_config: ConcurrencySearchConfig = ConcurrencySearchConfig() @field_validator("k")