diff --git a/tests/test_streaming_cli.py b/tests/test_streaming_cli.py new file mode 100644 index 000000000..1b46d8066 --- /dev/null +++ b/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 == {} diff --git a/vectordb_bench/__init__.py b/vectordb_bench/__init__.py index 1491630e0..d25fc3ba4 100644 --- a/vectordb_bench/__init__.py +++ b/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 diff --git a/vectordb_bench/cli/cli.py b/vectordb_bench/cli/cli.py index abcb91607..a0385d506 100644 --- a/vectordb_bench/cli/cli.py +++ b/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(