feat(cli): expose insert-rate, search-stages, streaming-concurrencies for StreamingPerformanceCase

tests/test_streaming_cli.py

@@ -0,0 +1,74 @@
+from click.testing import CliRunner
+from pytest import MonkeyPatch
+
+from vectordb_bench.backend.cases import CaseType
+from vectordb_bench.backend.clients.test import cli as test_cli
+from vectordb_bench.cli import cli as common_cli
+
+
+def invoke_test_command(monkeypatch: MonkeyPatch, args: list[str]):
+    captured = {}
+
+    def fake_run(tasks, task_label):
+        captured["task"] = tasks[0]
+        captured["task_label"] = task_label
+
+    monkeypatch.setattr(common_cli.benchmark_runner, "run", fake_run)
+    monkeypatch.setattr(common_cli.benchmark_runner, "has_running", lambda: False)
+    result = CliRunner().invoke(test_cli.Test, args)
+    return result, captured
+
+
+def test_common_cli_exposes_streaming_options() -> None:
+    result = CliRunner().invoke(test_cli.Test, ["--help"])
+
+    assert result.exit_code == 0, result.output
+    assert "--insert-rate INTEGER" in result.output
+    assert "--search-stages TEXT" in result.output
+    assert "--streaming-concurrencies TEXT" in result.output
+
+
+def test_streaming_case_type_accepted_with_defaults(monkeypatch: MonkeyPatch) -> None:
+    result, captured = invoke_test_command(
+        monkeypatch,
+        ["--case-type", "StreamingPerformanceCase"],
+    )
+
+    assert result.exit_code == 0, result.output
+    case_config = captured["task"].case_config
+    assert case_config.case_id == CaseType.StreamingPerformanceCase
+    assert case_config.custom_case["insert_rate"] == 500
+    assert case_config.custom_case["search_stages"] == [0.5, 0.8]
+    assert case_config.custom_case["concurrencies"] == [5, 10]
+
+
+def test_streaming_case_type_accepts_overrides(monkeypatch: MonkeyPatch) -> None:
+    result, captured = invoke_test_command(
+        monkeypatch,
+        [
+            "--case-type",
+            "StreamingPerformanceCase",
+            "--insert-rate",
+            "1000",
+            "--search-stages",
+            "0.3,0.6,0.9",
+            "--streaming-concurrencies",
+            "1,2,4",
+        ],
+    )
+
+    assert result.exit_code == 0, result.output
+    custom_case = captured["task"].case_config.custom_case
+    assert custom_case["insert_rate"] == 1000
+    assert custom_case["search_stages"] == [0.3, 0.6, 0.9]
+    assert custom_case["concurrencies"] == [1, 2, 4]
+
+
+def test_non_streaming_case_type_ignores_streaming_options(monkeypatch: MonkeyPatch) -> None:
+    result, captured = invoke_test_command(
+        monkeypatch,
+        ["--case-type", "Performance1536D50K", "--insert-rate", "999"],
+    )
+
+    assert result.exit_code == 0, result.output
+    assert captured["task"].case_config.custom_case == {}

vectordb_bench/__init__.py

@@ -32,6 +32,10 @@ class config:
 
     NUM_CONCURRENCY = env.list("NUM_CONCURRENCY", [1, 5, 10, 20, 30, 40, 60, 80], subcast=int)
 
+    STREAMING_INSERT_RATE = env.int("STREAMING_INSERT_RATE", 500)
+    STREAMING_SEARCH_STAGES = env.list("STREAMING_SEARCH_STAGES", [0.5, 0.8], subcast=float)
+    STREAMING_CONCURRENCIES = env.list("STREAMING_CONCURRENCIES", [5, 10], subcast=int)
+
     CONCURRENCY_DURATION = 30
 
     CONCURRENCY_TIMEOUT = 3600

vectordb_bench/cli/cli.py

@@ -233,6 +233,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"] == "StreamingPerformanceCase":
+        custom_case_config = {
+            "dataset_with_size_type": dataset_with_size_type,
+            "insert_rate": parameters["insert_rate"],
+            "search_stages": parameters["search_stages"],
+            "concurrencies": parameters["streaming_concurrencies"],
+        }
     elif parameters["case_type"] == "CloudPayloadSearchCase":
         custom_case_config = {
             "payload_profile": parameters["payload_profile"],
@@ -384,6 +391,39 @@ class CommonTypedDict(TypedDict):
             help="Number of concurrent workers for data loading in performance cases (0 = cpu_count)",
         ),
     ]
+    insert_rate: Annotated[
+        int,
+        click.option(
+            "--insert-rate",
+            type=int,
+            default=config.STREAMING_INSERT_RATE,
+            show_default=True,
+            help="Rows inserted per second for StreamingPerformanceCase",
+        ),
+    ]
+    search_stages: Annotated[
+        list[str],
+        click.option(
+            "--search-stages",
+            type=str,
+            help="Comma-separated stream-completion fractions (0-1) at which search is sampled, "
+            "for StreamingPerformanceCase",
+            show_default=True,
+            default=",".join(map(str, config.STREAMING_SEARCH_STAGES)),
+            callback=lambda *args: list(map(float, click_arg_split(*args))),
+        ),
+    ]
+    streaming_concurrencies: Annotated[
+        list[str],
+        click.option(
+            "--streaming-concurrencies",
+            type=str,
+            help="Comma-separated concurrency levels for the per-stage search sweep, for StreamingPerformanceCase",
+            show_default=True,
+            default=",".join(map(str, config.STREAMING_CONCURRENCIES)),
+            callback=lambda *args: list(map(int, click_arg_split(*args))),
+        ),
+    ]
     search_serial: Annotated[
         bool,
         click.option(