From 8d144e22f136215426c4a01016e39baa5cf5c74d Mon Sep 17 00:00:00 2001 From: James Gao Date: Tue, 30 Jun 2026 03:47:04 +0000 Subject: [PATCH 1/2] feat: add OSS OpenSearch FTS support Signed-off-by: jamesgao-jpg --- tests/test_oss_opensearch_fts.py | 180 ++++++++++++++++++ vectordb_bench/backend/clients/__init__.py | 4 + .../backend/clients/oss_opensearch/config.py | 34 ++++ .../clients/oss_opensearch/oss_opensearch.py | 135 ++++++++++++- .../frontend/config/dbCaseConfigs.py | 2 + 5 files changed, 351 insertions(+), 4 deletions(-) create mode 100644 tests/test_oss_opensearch_fts.py diff --git a/tests/test_oss_opensearch_fts.py b/tests/test_oss_opensearch_fts.py new file mode 100644 index 000000000..66480ea56 --- /dev/null +++ b/tests/test_oss_opensearch_fts.py @@ -0,0 +1,180 @@ +import sys +import types + +import pytest + + +class _FakeOpenSearch: + pass + + +sys.modules.setdefault("opensearchpy", types.SimpleNamespace(OpenSearch=_FakeOpenSearch)) + +from vectordb_bench.backend.clients import DB +from vectordb_bench.backend.clients.api import IndexType +from vectordb_bench.backend.clients.oss_opensearch.config import OSSOpenSearchFtsConfig +from vectordb_bench.backend.clients.oss_opensearch.oss_opensearch import OSSOpenSearch +from vectordb_bench.backend.payload import PayloadProfile + + +def make_fts_db(): + db = OSSOpenSearch.__new__(OSSOpenSearch) + db.index_name = "idx" + db.id_col_name = "doc_id" + db.text_col_name = "text" + db._is_fts = True + db.client = object() + return db + + +def test_oss_opensearch_fts_config_defaults(): + config = OSSOpenSearchFtsConfig() + + assert config.index_param()["properties"]["doc_id"] == {"type": "keyword"} + assert config.index_param()["properties"]["text"] == {"type": "text"} + assert config.search_param() == {} + + +def test_oss_opensearch_fts_config_supports_bm25_similarity(): + config = OSSOpenSearchFtsConfig(bm25_k1=1.2, bm25_b=0.75) + + assert config.index_param()["properties"]["text"]["similarity"] == "vdbbench_bm25" + assert config.similarity_settings() == { + "similarity": { + "vdbbench_bm25": { + "type": "BM25", + "k1": 1.2, + "b": 0.75, + } + } + } + + +def test_oss_opensearch_declares_full_text_support(): + assert OSSOpenSearch.supports_full_text_search() is True + assert DB.OSSOpenSearch.case_config_cls(IndexType.FTS) is OSSOpenSearchFtsConfig + + +def test_oss_opensearch_create_index_fts_uses_text_mappings_and_settings(): + db = OSSOpenSearch.__new__(OSSOpenSearch) + db._is_fts = True + db.case_config = OSSOpenSearchFtsConfig( + number_of_shards=2, + number_of_replicas=1, + refresh_interval="10s", + ) + db.index_name = "idx" + calls = {} + + class Indices: + def create(self, **kwargs): + calls.update(kwargs) + + class Client: + indices = Indices() + + db._create_index(Client()) + + assert calls == { + "index": "idx", + "body": { + "settings": { + "index": { + "number_of_shards": 2, + "number_of_replicas": 1, + "refresh_interval": "10s", + } + }, + "mappings": { + "properties": { + "doc_id": {"type": "keyword"}, + "text": {"type": "text"}, + } + }, + }, + } + + +def test_oss_opensearch_insert_documents_builds_bulk_body(): + db = make_fts_db() + captured = {} + + class Client: + def bulk(self, **kwargs): + captured.update(kwargs) + return {"errors": False} + + db.client = Client() + + assert db.insert_documents(["alpha", "beta"], ["d1", "d2"]) == (2, None) + assert captured["body"] == [ + {"index": {"_index": "idx", "_id": "d1"}}, + {"doc_id": "d1", "text": "alpha"}, + {"index": {"_index": "idx", "_id": "d2"}}, + {"doc_id": "d2", "text": "beta"}, + ] + + +def test_oss_opensearch_insert_documents_validates_lengths(): + db = make_fts_db() + + class Client: + def bulk(self, **kwargs): + raise AssertionError("bulk should not be called") + + db.client = Client() + + with pytest.raises(ValueError, match="Mismatch between texts .* and doc_ids .* lengths"): + db.insert_documents(["alpha", "beta"], ["d1"]) + + +def test_oss_opensearch_search_documents_builds_match_query(): + db = make_fts_db() + calls = {} + + class Client: + def search(self, **kwargs): + calls.update(kwargs) + return {"hits": {"hits": [{"fields": {"doc_id": ["d1"]}}]}} + + db.client = Client() + + assert db.search_documents("hello world", k=3) == ["d1"] + assert calls["index"] == "idx" + assert calls["body"] == {"query": {"match": {"text": "hello world"}}} + assert calls["size"] == 3 + assert calls["stored_fields"] == "_none_" + assert calls["filter_path"] == ["hits.hits._id", "hits.hits.fields.doc_id"] + + +def test_oss_opensearch_search_documents_requests_text_payload(): + db = make_fts_db() + calls = {} + + class Client: + def search(self, **kwargs): + calls.update(kwargs) + return {"hits": {"hits": [{"_id": "d1", "_source": {"text": "hello"}}]}} + + db.client = Client() + + assert db.search_documents("hello world", k=3, payload_profile=PayloadProfile.TEXT) == ["d1"] + assert calls["_source"] == ["text"] + assert "stored_fields" not in calls + assert calls["filter_path"] == [ + "hits.hits._id", + "hits.hits.fields.doc_id", + "hits.hits._source.text", + ] + + +def test_oss_opensearch_document_methods_guard_non_fts_mode(): + db = OSSOpenSearch.__new__(OSSOpenSearch) + db._is_fts = False + db.client = object() + + with pytest.raises(RuntimeError, match="OSSOpenSearch full-text insert requires OSSOpenSearchFtsConfig"): + db.insert_documents(["alpha"], ["d1"]) + + with pytest.raises(RuntimeError, match="OSSOpenSearch full-text search requires OSSOpenSearchFtsConfig"): + db.search_documents("alpha") diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index 4fd50871c..dd06a912e 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/vectordb_bench/backend/clients/__init__.py @@ -538,6 +538,10 @@ def case_config_cls( # noqa: C901, PLR0911, PLR0912, PLR0915 return AWSOpenSearchIndexConfig if self == DB.OSSOpenSearch: + if index_type == IndexType.FTS: + from .oss_opensearch.config import OSSOpenSearchFtsConfig + + return OSSOpenSearchFtsConfig from .oss_opensearch.config import OSSOpenSearchIndexConfig return OSSOpenSearchIndexConfig diff --git a/vectordb_bench/backend/clients/oss_opensearch/config.py b/vectordb_bench/backend/clients/oss_opensearch/config.py index 7a8b1d98e..5b843b19c 100644 --- a/vectordb_bench/backend/clients/oss_opensearch/config.py +++ b/vectordb_bench/backend/clients/oss_opensearch/config.py @@ -247,3 +247,37 @@ def index_param(self) -> dict: def search_param(self) -> dict: return {"ef_search": self.efSearch} + + +class OSSOpenSearchFtsConfig(BaseModel, DBCaseConfig): + number_of_shards: int = 1 + number_of_replicas: int = 0 + refresh_interval: str = "30s" + force_merge_enabled: bool = True + metric_type: MetricType = MetricType.BM25 + bm25_k1: float | None = None + bm25_b: float | None = None + + def index_param(self) -> dict: + text_mapping = {"type": "text"} + if self.bm25_k1 is not None or self.bm25_b is not None: + text_mapping["similarity"] = "vdbbench_bm25" + return { + "properties": { + "doc_id": {"type": "keyword"}, + "text": text_mapping, + }, + } + + def search_param(self) -> dict: + return {} + + def similarity_settings(self) -> dict: + if self.bm25_k1 is None and self.bm25_b is None: + return {} + bm25_settings = {"type": "BM25"} + if self.bm25_k1 is not None: + bm25_settings["k1"] = self.bm25_k1 + if self.bm25_b is not None: + bm25_settings["b"] = self.bm25_b + return {"similarity": {"vdbbench_bm25": bm25_settings}} diff --git a/vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py b/vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py index f71850a17..d29abe743 100644 --- a/vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py +++ b/vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py @@ -9,9 +9,10 @@ from packaging.version import parse as parse_version from vectordb_bench.backend.filter import Filter, FilterOp +from vectordb_bench.backend.payload import PayloadProfile from ..api import VectorDB -from .config import OSSOpenSearchIndexConfig, OSSOS_Engine +from .config import OSSOpenSearchFtsConfig, OSSOpenSearchIndexConfig, OSSOS_Engine log = logging.getLogger(__name__) @@ -194,7 +195,7 @@ def __init__( self, dim: int, db_config: dict[str, Any], - db_case_config: OSSOpenSearchIndexConfig, + db_case_config: OSSOpenSearchIndexConfig | OSSOpenSearchFtsConfig, index_name: str = "vdb_bench_index", # must be lowercase id_col_name: str = "_id", label_col_name: str = "label", @@ -212,6 +213,10 @@ def __init__( self.label_col_name = label_col_name self.vector_col_name = vector_col_name self.with_scalar_labels = with_scalar_labels + self._is_fts = isinstance(db_case_config, OSSOpenSearchFtsConfig) + self.text_col_name = "text" + if self._is_fts: + self.id_col_name = "doc_id" # Initialize client state self.client: OpenSearch | None = None @@ -236,13 +241,21 @@ def _handle_index_initialization(self, client: OpenSearch, drop_old: bool) -> No if not is_existed: self._create_index(client) log.info(f"OSS_OpenSearch client create index: {self.index_name}") - self._update_ef_search_before_search(client) - self._load_graphs_to_memory(client) + if not self._is_fts: + self._update_ef_search_before_search(client) + self._load_graphs_to_memory(client) def need_normalize_cosine(self) -> bool: """Whether this database needs to normalize dataset to support COSINE metric.""" return True + @classmethod + def supports_full_text_search(cls) -> bool: + return True + + def has_text_field(self) -> bool: + return bool(getattr(self, "_is_fts", False) and getattr(self, "text_col_name", None)) + def _get_cluster_version(self, client: OpenSearch) -> Version: """ Return the OpenSearch cluster version as a comparable Version object. @@ -307,7 +320,31 @@ def _get_bulk_manager(self, client: OpenSearch) -> BulkInsertManager: """Get bulk insert manager for the given client.""" return BulkInsertManager(client, self.index_name, self.case_config) + def _create_fts_index(self, client: OpenSearch) -> None: + mappings = self.case_config.index_param() + index_settings = { + "number_of_shards": self.case_config.number_of_shards, + "number_of_replicas": self.case_config.number_of_replicas, + "refresh_interval": self.case_config.refresh_interval, + } + index_settings.update(self.case_config.similarity_settings()) + settings = {"index": index_settings} + try: + log.info(f"Creating FTS index with settings: {settings}") + log.info(f"Creating FTS index with mappings: {mappings}") + client.indices.create( + index=self.index_name, + body={"settings": settings, "mappings": mappings}, + ) + except Exception as e: + log.warning(f"Failed to create FTS index: {self.index_name} error: {e!s}") + raise e from None + def _create_index(self, client: OpenSearch) -> None: + if self._is_fts: + self._create_fts_index(client) + return + cluster_version = self._get_cluster_version(client) if self.case_config.on_disk and cluster_version < Version("2.17"): @@ -420,6 +457,36 @@ def insert_embeddings( log.info(f"Using {num_clients} parallel clients for data insertion") return self._insert_with_multiple_clients(embeddings, metadata, num_clients, labels_data) + def insert_documents( + self, + texts: Iterable[str], + doc_ids: list[str], + **kwargs: Any, + ) -> tuple[int, Exception | None]: + if not getattr(self, "_is_fts", False): + msg = "OSSOpenSearch full-text insert requires OSSOpenSearchFtsConfig" + raise RuntimeError(msg) + assert self.client is not None, "should self.init() first" + docs = list(texts) + if len(docs) != len(doc_ids): + msg = f"Mismatch between texts ({len(docs)}) and doc_ids ({len(doc_ids)}) lengths" + raise ValueError(msg) + + insert_data: list[dict[str, Any]] = [] + for i, doc in enumerate(docs): + doc_id = str(doc_ids[i]) + insert_data.append({"index": {"_index": self.index_name, "_id": doc_id}}) + insert_data.append({self.id_col_name: doc_id, self.text_col_name: doc}) + + try: + response = self.client.bulk(body=insert_data) + if response.get("errors"): + log.warning(f"FTS bulk insert had errors: {response}") + return len(docs), None + except Exception as e: + log.warning(f"Failed to insert FTS docs: {self.index_name} error: {e!s}") + return 0, e + def _insert_with_single_client( self, embeddings: Iterable[list[float]], @@ -569,6 +636,47 @@ def search_embedding( log.warning(f"Failed to search: {self.index_name} error: {e!s}") raise e from None + def search_documents( + self, + query: str, + k: int = 100, + payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY, + **kwargs: Any, + ) -> list[str]: + if not getattr(self, "_is_fts", False): + msg = "OSSOpenSearch full-text search requires OSSOpenSearchFtsConfig" + raise RuntimeError(msg) + if not self.supports_document_payload_profile(payload_profile): + msg = f"OSSOpenSearch does not support document payload_profile={payload_profile.value}" + raise NotImplementedError(msg) + assert self.client is not None, "should self.init() first" + + source = [self.text_col_name] if payload_profile == PayloadProfile.TEXT else False + filter_path = ["hits.hits._id", f"hits.hits.fields.{self.id_col_name}"] + if payload_profile == PayloadProfile.TEXT: + filter_path.append(f"hits.hits._source.{self.text_col_name}") + search_kwargs: dict[str, Any] = { + "index": self.index_name, + "body": {"query": {"match": {self.text_col_name: query}}}, + "size": k, + "_source": source, + "docvalue_fields": [self.id_col_name], + "filter_path": filter_path, + } + if payload_profile != PayloadProfile.TEXT: + search_kwargs["stored_fields"] = "_none_" + response = self.client.search(**search_kwargs) + + doc_ids = [] + for hit in response.get("hits", {}).get("hits", []): + if hit.get("_id") is not None: + doc_ids.append(str(hit["_id"])) + continue + values = hit.get("fields", {}).get(self.id_col_name, []) + if values: + doc_ids.append(str(values[0])) + return doc_ids + def prepare_filter(self, filters: Filter) -> None: """Prepare filter conditions for search operations.""" self.routing_key = None @@ -587,6 +695,14 @@ def prepare_filter(self, filters: Filter) -> None: def optimize(self, data_size: int | None = None) -> None: """Optimize the index for better search performance.""" + if self._is_fts: + self._refresh_index() + if self.case_config.force_merge_enabled: + self._do_fts_force_merge() + self._refresh_index() + self._update_replicas() + self._refresh_index() + return self._update_ef_search() # Call refresh first to ensure that all segments are created self._refresh_index() @@ -645,6 +761,17 @@ def _refresh_index(self): continue log.debug(f"Completed refresh for index {self.index_name}") + def _do_fts_force_merge(self): + log.info(f"Starting FTS force merge for index {self.index_name}") + force_merge_endpoint = f"/{self.index_name}/_forcemerge?max_num_segments=1&wait_for_completion=false" + force_merge_task_id = self.client.transport.perform_request("POST", force_merge_endpoint)["task"] + while True: + time.sleep(WAITING_FOR_FORCE_MERGE_SEC) + task_status = self.client.tasks.get(task_id=force_merge_task_id) + if task_status["completed"]: + break + log.info(f"Completed FTS force merge for index {self.index_name}") + def _do_force_merge(self): log.info(f"Updating the Index thread qty to {self.case_config.index_thread_qty_during_force_merge}.") diff --git a/vectordb_bench/frontend/config/dbCaseConfigs.py b/vectordb_bench/frontend/config/dbCaseConfigs.py index 764dad2bc..76d02935c 100644 --- a/vectordb_bench/frontend/config/dbCaseConfigs.py +++ b/vectordb_bench/frontend/config/dbCaseConfigs.py @@ -2326,6 +2326,7 @@ class CaseConfigInput(BaseModel): ElasticCloudFtsConfig = [] VespaFtsConfig = [] +OSSOpenSearchFtsConfig = [] TurboPufferFtsConfig = [] WeaviateLoadConfig = [ @@ -3158,6 +3159,7 @@ class FilterType(Enum): DB.OSSOpenSearch: { CaseLabel.Load: OSSOpensearchLoadingConfig, CaseLabel.Performance: OSSOpenSearchPerformanceConfig, + CaseLabel.FullTextSearchPerformance: OSSOpenSearchFtsConfig, }, DB.PgVector: { CaseLabel.Load: PgVectorLoadingConfig, From dc700aa3c36501ec5479913eaf6dff5bf5e5dfcc Mon Sep 17 00:00:00 2001 From: jamesgao-jpg Date: Tue, 14 Jul 2026 00:08:39 +0800 Subject: [PATCH 2/2] feat: add insert batch and streaming rate controls Replace environment and CloudInsert-specific batch sizing with a task-level insert batch size shared by CLI, REST, frontend, datasets, and runners. Keep streaming insert rate case-specific and validate its relationship to batching. Signed-off-by: jamesgao-jpg --- .env.example | 1 - README.md | 10 +- docs/release/2026-05-cloud-leaderboard.md | 2 +- tests/test_case_runner_reuse.py | 12 + tests/test_cloud_insert_case.py | 45 ++-- tests/test_concurrent_runner.py | 16 +- tests/test_frontend_run_settings.py | 72 ++++++ tests/test_insert_control_cli.py | 131 +++++++++++ tests/test_insert_control_contract.py | 148 +++++++++++++ tests/test_insert_control_runners.py | 205 ++++++++++++++++++ tests/test_rate_runner.py | 160 +++++++------- vectordb_bench/__init__.py | 4 +- vectordb_bench/backend/cases.py | 28 +-- vectordb_bench/backend/clients/api.py | 3 +- vectordb_bench/backend/clients/doris/cli.py | 2 +- .../backend/clients/memorydb/cli.py | 12 +- .../backend/clients/memorydb/config.py | 7 +- .../backend/clients/memorydb/memorydb.py | 4 +- vectordb_bench/backend/dataset.py | 13 +- .../backend/runner/concurrent_runner.py | 6 +- vectordb_bench/backend/runner/mp_runner.py | 1 - vectordb_bench/backend/runner/rate_runner.py | 15 +- .../backend/runner/read_write_runner.py | 6 +- .../backend/runner/serial_runner.py | 19 +- vectordb_bench/backend/task_runner.py | 9 +- vectordb_bench/cli/cli.py | 39 +++- .../components/run_test/generateTasks.py | 9 +- .../components/run_test/runSettings.py | 59 +++++ .../frontend/config/dbCaseConfigs.py | 5 +- vectordb_bench/frontend/pages/run_test.py | 8 +- vectordb_bench/models.py | 49 ++++- vectordb_bench/restful/app.py | 2 + 32 files changed, 908 insertions(+), 194 deletions(-) create mode 100644 tests/test_frontend_run_settings.py create mode 100644 tests/test_insert_control_cli.py create mode 100644 tests/test_insert_control_contract.py create mode 100644 tests/test_insert_control_runners.py create mode 100644 vectordb_bench/frontend/components/run_test/runSettings.py diff --git a/.env.example b/.env.example index e495ea999..d0f5abbe3 100644 --- a/.env.example +++ b/.env.example @@ -2,7 +2,6 @@ LOG_LEVEL=INFO LOG_FILE="logs/vectordb_bench.log" # TIMEZONE= -# NUM_PER_BATCH= # DEFAULT_DATASET_URL= DATASET_LOCAL_DIR="/tmp/vectordb_bench/dataset" diff --git a/README.md b/README.md index 72e7f8ca7..1c919fdee 100644 --- a/README.md +++ b/README.md @@ -426,7 +426,7 @@ pip install 'vectordb-bench[hologres]' 'psycopg[binary]' pgvector Execute tests for the index types: HGraph. ```shell -NUM_PER_BATCH=10000 vectordbbench hologreshgraph --host Hologres_Endpoint --port 80 \ +vectordbbench hologreshgraph --host Hologres_Endpoint --port 80 --insert-batch-size 10000 \ --user ACCESS_ID --password ACCESS_KEY --database DATABASE_NAME \ --m 64 --ef-construction 400 --case-type Performance768D10M \ --index-type HGraph --ef-search 400 --k 10 --num-concurrency 1,60,70,75,80,90,95,100,110,120 @@ -479,12 +479,12 @@ To list the options for zvec, execute vectordbbench zvec --help Doris supports ann index with type hnsw from version 4.0.x ```shell -NUM_PER_BATCH=1000000 vectordbbench doris --http-port=8030 --port=9030 --db-name=vector_test --case-type=Performance768D1M --stream-load-rows-per-batch=500000 +vectordbbench doris --http-port=8030 --port=9030 --db-name=vector_test --case-type=Performance768D1M --insert-batch-size=1000000 --stream-load-rows-per-batch=500000 ``` Using flag `--session-var`, if you want to test doris with some customized session variables. For example: ```shell -NUM_PER_BATCH=1000000 vectordbbench doris --http-port=8030 --port=9030 --db-name=vector_test --case-type=Performance768D1M --stream-load-rows-per-batch=500000 --session-var enable_profile=True +vectordbbench doris --http-port=8030 --port=9030 --db-name=vector_test --case-type=Performance768D1M --insert-batch-size=1000000 --stream-load-rows-per-batch=500000 --session-var enable_profile=True ``` Mote options: @@ -505,8 +505,8 @@ Mote options: --session-var TEXT Session variable key=value applied to each SQL session (repeatable) --stream-load-rows-per-batch INTEGER - Rows per single stream load request; default - uses NUM_PER_BATCH + Rows per Doris stream-load request; when + omitted, the Doris client default is used --no-index Create table without ANN index ``` diff --git a/docs/release/2026-05-cloud-leaderboard.md b/docs/release/2026-05-cloud-leaderboard.md index c9dc76a7b..252b6cd3b 100644 --- a/docs/release/2026-05-cloud-leaderboard.md +++ b/docs/release/2026-05-cloud-leaderboard.md @@ -65,7 +65,7 @@ vectordbbench zillizautoindex \ --uri "$ZILLIZ_URI" \ --token "$ZILLIZ_TOKEN" \ --collection-name cloud_insert_laion100m_bs10k \ - --cloud-insert-batch-size 10000 \ + --insert-batch-size 10000 \ --load-concurrency 16 \ --skip-search-serial \ --skip-search-concurrent \ diff --git a/tests/test_case_runner_reuse.py b/tests/test_case_runner_reuse.py index 55dfbda0b..ce06e5ba0 100644 --- a/tests/test_case_runner_reuse.py +++ b/tests/test_case_runner_reuse.py @@ -1,5 +1,6 @@ from pydantic import SecretStr +from vectordb_bench import config from vectordb_bench.backend.clients import DB from vectordb_bench.backend.clients.api import EmptyDBCaseConfig, MetricType from vectordb_bench.backend.clients.doris.config import DorisCaseConfig, DorisConfig @@ -12,6 +13,8 @@ from vectordb_bench.metric import Metric from vectordb_bench.models import CaseConfig, CaseType, TaskConfig, TaskStage, TestResult +DEFAULT_INSERT_BATCH_SIZE = config.DEFAULT_INSERT_BATCH_SIZE + def make_runner( *, @@ -21,6 +24,7 @@ def make_runner( db_config=None, db_case_config=None, stages: list[TaskStage] | None = None, + insert_batch_size: int = DEFAULT_INSERT_BATCH_SIZE, ) -> CaseRunner: if db_config is None: if db == DB.TurboPuffer: @@ -45,6 +49,7 @@ def make_runner( db_case_config=db_case_config, case_config=CaseConfig(case_id=case_id, custom_case=custom_case or {}), stages=stages or [TaskStage.DROP_OLD, TaskStage.LOAD, TaskStage.SEARCH_SERIAL], + insert_batch_size=insert_batch_size, ) return CaseRunner( run_id="run-id", @@ -111,6 +116,13 @@ def test_reuse_key_preserves_safe_payload_reuse(): assert hash(ids_only) == hash(vector) +def test_reuse_key_distinguishes_insert_batch_size(): + assert_not_reusable( + make_runner(insert_batch_size=100), + make_runner(insert_batch_size=200), + ) + + def test_reuse_key_distinguishes_physical_db_targets(): assert_not_reusable( make_runner(db_config=TurboPufferConfig(api_key="key", region="aws-us-east-1", namespace="namespace_a")), diff --git a/tests/test_cloud_insert_case.py b/tests/test_cloud_insert_case.py index 28e70f7bd..9fc8ba1a6 100644 --- a/tests/test_cloud_insert_case.py +++ b/tests/test_cloud_insert_case.py @@ -69,13 +69,13 @@ def iter_batches(self, batch_size): def test_cloud_insert_case_defaults_to_laion_100m(): - case = CloudInsertCase(batch_size=1000) + case = CloudInsertCase() assert case.case_id == CaseType.CloudInsertCase assert case.label == CaseLabel.CloudInsert assert case.dataset.data.name == "LAION" assert case.dataset.data.size == 100_000_000 - assert case.batch_size == 1000 + assert not hasattr(case, "batch_size") assert case.duration is None assert case.readiness_timeout is None @@ -84,14 +84,13 @@ def test_case_config_builds_cloud_insert_case_from_custom_case(): case = CaseConfig( case_id=CaseType.CloudInsertCase, custom_case={ - "batch_size": 5000, "duration": 1800, "dataset_with_size_type": DatasetWithSizeType.CohereMedium.value, }, ).case assert isinstance(case, CloudInsertCase) - assert case.batch_size == 5000 + assert not hasattr(case, "batch_size") assert case.duration == 1800 assert case.dataset.data.name == "Cohere" assert case.dataset.data.size == 1_000_000 @@ -101,13 +100,12 @@ def test_case_config_builds_cloud_insert_case_from_laion_100m_dataset_option(): case = CaseConfig( case_id=CaseType.CloudInsertCase, custom_case={ - "batch_size": 10_000, "dataset_with_size_type": "Large LAION (768dim, 100M)", }, ).case assert isinstance(case, CloudInsertCase) - assert case.batch_size == 10_000 + assert not hasattr(case, "batch_size") assert case.dataset_with_size_type == DatasetWithSizeType.LAIONLarge assert case.dataset.data.name == "LAION" assert case.dataset.data.size == 100_000_000 @@ -122,7 +120,6 @@ def test_laion_100m_dataset_option_uses_100m_timeouts(): def test_cli_builds_cloud_insert_custom_case_config(): params = { "case_type": "CloudInsertCase", - "cloud_insert_batch_size": 10_000, "cloud_insert_duration": 1800, "cloud_insert_readiness_timeout": 7200, "cloud_insert_readiness_poll_interval": 10, @@ -130,7 +127,6 @@ def test_cli_builds_cloud_insert_custom_case_config(): } assert get_custom_case_config(params) == { - "batch_size": 10_000, "duration": 1800, "readiness_timeout": 7200, "readiness_poll_interval": 10, @@ -142,7 +138,6 @@ def test_cli_builds_cloud_insert_custom_case_config_with_laion_100m_dataset(): cfg = get_custom_case_config( { "case_type": "CloudInsertCase", - "cloud_insert_batch_size": 10_000, "cloud_insert_duration": None, "cloud_insert_readiness_timeout": None, "cloud_insert_readiness_poll_interval": None, @@ -151,7 +146,6 @@ def test_cli_builds_cloud_insert_custom_case_config_with_laion_100m_dataset(): ) assert cfg == { - "batch_size": 10_000, "duration": None, "dataset_with_size_type": DatasetWithSizeType.LAIONLarge.value, } @@ -166,7 +160,6 @@ def test_cli_builds_cloud_insert_custom_case_config_with_default_dataset(): cfg = get_custom_case_config( { "case_type": "CloudInsertCase", - "cloud_insert_batch_size": 10_000, "cloud_insert_duration": None, "cloud_insert_readiness_timeout": None, "cloud_insert_readiness_poll_interval": None, @@ -175,7 +168,6 @@ def test_cli_builds_cloud_insert_custom_case_config_with_default_dataset(): ) assert cfg == { - "batch_size": 10_000, "duration": None, "dataset_with_size_type": DatasetWithSizeType.CohereMedium.value, } @@ -243,11 +235,11 @@ def test_assembler_schedules_cloud_insert_case(): case_config=CaseConfig( case_id=CaseType.CloudInsertCase, custom_case={ - "batch_size": 1000, "dataset_with_size_type": DatasetWithSizeType.CohereMedium.value, }, ), stages=[TaskStage.DROP_OLD, TaskStage.LOAD], + insert_batch_size=1000, ) runner = Assembler.assemble_all("run-id", "task-label", [task], DatasetSource.S3) @@ -311,10 +303,11 @@ def test_cloud_insert_result_file_uses_insert_only_metrics(tmp_path: Path): db_case_config=EmptyDBCaseConfig(), case_config=CaseConfig( case_id=CaseType.CloudInsertCase, - custom_case={"batch_size": 1000, "duration": None}, + custom_case={"duration": None}, ), stages=[TaskStage.DROP_OLD, TaskStage.LOAD], load_concurrency=0, + insert_batch_size=1000, ), metrics=Metric( inserted_count=100_000_000, @@ -343,14 +336,16 @@ def test_cloud_insert_result_file_uses_insert_only_metrics(tmp_path: Path): } assert written["results"][0]["task_config"]["db_config"]["api_key"] == "**********" assert written["results"][0]["task_config"]["db_config"]["index_name"] == "laion100m" + assert written["results"][0]["task_config"]["insert_batch_size"] == 1000 assert written["results"][0]["task_config"]["case_config"] == { "case_id": 600, - "custom_case": {"batch_size": 1000, "duration": None}, + "custom_case": {"duration": None}, } read_back = TestResult.read_file(result_file) assert read_back.results[0].task_config.case_config.case_id == CaseType.CloudInsertCase - assert read_back.results[0].task_config.case_config.custom_case == {"batch_size": 1000, "duration": None} + assert read_back.results[0].task_config.case_config.custom_case == {"duration": None} + assert read_back.results[0].task_config.insert_batch_size == 1000 collected = ResultCollector.collect(tmp_path) assert len(collected) == 1 @@ -423,7 +418,7 @@ def write(self, **kwargs): def test_milvus_insert_readiness_uses_entity_count_and_index_progress(): db = Milvus.__new__(Milvus) db.collection_name = "c" - db._vector_index_name = "vector_idx" + db._main_index_name = "vector_idx" db.client = type( "Client", (), @@ -697,9 +692,9 @@ def poll_insert_readiness(self, expected_count): db = DB() monkeypatch.setattr("vectordb_bench.backend.task_runner.time.sleep", lambda _: None) - case = CloudInsertCase(batch_size=2) + case = CloudInsertCase() case.dataset = Dataset() - config = type("Config", (), {"load_concurrency": 1})() + config = type("Config", (), {"load_concurrency": 1, "insert_batch_size": 2})() runner = CaseRunner.construct(ca=case, db=db, config=config) metric = runner._run_cloud_insert_case() @@ -752,9 +747,13 @@ def fail_on_sleep(_seconds): monkeypatch.setattr("vectordb_bench.backend.task_runner.ConcurrentInsertRunner", FakeConcurrentInsertRunner) monkeypatch.setattr("vectordb_bench.backend.task_runner.time.sleep", fail_on_sleep) - case = CloudInsertCase(batch_size=1, readiness_timeout=0, readiness_poll_interval=0) + case = CloudInsertCase(readiness_timeout=0, readiness_poll_interval=0) case.dataset = Dataset() - runner = CaseRunner.construct(ca=case, db=DB(), config=type("Config", (), {"load_concurrency": 1})()) + runner = CaseRunner.construct( + ca=case, + db=DB(), + config=type("Config", (), {"load_concurrency": 1, "insert_batch_size": 1})(), + ) with pytest.raises(TimeoutError, match="fully_searchable.*last_status.*stalled"): runner._run_cloud_insert_case() @@ -798,9 +797,9 @@ def poll_insert_readiness(self, expected_count): return {"fully_searchable": True, "fully_indexed": True, "additional_parameters": {}} monkeypatch.setattr("vectordb_bench.backend.task_runner.ConcurrentInsertRunner", FakeConcurrentInsertRunner) - case = CloudInsertCase(batch_size=1000, duration=60) + case = CloudInsertCase(duration=60) case.dataset = Dataset() - config = type("Config", (), {"load_concurrency": 7})() + config = type("Config", (), {"load_concurrency": 7, "insert_batch_size": 1000})() runner = CaseRunner.construct(ca=case, db=DB(), config=config) metric = runner._run_cloud_insert_case() diff --git a/tests/test_concurrent_runner.py b/tests/test_concurrent_runner.py index c9e5d9267..8ad2a06b1 100644 --- a/tests/test_concurrent_runner.py +++ b/tests/test_concurrent_runner.py @@ -4,8 +4,8 @@ - Correctness tests (threading & async backends) - Parameterized benchmark: serial vs concurrent across (batch_size, workers) matrix -NUM_PER_BATCH is set via os.environ before each run. Since runners execute -task() in a spawn subprocess that re-imports config, the env var takes effect. +Batch size is passed directly to each runner so subprocess execution uses the +same explicit benchmark value. Requires: - Milvus running at localhost:19530 @@ -21,7 +21,6 @@ from __future__ import annotations import logging -import os import time from vectordb_bench.backend.clients import DB @@ -55,10 +54,6 @@ def prepare_dataset(): return dataset -def set_batch_size(batch_size: int) -> None: - os.environ["NUM_PER_BATCH"] = str(batch_size) - - def timed_run(runner: SerialInsertRunner | ConcurrentInsertRunner) -> tuple[int, float]: start = time.perf_counter() count = runner.run() @@ -100,23 +95,23 @@ def test_concurrent_insert_async(): def run_serial(batch_size: int) -> tuple[int, float]: - set_batch_size(batch_size) runner = SerialInsertRunner( db=get_milvus_db(f"bench_serial_b{batch_size}"), dataset=prepare_dataset(), normalize=False, + batch_size=batch_size, ) return timed_run(runner) def run_concurrent(batch_size: int, workers: int) -> tuple[int, float]: - set_batch_size(batch_size) runner = ConcurrentInsertRunner( db=get_milvus_db(f"bench_conc_b{batch_size}_w{workers}"), dataset=prepare_dataset(), normalize=False, max_workers=workers, backend=ExecutorBackend.THREADING, + batch_size=batch_size, ) return timed_run(runner) @@ -151,9 +146,6 @@ def bench_matrix(): print(f" {dur_s / dur_c:>11.2f}x", end="") print() - # restore default - set_batch_size(100) - if __name__ == "__main__": bench_matrix() diff --git a/tests/test_frontend_run_settings.py b/tests/test_frontend_run_settings.py new file mode 100644 index 000000000..4c421aa92 --- /dev/null +++ b/tests/test_frontend_run_settings.py @@ -0,0 +1,72 @@ +from collections import defaultdict + +import pytest + +from vectordb_bench.backend.cases import CaseType +from vectordb_bench.backend.clients import DB +from vectordb_bench.frontend.components.run_test import generateTasks +from vectordb_bench.frontend.components.run_test.runSettings import ( + DEFAULT_STREAMING_INSERT_RATE, + validate_streaming_insert_rates, +) +from vectordb_bench.models import CaseConfig + + +def streaming_case(insert_rate: int | None = None) -> CaseConfig: + custom_case = {} if insert_rate is None else {"insert_rate": insert_rate} + return CaseConfig(case_id=CaseType.StreamingPerformanceCase, custom_case=custom_case) + + +@pytest.mark.parametrize( + ("insert_rate", "batch_size", "expected_message"), + [ + (400, 500, "must be greater than or equal to"), + (750, 500, "must be divisible by"), + ], +) +def test_validate_streaming_insert_rates_rejects_invalid_rate( + insert_rate: int, + batch_size: int, + expected_message: str, +): + is_valid, errors = validate_streaming_insert_rates([streaming_case(insert_rate)], batch_size) + + assert not is_valid + assert len(errors) == 1 + assert expected_message in errors[0] + + +def test_validate_streaming_insert_rates_checks_each_streaming_case_and_uses_default(): + cases = [ + CaseConfig(case_id=CaseType.Performance768D1M), + streaming_case(), + CaseConfig(case_id=CaseType.StreamingCustomDataset, custom_case={"insert_rate": 1_000}), + ] + + is_valid, errors = validate_streaming_insert_rates(cases, DEFAULT_STREAMING_INSERT_RATE) + + assert is_valid + assert errors == [] + + +def test_generate_tasks_passes_batch_size_to_task_config(monkeypatch: pytest.MonkeyPatch): + captured_task_configs: list[dict[str, object]] = [] + + class CapturedTaskConfig: + def __init__(self, **kwargs: object): + captured_task_configs.append(kwargs) + + monkeypatch.setattr(generateTasks, "TaskConfig", CapturedTaskConfig) + case = CaseConfig(case_id=CaseType.Performance768D1M) + all_case_configs = defaultdict(lambda: defaultdict(dict)) + + tasks = generateTasks.generate_tasks( + [DB.Test], + {DB.Test: DB.Test.config_cls()}, + [case], + all_case_configs, + batch_size=250, + ) + + assert len(tasks) == 1 + assert captured_task_configs[0]["insert_batch_size"] == 250 diff --git a/tests/test_insert_control_cli.py b/tests/test_insert_control_cli.py new file mode 100644 index 000000000..df9ce7c6a --- /dev/null +++ b/tests/test_insert_control_cli.py @@ -0,0 +1,131 @@ +import pytest +from click.testing import CliRunner + +from vectordb_bench.backend.clients.memorydb import cli as memorydb_cli +from vectordb_bench.backend.clients.memorydb.config import MemoryDBHNSWConfig +from vectordb_bench.backend.clients.test import cli as test_cli +from vectordb_bench.cli import cli as core_cli + + +@pytest.mark.parametrize( + ("args", "expected_batch_size"), + [ + (["--dry-run"], 100), + (["--dry-run", "--insert-batch-size", "250"], 250), + ], +) +def test_common_insert_batch_size_is_forwarded_to_task_config( + monkeypatch: pytest.MonkeyPatch, + args: list[str], + expected_batch_size: int, +) -> None: + captured = {} + + class FakeTaskConfig: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(core_cli, "TaskConfig", FakeTaskConfig) + + result = CliRunner().invoke(test_cli.Test, args) + + assert result.exit_code == 0, result.output + assert captured["insert_batch_size"] == expected_batch_size + + +@pytest.mark.parametrize("option", ["--insert-batch-size", "--streaming-insert-rate"]) +def test_positive_insert_controls_reject_zero(option: str) -> None: + result = CliRunner().invoke(test_cli.Test, ["--dry-run", option, "0"]) + + assert result.exit_code == 2 + assert "x>=1" in result.output + + +@pytest.mark.parametrize( + "case_type", + ["StreamingPerformanceCase", "StreamingCustomDataset"], +) +def test_streaming_insert_rate_only_maps_to_streaming_cases(case_type: str) -> None: + streaming = core_cli.get_custom_case_config( + { + "case_type": case_type, + "dataset_with_size_type": None, + "streaming_insert_rate": 750, + }, + ) + non_streaming = core_cli.get_custom_case_config( + { + "case_type": "Performance1536D50K", + "dataset_with_size_type": None, + "streaming_insert_rate": 750, + }, + ) + + assert streaming == {"insert_rate": 750} + assert non_streaming == {} + + +def test_cloud_insert_no_longer_has_a_custom_batch_mapping() -> None: + custom_case = core_cli.get_custom_case_config( + { + "case_type": "CloudInsertCase", + "dataset_with_size_type": None, + "cloud_insert_duration": None, + "cloud_insert_readiness_timeout": None, + "cloud_insert_readiness_poll_interval": None, + }, + ) + + assert "batch_size" not in custom_case + + +def test_memorydb_cli_keeps_task_and_pipeline_batch_sizes_distinct(monkeypatch: pytest.MonkeyPatch) -> None: + captured = {} + + def fake_run(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(memorydb_cli, "run", fake_run) + + result = CliRunner().invoke( + memorydb_cli.MemoryDB, + [ + "--host", + "localhost", + "--dry-run", + "--insert-batch-size", + "200", + "--memorydb-pipeline-batch-size", + "8", + ], + ) + + assert result.exit_code == 0, result.output + assert captured["insert_batch_size"] == 200 + assert captured["db_case_config"].pipeline_batch_size == 8 + + +def test_memorydb_config_accepts_legacy_insert_batch_size() -> None: + config = MemoryDBHNSWConfig.model_validate({"insert_batch_size": 12}) + + assert config.pipeline_batch_size == 12 + assert config.model_dump()["pipeline_batch_size"] == 12 + assert "insert_batch_size" not in config.model_dump() + + +def test_memorydb_config_prefers_canonical_pipeline_batch_size() -> None: + config = MemoryDBHNSWConfig.model_validate( + { + "pipeline_batch_size": 8, + "insert_batch_size": 12, + }, + ) + + assert config.pipeline_batch_size == 8 + + +def test_removed_cloud_insert_batch_option_is_rejected() -> None: + result = CliRunner().invoke(test_cli.Test, ["--dry-run", "--cloud-insert-batch-size", "5000"]) + + assert result.exit_code == 2 + assert "No such option '--cloud-insert-batch-size'" in result.output diff --git a/tests/test_insert_control_contract.py b/tests/test_insert_control_contract.py new file mode 100644 index 000000000..c4ce5f8b8 --- /dev/null +++ b/tests/test_insert_control_contract.py @@ -0,0 +1,148 @@ +import importlib +import json +from pathlib import Path +from typing import Any + +import pytest +from pydantic import ValidationError + +from vectordb_bench import config +from vectordb_bench.backend.cases import CaseType +from vectordb_bench.backend.clients import DB, EmptyDBCaseConfig +from vectordb_bench.backend.clients.test.config import TestConfig +from vectordb_bench.metric import Metric +from vectordb_bench.models import CaseConfig, CaseResult, TaskConfig, TestResult + + +def make_task( + case_id: CaseType = CaseType.Performance768D1M, + custom_case: dict | None = None, + **overrides: Any, +) -> TaskConfig: + values = { + "db": DB.Test, + "db_config": TestConfig(), + "db_case_config": EmptyDBCaseConfig(), + "case_config": CaseConfig(case_id=case_id, custom_case=custom_case), + } + values.update(overrides) + return TaskConfig(**values) + + +def write_legacy_result(tmp_path: Path, task: TaskConfig, metrics: Metric) -> Path: + result = TestResult( + run_id="legacy-run", + task_label="legacy", + results=[CaseResult(task_config=task, metrics=metrics)], + ) + raw = result.model_dump(mode="json", serialize_as_any=True) + raw["results"][0]["task_config"].pop("insert_batch_size") + result_path = tmp_path / "legacy.json" + result_path.write_text(json.dumps(raw)) + return result_path + + +def test_insert_control_defaults_and_serialization(): + task = make_task() + + assert config.DEFAULT_INSERT_BATCH_SIZE == 100 + assert config.DEFAULT_STREAMING_INSERT_RATE == 500 + assert task.insert_batch_size == 100 + assert task.model_dump(mode="json")["insert_batch_size"] == 100 + assert "insert_rate" not in TaskConfig.model_fields + + +@pytest.mark.parametrize("insert_batch_size", [0, -1]) +def test_insert_batch_size_must_be_positive(insert_batch_size: int): + with pytest.raises(ValidationError, match="greater than 0"): + make_task(insert_batch_size=insert_batch_size) + + +@pytest.mark.parametrize( + "case_id", + [CaseType.StreamingPerformanceCase, CaseType.StreamingCustomDataset], +) +def test_streaming_rate_must_cover_and_divide_batch(case_id: CaseType): + valid = make_task(case_id, {"insert_rate": 1000}, insert_batch_size=250) + assert valid.case_config.custom_case["insert_rate"] == 1000 + + with pytest.raises(ValidationError, match="greater than or equal"): + make_task(case_id, {"insert_rate": 200}, insert_batch_size=250) + with pytest.raises(ValidationError, match="divisible"): + make_task(case_id, {"insert_rate": 550}, insert_batch_size=100) + + +def test_streaming_rate_uses_stable_default(): + assert make_task(CaseType.StreamingPerformanceCase, insert_batch_size=250).insert_batch_size == 250 + with pytest.raises(ValidationError, match="divisible"): + make_task(CaseType.StreamingPerformanceCase, insert_batch_size=300) + + +def test_read_file_migrates_cloud_insert_batch_size(tmp_path: Path): + path = write_legacy_result( + tmp_path, + make_task( + CaseType.CloudInsertCase, + {"batch_size": 250}, + insert_batch_size=100, + ), + Metric(additional_parameters={"num_per_batch": 100}), + ) + + result = TestResult.read_file(path) + + assert result.results[0].task_config.insert_batch_size == 250 + + +def test_read_file_migrates_metrics_batch_size(tmp_path: Path): + path = write_legacy_result( + tmp_path, + make_task( + CaseType.StreamingPerformanceCase, + {"insert_rate": 500}, + insert_batch_size=100, + ), + Metric(additional_parameters={"num_per_batch": 250}), + ) + + result = TestResult.read_file(path) + + assert result.results[0].task_config.insert_batch_size == 250 + + +def test_rest_run_accepts_insert_batch_size(monkeypatch: pytest.MonkeyPatch): + pytest.importorskip("flask") + restful_app = importlib.import_module("vectordb_bench.restful.app") + + captured: dict[str, Any] = {} + monkeypatch.setattr(restful_app.benchmark_runner, "has_running", lambda: False) + monkeypatch.setattr(restful_app.benchmark_runner, "set_download_address", lambda _value: None) + monkeypatch.setattr( + restful_app.benchmark_runner, + "run", + lambda tasks, task_label: captured.update(tasks=tasks, task_label=task_label), + ) + + response = restful_app.app.test_client().post( + "/run", + json={ + "task_label": "contract", + "tasks": [ + { + "db": DB.Test.value, + "db_config": {}, + "db_case_config": {}, + "case_config": { + "case_id": CaseType.StreamingPerformanceCase.value, + "custom_case": {"insert_rate": 1000}, + }, + "stages": [], + "insert_batch_size": 250, + } + ], + }, + ) + + assert response.get_json()["code"] == 0 + assert captured["task_label"] == "contract" + assert captured["tasks"][0].insert_batch_size == 250 diff --git a/tests/test_insert_control_runners.py b/tests/test_insert_control_runners.py new file mode 100644 index 000000000..db69e0fa0 --- /dev/null +++ b/tests/test_insert_control_runners.py @@ -0,0 +1,205 @@ +from contextlib import contextmanager +from types import SimpleNamespace +from typing import Any + +import pytest + +from vectordb_bench import config +from vectordb_bench.backend import task_runner as task_runner_module +from vectordb_bench.backend.cases import CaseLabel, StreamingPerformanceCase +from vectordb_bench.backend.dataset import DataSetIterator, FtsDocumentIterator +from vectordb_bench.backend.filter import non_filter +from vectordb_bench.backend.runner.concurrent_runner import ConcurrentInsertRunner +from vectordb_bench.backend.runner.serial_runner import SerialInsertRunner +from vectordb_bench.backend.task_runner import CaseRunner +from vectordb_bench.backend.workload import WorkloadKind + +DEFAULT_INSERT_BATCH_SIZE = config.DEFAULT_INSERT_BATCH_SIZE + + +class FakeDB: + name = "FakeDB" + thread_safe = True + + @contextmanager + def init(self): + yield + + def need_normalize_cosine(self): + return False + + +def make_case_runner( + case: Any, + *, + batch_size: int = 17, + load_concurrency: int = 3, + db: Any | None = None, +) -> CaseRunner: + task_config = SimpleNamespace( + insert_batch_size=batch_size, + load_concurrency=load_concurrency, + case_config=SimpleNamespace(k=10), + ) + return CaseRunner.model_construct(ca=case, config=task_config, db=db or FakeDB()) + + +def test_streaming_case_preserves_requested_insert_rate(): + case = StreamingPerformanceCase(insert_rate=550) + + assert case.insert_rate == 550 + assert "550 rows/s" in case.name + + +def test_direct_callers_use_stable_batch_default(): + dataset = SimpleNamespace(train_files=[]) + concurrent_dataset = SimpleNamespace(data=SimpleNamespace()) + + assert DataSetIterator(dataset)._batch_size == DEFAULT_INSERT_BATCH_SIZE + assert FtsDocumentIterator(SimpleNamespace())._batch_size == DEFAULT_INSERT_BATCH_SIZE + assert ConcurrentInsertRunner(FakeDB(), concurrent_dataset, normalize=False).batch_size == DEFAULT_INSERT_BATCH_SIZE + assert SerialInsertRunner(FakeDB(), dataset, normalize=False).batch_size == DEFAULT_INSERT_BATCH_SIZE + + +def test_serial_insert_runner_groups_rows_by_explicit_batch_size(): + class InsertDB(FakeDB): + def __init__(self): + self.metadata_batches = [] + + def insert_embeddings( + self, + embeddings: list[Any], + metadata: list[Any], + ) -> tuple[int, None]: + self.metadata_batches.append(metadata) + return len(metadata), None + + db = InsertDB() + runner = SerialInsertRunner(db, SimpleNamespace(), normalize=False, batch_size=2) + + inserted = runner.endless_insert_data( + all_embeddings=[[0.1], [0.2], [0.3], [0.4], [0.5]], + all_metadata=[0, 1, 2, 3, 4], + ) + + assert inserted == 5 + assert db.metadata_batches == [[0, 1], [2, 3], [4]] + + +@pytest.mark.parametrize( + ("label", "workload_kind"), + [ + (CaseLabel.Performance, WorkloadKind.VECTOR), + (CaseLabel.FullTextSearchPerformance, WorkloadKind.FULL_TEXT), + ], +) +def test_performance_load_propagates_task_batch( + monkeypatch: pytest.MonkeyPatch, + label: CaseLabel, + workload_kind: WorkloadKind, +): + created: dict[str, Any] = {} + + class FakeConcurrentInsertRunner: + def __init__(self, *args, **kwargs): + created.update(kwargs) + + def run(self): + return 9, 1.25 + + case = SimpleNamespace( + label=label, + is_multitenant=False, + dataset=SimpleNamespace(data=SimpleNamespace(metric_type="L2")), + filters=non_filter, + load_timeout=30, + with_scalar_labels=False, + ) + monkeypatch.setattr(task_runner_module, "ConcurrentInsertRunner", FakeConcurrentInsertRunner) + + result = make_case_runner(case)._load_train_data() + + assert result == (9, 1.25) + assert created["batch_size"] == 17 + assert created["workload_kind"] == workload_kind + + +def test_capacity_load_propagates_task_batch(monkeypatch: pytest.MonkeyPatch): + created: dict[str, Any] = {} + + class FakeSerialInsertRunner: + def __init__(self, *args, **kwargs): + created.update(kwargs) + + def run_endlessness(self): + return 123 + + case = SimpleNamespace( + label=CaseLabel.Load, + dataset=SimpleNamespace(data=SimpleNamespace(metric_type="L2")), + filters=non_filter, + load_timeout=30, + ) + monkeypatch.setattr(task_runner_module, "SerialInsertRunner", FakeSerialInsertRunner) + + metric = make_case_runner(case)._run_capacity_case() + + assert metric.max_load_count == 123 + assert created["batch_size"] == 17 + + +def test_cloud_insert_propagates_task_batch(monkeypatch: pytest.MonkeyPatch): + created: dict[str, Any] = {} + + class FakeConcurrentInsertRunner: + def __init__(self, *args, **kwargs): + created.update(kwargs) + + def task(self): + return 3 + + class ReadinessDB(FakeDB): + def poll_insert_readiness(self, expected_count: int) -> dict[str, Any]: + assert expected_count == 3 + return {"fully_searchable": True, "fully_indexed": True, "additional_parameters": {}} + + case = SimpleNamespace( + label=CaseLabel.CloudInsert, + is_multitenant=False, + dataset=SimpleNamespace(data=SimpleNamespace(metric_type="L2")), + filters=non_filter, + duration=60, + readiness_timeout=None, + readiness_poll_interval=0, + ) + monkeypatch.setattr(task_runner_module, "ConcurrentInsertRunner", FakeConcurrentInsertRunner) + + metric = make_case_runner(case, db=ReadinessDB())._run_cloud_insert_case() + + assert metric.inserted_count == 3 + assert created["batch_size"] == 17 + assert created["duration"] == 60 + + +def test_streaming_runner_propagates_task_batch(monkeypatch: pytest.MonkeyPatch): + created: dict[str, Any] = {} + + class FakeReadWriteRunner: + def __init__(self, **kwargs): + created.update(kwargs) + + case = SimpleNamespace( + label=CaseLabel.Streaming, + dataset=SimpleNamespace(data=SimpleNamespace(metric_type="L2")), + insert_rate=34, + search_stages=[0.5], + optimize_after_write=False, + read_dur_after_write=10, + concurrencies=[1], + ) + monkeypatch.setattr(task_runner_module, "ReadWriteRunner", FakeReadWriteRunner) + + make_case_runner(case)._init_read_write_runner() + + assert created["insert_rate"] == 34 + assert created["batch_size"] == 17 diff --git a/tests/test_rate_runner.py b/tests/test_rate_runner.py index df92b0dd7..a02b9da36 100644 --- a/tests/test_rate_runner.py +++ b/tests/test_rate_runner.py @@ -1,88 +1,90 @@ -from typing import Iterable -import argparse -from vectordb_bench.backend.dataset import Dataset, DatasetSource +from types import SimpleNamespace + +import pytest + +from vectordb_bench import config +from vectordb_bench.backend.runner import read_write_runner as read_write_runner_module +from vectordb_bench.backend.runner.mp_runner import MultiProcessingSearchRunner from vectordb_bench.backend.runner.rate_runner import RatedMultiThreadingInsertRunner from vectordb_bench.backend.runner.read_write_runner import ReadWriteRunner -from vectordb_bench.backend.clients import DB, VectorDB -from vectordb_bench.backend.clients.milvus.config import FLATConfig -from vectordb_bench.backend.clients.zilliz_cloud.config import AutoIndexConfig -import logging +DEFAULT_INSERT_BATCH_SIZE = config.DEFAULT_INSERT_BATCH_SIZE + + +class FakeDB: + name = "FakeDB" + -log = logging.getLogger("vectordb_bench") -log.setLevel(logging.DEBUG) +def test_rate_runner_uses_explicit_batch_size(): + runner = RatedMultiThreadingInsertRunner( + rate=30, + db=FakeDB(), + dataset_iter=iter(()), + batch_size=5, + ) + + assert runner.insert_rate == 30 + assert runner.batch_size == 5 + assert runner.batch_rate == 6 -def get_rate_runner(db): - cohere = Dataset.COHERE.manager(100_000) - prepared = cohere.prepare(DatasetSource.AliyunOSS) - assert prepared + +def test_rate_runner_direct_caller_uses_stable_batch_default(): runner = RatedMultiThreadingInsertRunner( - rate = 10, - db = db, - dataset = cohere, + rate=DEFAULT_INSERT_BATCH_SIZE, + db=FakeDB(), + dataset_iter=iter(()), ) - return runner - -def test_rate_runner(db, insert_rate): - runner = get_rate_runner(db) - - _, t = runner.run_with_rate() - log.info(f"insert run done, time={t}") - -def test_read_write_runner(db, insert_rate, conc: list, search_stage: Iterable[float], read_dur_after_write: int, local: bool=False): - cohere = Dataset.COHERE.manager(1_000_000) - if local is True: - source = DatasetSource.AliyunOSS - else: - source = DatasetSource.S3 - prepared = cohere.prepare(source) - assert prepared - - rw_runner = ReadWriteRunner( - db=db, - dataset=cohere, - insert_rate=insert_rate, - search_stage=search_stage, - read_dur_after_write=read_dur_after_write, - concurrencies=conc + assert runner.batch_size == DEFAULT_INSERT_BATCH_SIZE + assert runner.batch_rate == 1 + + +@pytest.mark.parametrize( + ("rate", "batch_size", "message"), + [ + (0, 10, "insert rate must be greater than 0"), + (-10, 10, "insert rate must be greater than 0"), + (10, 0, "insert batch size must be greater than 0"), + (10, -1, "insert batch size must be greater than 0"), + (10, 4, "insert rate 10 must be divisible by insert batch size 4"), + ], +) +def test_rate_runner_rejects_invalid_rate_batch_combinations(rate, batch_size, message): + with pytest.raises(ValueError, match=message): + RatedMultiThreadingInsertRunner( + rate=rate, + db=FakeDB(), + dataset_iter=iter(()), + batch_size=batch_size, + ) + + +def test_read_write_runner_requests_task_batch_from_dataset(monkeypatch): + requested_batch_sizes = [] + + class Dataset: + data = SimpleNamespace(size=100) + test_data = [] + gt_data = [] + + def iter_batches(self, batch_size): + requested_batch_sizes.append(batch_size) + return iter(()) + + class FakeSerialSearchRunner: + def __init__(self, **kwargs): + pass + + monkeypatch.setattr(MultiProcessingSearchRunner, "__init__", lambda self, **kwargs: None) + monkeypatch.setattr(read_write_runner_module, "SerialSearchRunner", FakeSerialSearchRunner) + + runner = ReadWriteRunner( + db=FakeDB(), + dataset=Dataset(), + insert_rate=30, + batch_size=5, ) - rw_runner.run_read_write() - - -def get_db(db: str, config: dict) -> VectorDB: - if db == DB.Milvus.name: - return DB.Milvus.init_cls(dim=768, db_config=config, db_case_config=FLATConfig(metric_type="COSINE"), drop_old=True) - elif db == DB.ZillizCloud.name: - return DB.ZillizCloud.init_cls(dim=768, db_config=config, db_case_config=AutoIndexConfig(metric_type="COSINE"), drop_old=True) - else: - raise ValueError(f"unknown db: {db}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-r", "--insert_rate", type=int, default="1000", help="insert entity row count per seconds, cps") - parser.add_argument("-d", "--db", type=str, default=DB.Milvus.name, help="db name") - parser.add_argument("-t", "--duration", type=int, default=300, help="stage search duration in seconds") - parser.add_argument("--use_s3", action='store_true', help="whether to use S3 dataset") - - flags = parser.parse_args() - - # TODO read uri, user, password from .env - config = { - "uri": "http://localhost:19530", - "user": "", - "password": "", - } - - conc = (1, 15, 50) - search_stage = (0.5, 0.6, 0.7, 0.8, 0.9) - - db = get_db(flags.db, config) - test_read_write_runner( - db=db, - insert_rate=flags.insert_rate, - conc=conc, - search_stage=search_stage, - read_dur_after_write=flags.duration, - local=flags.use_s3) + + assert requested_batch_sizes == [5] + assert runner.batch_size == 5 + assert runner.batch_rate == 6 diff --git a/vectordb_bench/__init__.py b/vectordb_bench/__init__.py index 3e5c1e69e..c053f9604 100644 --- a/vectordb_bench/__init__.py +++ b/vectordb_bench/__init__.py @@ -13,13 +13,15 @@ class config: ALIYUN_OSS_URL = "assets.zilliz.com.cn/benchmark/" AWS_S3_URL = "assets.zilliz.com/benchmark/" + DEFAULT_INSERT_BATCH_SIZE = 100 + DEFAULT_STREAMING_INSERT_RATE = 500 + LOG_LEVEL = env.str("LOG_LEVEL", "INFO") LOG_FILE = env.str("LOG_FILE", "logs/vectordb_bench.log") DEFAULT_DATASET_URL = env.str("DEFAULT_DATASET_URL", AWS_S3_URL) DATASET_SOURCE = env.str("DATASET_SOURCE", "S3") # Options "S3", "AliyunOSS", or "IR_DATASETS" DATASET_LOCAL_DIR = env.path("DATASET_LOCAL_DIR", "/tmp/vectordb_bench/dataset") - NUM_PER_BATCH = env.int("NUM_PER_BATCH", 100) LOAD_CONCURRENCY = env.int("LOAD_CONCURRENCY", 0) # 0 = cpu_count TIME_PER_BATCH = 1 # 1s. for streaming insertion. MAX_INSERT_RETRY = 5 diff --git a/vectordb_bench/backend/cases.py b/vectordb_bench/backend/cases.py index 0253168e0..ff592525b 100644 --- a/vectordb_bench/backend/cases.py +++ b/vectordb_bench/backend/cases.py @@ -464,17 +464,6 @@ def __init__( concurrencies: list[int] | str = (5, 10), **kwargs, ): - num_per_batch = config.NUM_PER_BATCH - if insert_rate % config.NUM_PER_BATCH != 0: - _insert_rate = max( - num_per_batch, - insert_rate // num_per_batch * num_per_batch, - ) - log.warning( - f"[streaming_case init] insert_rate(={insert_rate}) should be " - f"divisible by NUM_PER_BATCH={num_per_batch}), reset to {_insert_rate}", - ) - insert_rate = _insert_rate if not isinstance(dataset_with_size_type, DatasetWithSizeType): dataset_with_size_type = DatasetWithSizeType(dataset_with_size_type) dataset = dataset_with_size_type.get_manager() @@ -524,18 +513,6 @@ def __init__( read_dur_after_write: int = 30, **kwargs, ): - num_per_batch = config.NUM_PER_BATCH - if insert_rate % config.NUM_PER_BATCH != 0: - _insert_rate = max( - num_per_batch, - insert_rate // num_per_batch * num_per_batch, - ) - log.warning( - f"[streaming_case init] insert_rate(={insert_rate}) should be " - f"divisible by NUM_PER_BATCH={num_per_batch}), reset to {_insert_rate}", - ) - insert_rate = _insert_rate - dataset_config = CustomDatasetConfig(**dataset_config) dataset = CustomDataset( name=dataset_config.name, @@ -756,7 +733,6 @@ def filters(self) -> Filter: class CloudInsertCase(Case): case_id: CaseType = CaseType.CloudInsertCase label: CaseLabel = CaseLabel.CloudInsert - batch_size: int duration: float | None = None readiness_timeout: float | None = config.CLOUD_INSERT_READINESS_TIMEOUT readiness_poll_interval: float = config.CLOUD_INSERT_READINESS_POLL_INTERVAL @@ -764,7 +740,6 @@ class CloudInsertCase(Case): def __init__( self, - batch_size: int, duration: float | None = None, readiness_timeout: float | None = config.CLOUD_INSERT_READINESS_TIMEOUT, readiness_poll_interval: float = config.CLOUD_INSERT_READINESS_POLL_INTERVAL, @@ -779,10 +754,9 @@ def __init__( else dataset_with_size_type.get_manager() ) super().__init__( - name=f"Cloud Insert - batch {batch_size}", + name="Cloud Insert", description="Cloud leaderboard insert-only case with readiness polling.", dataset=dataset, - batch_size=batch_size, duration=duration, readiness_timeout=readiness_timeout, readiness_poll_interval=readiness_poll_interval, diff --git a/vectordb_bench/backend/clients/api.py b/vectordb_bench/backend/clients/api.py index c5474fba7..6d8239ce9 100644 --- a/vectordb_bench/backend/clients/api.py +++ b/vectordb_bench/backend/clients/api.py @@ -358,8 +358,7 @@ def insert_embeddings( tenant_labels_data: list[str] | None = None, **kwargs, ) -> tuple[int, Exception]: - """Insert the embeddings to the vector database. The default number of embeddings for - each insert_embeddings is 5000. + """Insert one task-configured batch of embeddings into the vector database. Args: embeddings(list[list[float]]): list of embedding to add to the vector database. diff --git a/vectordb_bench/backend/clients/doris/cli.py b/vectordb_bench/backend/clients/doris/cli.py index 8153b412d..9e796f937 100644 --- a/vectordb_bench/backend/clients/doris/cli.py +++ b/vectordb_bench/backend/clients/doris/cli.py @@ -144,7 +144,7 @@ class DorisTypedDict(CommonTypedDict, HNSWBaseTypedDict): "--stream-load-rows-per-batch", type=int, required=False, - help="Rows per single stream load request; default uses NUM_PER_BATCH", + help="Rows per Doris stream-load request; when omitted, the Doris client default is used", ), ] no_index: Annotated[ diff --git a/vectordb_bench/backend/clients/memorydb/cli.py b/vectordb_bench/backend/clients/memorydb/cli.py index 568eec2a3..8e43f0843 100644 --- a/vectordb_bench/backend/clients/memorydb/cli.py +++ b/vectordb_bench/backend/clients/memorydb/cli.py @@ -48,13 +48,15 @@ class MemoryDBTypedDict(TypedDict): ), ), ] - insert_batch_size: Annotated[ + pipeline_batch_size: Annotated[ int, click.option( - "--insert-batch-size", - type=int, + "--memorydb-pipeline-batch-size", + "pipeline_batch_size", + type=click.IntRange(min=1), default=10, - help="Batch size for inserting data. Adjust this as needed, but don't make it too big", + show_default=True, + help="Commands buffered in each MemoryDB pipeline execution", ), ] @@ -82,7 +84,7 @@ def MemoryDB(**parameters: Unpack[MemoryDBHNSWTypedDict]): M=parameters["m"], ef_construction=parameters["ef_construction"], ef_runtime=parameters["ef_runtime"], - insert_batch_size=parameters["insert_batch_size"], + pipeline_batch_size=parameters["pipeline_batch_size"], ), **parameters, ) diff --git a/vectordb_bench/backend/clients/memorydb/config.py b/vectordb_bench/backend/clients/memorydb/config.py index 2c40ff546..befa3c23e 100644 --- a/vectordb_bench/backend/clients/memorydb/config.py +++ b/vectordb_bench/backend/clients/memorydb/config.py @@ -1,4 +1,4 @@ -from pydantic import BaseModel, SecretStr +from pydantic import AliasChoices, BaseModel, Field, PositiveInt, SecretStr from ..api import DBCaseConfig, DBConfig, IndexType, MetricType @@ -24,7 +24,10 @@ def to_dict(self) -> dict: class MemoryDBIndexConfig(BaseModel, DBCaseConfig): metric_type: MetricType | None = None - insert_batch_size: int | None = None + pipeline_batch_size: PositiveInt = Field( + default=10, + validation_alias=AliasChoices("pipeline_batch_size", "insert_batch_size"), + ) def parse_metric(self) -> str: if self.metric_type == MetricType.L2: diff --git a/vectordb_bench/backend/clients/memorydb/memorydb.py b/vectordb_bench/backend/clients/memorydb/memorydb.py index 7e7a8650b..d816d0fd4 100644 --- a/vectordb_bench/backend/clients/memorydb/memorydb.py +++ b/vectordb_bench/backend/clients/memorydb/memorydb.py @@ -32,7 +32,7 @@ def __init__( self.case_config = db_case_config self.collection_name = INDEX_NAME self.target_nodes = RedisCluster.RANDOM if not self.db_config["cmd"] else None - self.insert_batch_size = db_case_config.insert_batch_size + self.pipeline_batch_size = db_case_config.pipeline_batch_size self.dbsize = kwargs.get("num_rows") # Create a MemoryDB connection, if db has password configured, add it to the connection here and in init(): @@ -190,7 +190,7 @@ def insert_embeddings( }, ) # Execute the pipe so we don't keep too much in memory at once - if (i + 1) % self.insert_batch_size == 0: + if (i + 1) % self.pipeline_batch_size == 0: pipe.execute() pipe.execute() diff --git a/vectordb_bench/backend/dataset.py b/vectordb_bench/backend/dataset.py index 6a6f262cc..35a11a20a 100644 --- a/vectordb_bench/backend/dataset.py +++ b/vectordb_bench/backend/dataset.py @@ -31,6 +31,7 @@ from .filter import Filter, FilterOp, non_filter log = logging.getLogger(__name__) +DEFAULT_INSERT_BATCH_SIZE = config.DEFAULT_INSERT_BATCH_SIZE class SizeLabel(NamedTuple): @@ -414,7 +415,10 @@ def _read_file(self, file_name: str) -> pl.DataFrame: class DataSetIterator: - def __init__(self, dataset: DatasetManager, batch_size: int = config.NUM_PER_BATCH): + def __init__(self, dataset: DatasetManager, batch_size: int = DEFAULT_INSERT_BATCH_SIZE): + if batch_size <= 0: + msg = f"insert batch size must be greater than 0, got {batch_size}" + raise ValueError(msg) self._ds = dataset self._batch_size = batch_size self._idx = 0 # file number @@ -880,7 +884,7 @@ def prepare( log.info(f"FTS dataset preparation completed: {self.data.full_name}") return True - def iter_batches(self, batch_size: int = config.NUM_PER_BATCH): + def iter_batches(self, batch_size: int = DEFAULT_INSERT_BATCH_SIZE): """Return an iterator for streaming FTS document batches.""" return FtsDocumentIterator(self, batch_size=batch_size) @@ -907,7 +911,10 @@ class FtsDocumentIterator: processing of large datasets. """ - def __init__(self, dataset: FtsDatasetManager, batch_size: int = config.NUM_PER_BATCH): + def __init__(self, dataset: FtsDatasetManager, batch_size: int = DEFAULT_INSERT_BATCH_SIZE): + if batch_size <= 0: + msg = f"insert batch size must be greater than 0, got {batch_size}" + raise ValueError(msg) self._ds = dataset self._batch_size = batch_size self._finished = False diff --git a/vectordb_bench/backend/runner/concurrent_runner.py b/vectordb_bench/backend/runner/concurrent_runner.py index 194fbf53c..6e33fe9cc 100644 --- a/vectordb_bench/backend/runner/concurrent_runner.py +++ b/vectordb_bench/backend/runner/concurrent_runner.py @@ -33,6 +33,7 @@ from .executor import TaskExecutor log = logging.getLogger(__name__) +DEFAULT_INSERT_BATCH_SIZE = config.DEFAULT_INSERT_BATCH_SIZE class ExecutorBackend(StrEnum): @@ -65,12 +66,15 @@ def __init__( timeout: float | None = None, max_workers: int | None = None, backend: ExecutorBackend = ExecutorBackend.THREADING, - batch_size: int = config.NUM_PER_BATCH, + batch_size: int = DEFAULT_INSERT_BATCH_SIZE, duration: float | None = None, with_scalar_labels: bool = False, tenant_case=None, # noqa: ANN001 workload_kind: WorkloadKind = WorkloadKind.VECTOR, ): + if batch_size <= 0: + msg = f"insert batch size must be greater than 0, got {batch_size}" + raise ValueError(msg) self.timeout = timeout if isinstance(timeout, int | float) else None self.dataset: DatasetManager | FtsDatasetManager = dataset self.db = db diff --git a/vectordb_bench/backend/runner/mp_runner.py b/vectordb_bench/backend/runner/mp_runner.py index bf81d2a7e..b5a7fe17d 100644 --- a/vectordb_bench/backend/runner/mp_runner.py +++ b/vectordb_bench/backend/runner/mp_runner.py @@ -19,7 +19,6 @@ from ...models import ConcurrencySlotTimeoutError from ..clients import api -NUM_PER_BATCH = config.NUM_PER_BATCH log = logging.getLogger(__name__) # HDR Histogram constants diff --git a/vectordb_bench/backend/runner/rate_runner.py b/vectordb_bench/backend/runner/rate_runner.py index 2387abfcb..26a8b5649 100644 --- a/vectordb_bench/backend/runner/rate_runner.py +++ b/vectordb_bench/backend/runner/rate_runner.py @@ -13,6 +13,7 @@ from .util import get_data log = logging.getLogger(__name__) +DEFAULT_INSERT_BATCH_SIZE = config.DEFAULT_INSERT_BATCH_SIZE class RatedMultiThreadingInsertRunner: @@ -23,13 +24,25 @@ def __init__( dataset_iter: DataSetIterator, normalize: bool = False, timeout: float | None = None, + batch_size: int = DEFAULT_INSERT_BATCH_SIZE, ): + if batch_size <= 0: + msg = f"insert batch size must be greater than 0, got {batch_size}" + raise ValueError(msg) + if rate <= 0: + msg = f"insert rate must be greater than 0, got {rate}" + raise ValueError(msg) + if rate % batch_size != 0: + msg = f"insert rate {rate} must be divisible by insert batch size {batch_size}" + raise ValueError(msg) + self.timeout = timeout if isinstance(timeout, int | float) else None self.dataset = dataset_iter self.db = db self.normalize = normalize self.insert_rate = rate - self.batch_rate = rate // config.NUM_PER_BATCH + self.batch_size = batch_size + self.batch_rate = rate // batch_size self.executing_futures = [] self.sig_idx = 0 diff --git a/vectordb_bench/backend/runner/read_write_runner.py b/vectordb_bench/backend/runner/read_write_runner.py index d3d1df2fa..8293128fe 100644 --- a/vectordb_bench/backend/runner/read_write_runner.py +++ b/vectordb_bench/backend/runner/read_write_runner.py @@ -8,6 +8,7 @@ import numpy as np +from vectordb_bench import config from vectordb_bench.backend.clients import api from vectordb_bench.backend.dataset import DatasetManager from vectordb_bench.backend.filter import Filter, non_filter @@ -19,6 +20,7 @@ from .serial_runner import SerialSearchRunner log = logging.getLogger(__name__) +DEFAULT_INSERT_BATCH_SIZE = config.DEFAULT_INSERT_BATCH_SIZE class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunner): @@ -41,6 +43,7 @@ def __init__( optimize_after_write: bool = True, read_dur_after_write: int = 300, # seconds, search duration when insertion is done timeout: float | None = None, + batch_size: int = DEFAULT_INSERT_BATCH_SIZE, ): self.insert_rate = insert_rate self.data_volume = dataset.data.size @@ -75,8 +78,9 @@ def __init__( self, rate=insert_rate, db=db, - dataset_iter=iter(dataset), + dataset_iter=dataset.iter_batches(batch_size), normalize=normalize, + batch_size=batch_size, ) self.serial_search_runner = SerialSearchRunner( db=db, diff --git a/vectordb_bench/backend/runner/serial_runner.py b/vectordb_bench/backend/runner/serial_runner.py index fc46aed23..646872c33 100644 --- a/vectordb_bench/backend/runner/serial_runner.py +++ b/vectordb_bench/backend/runner/serial_runner.py @@ -19,8 +19,8 @@ from .. import utils from ..clients import api -NUM_PER_BATCH = config.NUM_PER_BATCH LOAD_MAX_TRY_COUNT = config.LOAD_MAX_TRY_COUNT +DEFAULT_INSERT_BATCH_SIZE = config.DEFAULT_INSERT_BATCH_SIZE log = logging.getLogger(__name__) @@ -36,29 +36,34 @@ def __init__( normalize: bool, filters: Filter = non_filter, timeout: float | None = None, + batch_size: int = DEFAULT_INSERT_BATCH_SIZE, ): + if batch_size <= 0: + msg = f"insert batch size must be greater than 0, got {batch_size}" + raise ValueError(msg) self.timeout = timeout if isinstance(timeout, int | float) else None self.dataset = dataset self.db = db self.normalize = normalize self.filters = filters + self.batch_size = batch_size def endless_insert_data(self, all_embeddings: list, all_metadata: list, left_id: int = 0) -> int: with self.db.init(): # unique id for endlessness insertion all_metadata = [i + left_id for i in all_metadata] - num_batches = math.ceil(len(all_embeddings) / NUM_PER_BATCH) + num_batches = math.ceil(len(all_embeddings) / self.batch_size) log.info( f"({mp.current_process().name:16}) Start inserting {len(all_embeddings)} " - f"embeddings in batch {NUM_PER_BATCH}" + f"embeddings in batch {self.batch_size}" ) count = 0 for batch_id in range(num_batches): retry_count = 0 already_insert_count = 0 - metadata = all_metadata[batch_id * NUM_PER_BATCH : (batch_id + 1) * NUM_PER_BATCH] - embeddings = all_embeddings[batch_id * NUM_PER_BATCH : (batch_id + 1) * NUM_PER_BATCH] + metadata = all_metadata[batch_id * self.batch_size : (batch_id + 1) * self.batch_size] + embeddings = all_embeddings[batch_id * self.batch_size : (batch_id + 1) * self.batch_size] log.debug( f"({mp.current_process().name:16}) batch [{batch_id:3}/{num_batches}], " @@ -88,7 +93,7 @@ def endless_insert_data(self, all_embeddings: list, all_metadata: list, left_id: count += already_insert_count log.info( f"({mp.current_process().name:16}) Finish inserting {len(all_embeddings)} embeddings in " - f"batch {NUM_PER_BATCH}" + f"batch {self.batch_size}" ) return count @@ -96,7 +101,7 @@ def run_endlessness(self) -> int: """run forever util DB raises exception or crash""" # datasets for load tests are quite small, can fit into memory # only 1 file - data_df = next(iter(self.dataset)) + data_df = next(self.dataset.iter_batches(self.batch_size)) all_embeddings, all_metadata = ( np.stack(data_df[self.dataset.data.train_vector_field]).tolist(), data_df[self.dataset.data.train_id_field].tolist(), diff --git a/vectordb_bench/backend/task_runner.py b/vectordb_bench/backend/task_runner.py index 6f9025d4a..8bf8adc76 100644 --- a/vectordb_bench/backend/task_runner.py +++ b/vectordb_bench/backend/task_runner.py @@ -9,7 +9,6 @@ import numpy as np from pydantic import PrivateAttr -from .. import config from ..base import BaseModel from ..metric import Metric from ..models import PerformanceTimeoutError, TaskConfig, TaskStage @@ -85,6 +84,7 @@ def load_reuse_key(self) -> tuple | None: self._db_case_config_hash_key(), self._collection_name_hash_key(), self._dataset_hash_key(), + self.config.insert_batch_size, self.ca.with_scalar_labels, self.ca.is_multitenant, self._multitenant_routing_hash_key(), @@ -317,6 +317,7 @@ def _run_capacity_case(self) -> Metric: self.normalize, self.ca.filters, self.ca.load_timeout, + batch_size=self.config.insert_batch_size, ) count = runner.run_endlessness() except Exception as e: @@ -346,7 +347,7 @@ def _run_perf_case(self, drop_old: bool = True) -> Metric: m.load_duration = round(load_dur + build_dur, 4) m.additional_parameters.update( { - "num_per_batch": config.NUM_PER_BATCH, + "insert_batch_size": self.config.insert_batch_size, "load_concurrency": self.config.load_concurrency, } ) @@ -416,7 +417,7 @@ def _run_cloud_insert_case(self) -> Metric: self.normalize, self.ca.filters, max_workers=self.config.load_concurrency or None, - batch_size=self.ca.batch_size, + batch_size=self.config.insert_batch_size, duration=self.ca.duration, **runner_kwargs, ) @@ -519,6 +520,7 @@ def _load_train_data(self): self.ca.filters, self.ca.load_timeout, max_workers=self.config.load_concurrency or None, + batch_size=self.config.insert_batch_size, with_scalar_labels=self.ca.with_scalar_labels, workload_kind=self.workload_kind, **runner_kwargs, @@ -680,6 +682,7 @@ def _init_read_write_runner(self): concurrencies=ca.concurrencies, k=self.config.case_config.k, normalize=self.normalize, + batch_size=self.config.insert_batch_size, ) def stop(self): diff --git a/vectordb_bench/cli/cli.py b/vectordb_bench/cli/cli.py index d7854be6c..aaa87530b 100644 --- a/vectordb_bench/cli/cli.py +++ b/vectordb_bench/cli/cli.py @@ -209,6 +209,13 @@ def get_custom_case_config(parameters: dict) -> dict: "dataset_with_size_type": dataset_with_size_type, "label_percentage": parameters["label_percentage"], } + elif parameters["case_type"] in { + "StreamingPerformanceCase", + "StreamingCustomDataset", + }: + custom_case_config = { + "insert_rate": parameters["streaming_insert_rate"], + } elif parameters["case_type"] == "CloudPayloadSearchCase": custom_case_config = { "payload_profile": parameters["payload_profile"], @@ -228,7 +235,6 @@ def get_custom_case_config(parameters: dict) -> dict: copy_if_not_none(custom_case_config, parameters, "cloud_label_percentage", "label_percentage") elif parameters["case_type"] == "CloudInsertCase": custom_case_config = { - "batch_size": parameters["cloud_insert_batch_size"], "duration": parameters["cloud_insert_duration"], "dataset_with_size_type": dataset_with_size_type, } @@ -320,6 +326,26 @@ class CommonTypedDict(TypedDict): help="Number of concurrent workers for data loading in performance cases (0 = cpu_count)", ), ] + insert_batch_size: Annotated[ + int, + click.option( + "--insert-batch-size", + type=click.IntRange(min=1), + default=config.DEFAULT_INSERT_BATCH_SIZE, + show_default=True, + help="Rows or documents in each logical VDBBench insert batch; backends may split it further", + ), + ] + streaming_insert_rate: Annotated[ + int, + click.option( + "--streaming-insert-rate", + type=click.IntRange(min=1), + default=config.DEFAULT_STREAMING_INSERT_RATE, + show_default=True, + help="Rows inserted per second for StreamingPerformanceCase", + ), + ] search_serial: Annotated[ bool, click.option( @@ -584,16 +610,6 @@ class CommonTypedDict(TypedDict): help="Number of serial queries per cold/warm pass for CloudColdLatencyCase", ), ] - cloud_insert_batch_size: Annotated[ - int, - click.option( - "--cloud-insert-batch-size", - type=int, - default=5000, - show_default=True, - help="Insert batch size for CloudInsertCase", - ), - ] cloud_insert_duration: Annotated[ float | None, click.option( @@ -840,6 +856,7 @@ def run( parameters["search_concurrent"], ), load_concurrency=parameters["load_concurrency"], + insert_batch_size=parameters["insert_batch_size"], ) task_label = parameters["task_label"] diff --git a/vectordb_bench/frontend/components/run_test/generateTasks.py b/vectordb_bench/frontend/components/run_test/generateTasks.py index 5d848bb94..e78726381 100644 --- a/vectordb_bench/frontend/components/run_test/generateTasks.py +++ b/vectordb_bench/frontend/components/run_test/generateTasks.py @@ -3,7 +3,13 @@ from vectordb_bench.models import CaseConfig, CaseConfigParamType, TaskConfig -def generate_tasks(activedDbList: list[DB], dbConfigs, activedCaseList: list[CaseConfig], allCaseConfigs): +def generate_tasks( + activedDbList: list[DB], + dbConfigs, + activedCaseList: list[CaseConfig], + allCaseConfigs, + batch_size: int, +): tasks = [] for db in activedDbList: for case in activedCaseList: @@ -35,6 +41,7 @@ def generate_tasks(activedDbList: list[DB], dbConfigs, activedCaseList: list[Cas db_config=dbConfigs[db], case_config=case, db_case_config=db_case_config, + insert_batch_size=batch_size, ) tasks.append(task) diff --git a/vectordb_bench/frontend/components/run_test/runSettings.py b/vectordb_bench/frontend/components/run_test/runSettings.py new file mode 100644 index 000000000..ea459af5e --- /dev/null +++ b/vectordb_bench/frontend/components/run_test/runSettings.py @@ -0,0 +1,59 @@ +from vectordb_bench import config +from vectordb_bench.backend.cases import CaseType +from vectordb_bench.models import CaseConfig + +DEFAULT_INSERT_BATCH_SIZE = config.DEFAULT_INSERT_BATCH_SIZE +DEFAULT_STREAMING_INSERT_RATE = config.DEFAULT_STREAMING_INSERT_RATE +MAX_STREAMLIT_INT = (1 << 53) - 1 +STREAMING_CASE_TYPES = { + CaseType.StreamingPerformanceCase, + CaseType.StreamingCustomDataset, +} + + +def validate_streaming_insert_rates( + activedCaseList: list[CaseConfig], + batch_size: int, +) -> tuple[bool, list[str]]: + errors = [] + for case_config in activedCaseList: + if case_config.case_id not in STREAMING_CASE_TYPES: + continue + + custom_case = case_config.custom_case or {} + insert_rate = custom_case.get("insert_rate", DEFAULT_STREAMING_INSERT_RATE) + case_name = case_config.case_id.name + if insert_rate < batch_size: + errors.append( + f"{case_name}: Streaming Insert Rate ({insert_rate}) must be greater than or equal to " + f"Insert Batch Size ({batch_size})." + ) + elif insert_rate % batch_size != 0: + errors.append( + f"{case_name}: Streaming Insert Rate ({insert_rate}) must be divisible by " + f"Insert Batch Size ({batch_size})." + ) + + return len(errors) == 0, errors + + +def runSettings(container, activedCaseList: list[CaseConfig]) -> tuple[int, bool]: + container.markdown( + "
", + unsafe_allow_html=True, + ) + container.subheader("Run Settings") + batch_size = container.number_input( + "Insert Batch Size", + min_value=1, + max_value=MAX_STREAMLIT_INT, + value=DEFAULT_INSERT_BATCH_SIZE, + step=100, + help="Rows or documents in each logical VDBBench insert batch. Backends may split it further.", + ) + + is_valid, errors = validate_streaming_insert_rates(activedCaseList, batch_size) + for error in errors: + container.error(error) + + return batch_size, is_valid diff --git a/vectordb_bench/frontend/config/dbCaseConfigs.py b/vectordb_bench/frontend/config/dbCaseConfigs.py index 76d02935c..063ab7e02 100644 --- a/vectordb_bench/frontend/config/dbCaseConfigs.py +++ b/vectordb_bench/frontend/config/dbCaseConfigs.py @@ -240,9 +240,10 @@ def generate_custom_streaming_case() -> CaseConfig: ), ConfigInput( label=CaseConfigParamType.insert_rate, + displayLabel="Streaming Insert Rate", inputType=InputType.Number, - inputConfig=dict(step=100, min=100, max=4_000, value=200), - inputHelp="fixed insertion rate (rows/s), must be divisible by 100", + inputConfig=dict(step=100, min=100, max=MAX_STREAMLIT_INT, value=500), + inputHelp="Fixed streaming insertion rate (rows/s); must be at least and divisible by Insert Batch Size.", ), ConfigInput( label=CaseConfigParamType.search_stages, diff --git a/vectordb_bench/frontend/pages/run_test.py b/vectordb_bench/frontend/pages/run_test.py index 64115ff17..6a995f48e 100644 --- a/vectordb_bench/frontend/pages/run_test.py +++ b/vectordb_bench/frontend/pages/run_test.py @@ -5,6 +5,7 @@ from vectordb_bench.frontend.components.run_test.generateTasks import generate_tasks from vectordb_bench.frontend.components.run_test.hideSidebar import hideSidebar from vectordb_bench.frontend.components.run_test.initStyle import initStyle +from vectordb_bench.frontend.components.run_test.runSettings import runSettings from vectordb_bench.frontend.components.run_test.submitTask import submitTask from vectordb_bench.frontend.components.check_results.nav import NavToResults, NavToPages from vectordb_bench.frontend.components.check_results.headerIcon import drawHeaderIcon @@ -46,8 +47,13 @@ def main(): caseSelectorContainer = st.container() activedCaseList, allCaseConfigs = caseSelector(caseSelectorContainer, activedDbList) + # run settings + runSettingsContainer = st.container() + batch_size, areRunSettingsValid = runSettings(runSettingsContainer, activedCaseList) + isAllValid = isAllValid and areRunSettingsValid + # generate tasks - tasks = generate_tasks(activedDbList, dbConfigs, activedCaseList, allCaseConfigs) if isAllValid else [] + tasks = generate_tasks(activedDbList, dbConfigs, activedCaseList, allCaseConfigs, batch_size) if isAllValid else [] # submit submitContainer = st.container() diff --git a/vectordb_bench/models.py b/vectordb_bench/models.py index 855268cd1..f3eb24eeb 100644 --- a/vectordb_bench/models.py +++ b/vectordb_bench/models.py @@ -6,6 +6,7 @@ from typing import Any, ClassVar, Self import ujson +from pydantic import PositiveInt, model_validator from vectordb_bench.backend.cases import type2case from vectordb_bench.backend.dataset import DatasetWithSizeMap @@ -257,6 +258,38 @@ class TaskConfig(BaseModel): case_config: CaseConfig stages: list[TaskStage] = ALL_TASK_STAGES load_concurrency: int = config.LOAD_CONCURRENCY + insert_batch_size: PositiveInt = config.DEFAULT_INSERT_BATCH_SIZE + + @model_validator(mode="after") + def validate_streaming_insert_rate(self) -> Self: + streaming_case_types = { + CaseType.StreamingPerformanceCase, + CaseType.StreamingCustomDataset, + } + if self.case_config.case_id not in streaming_case_types: + return self + + custom_case = self.case_config.custom_case or {} + insert_rate = custom_case.get("insert_rate", config.DEFAULT_STREAMING_INSERT_RATE) + if not isinstance(insert_rate, int) or isinstance(insert_rate, bool) or insert_rate <= 0: + raise ValueError("streaming insert_rate must be a positive integer") + + rate_is_too_low = insert_rate < self.insert_batch_size + rate_is_divisible = insert_rate % self.insert_batch_size == 0 + + if rate_is_too_low: + msg = ( + f"streaming insert_rate ({insert_rate}) must be greater than or equal to " + f"insert_batch_size ({self.insert_batch_size})" + ) + raise ValueError(msg) + if not rate_is_divisible: + msg = ( + f"streaming insert_rate ({insert_rate}) must be divisible by " + f"insert_batch_size ({self.insert_batch_size})" + ) + raise ValueError(msg) + return self @property def db_name(self): @@ -420,8 +453,23 @@ def read_file(cls, full_path: pathlib.Path, trans_unit: bool = False) -> Self: for case_result in test_result["results"]: task_config = case_result.get("task_config") case_config = task_config.get("case_config") + metrics = case_result.get("metrics") db = DB(task_config.get("db")) + if "insert_batch_size" not in task_config: + insert_batch_size = None + if CaseType(case_config.get("case_id")) == CaseType.CloudInsertCase: + custom_case = case_config.get("custom_case") or {} + insert_batch_size = custom_case.get("batch_size") + if insert_batch_size is None and metrics: + additional_parameters = metrics.get("additional_parameters") or {} + insert_batch_size = additional_parameters.get( + "insert_batch_size", + additional_parameters.get("num_per_batch"), + ) + if insert_batch_size is not None: + task_config["insert_batch_size"] = insert_batch_size + task_config["db_config"] = db.config_cls(**task_config["db_config"]) # Safely instantiate DBCaseConfig (fallback to EmptyDBCaseConfig on None) @@ -439,7 +487,6 @@ def read_file(cls, full_path: pathlib.Path, trans_unit: bool = False) -> Self: task_config["case_config"] = cls.get_case_config(case_config=case_config) case_result["task_config"] = task_config - metrics = case_result.get("metrics") if ( metrics and CaseType(case_config.get("case_id")) == CaseType.CloudColdLatencyCase diff --git a/vectordb_bench/restful/app.py b/vectordb_bench/restful/app.py index ad0336501..a7de6cd2f 100644 --- a/vectordb_bench/restful/app.py +++ b/vectordb_bench/restful/app.py @@ -1,5 +1,6 @@ from flask import Flask, jsonify, request +from vectordb_bench import config from vectordb_bench.backend.clients import DB from vectordb_bench.interface import benchmark_runner from vectordb_bench.models import ALL_TASK_STAGES, CaseConfig, TaskConfig, TaskStage @@ -82,6 +83,7 @@ def run(): case_config=case_config, db_case_config=db_case_config, stages=stages, + insert_batch_size=task.get("insert_batch_size", config.DEFAULT_INSERT_BATCH_SIZE), ) task_configs.append(task_config) except Exception as e: