feat: add insert batch and streaming rate controls

.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"

README.md

@@ -466,7 +466,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,105,110,115,120,125,130 \
@@ -520,12 +520,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:
@@ -546,8 +546,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
 ```
 

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 \

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")),

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,15 +120,13 @@ 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,
         "dataset_with_size_type": DatasetWithSizeType.CohereMedium.value,
     }
 
     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()

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()