369 lines
13 KiB
Python
369 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for InteractiveState."""
|
|
|
|
# Standard
|
|
import argparse
|
|
import json
|
|
|
|
# First Party
|
|
from lmcache.cli.commands.bench.engine_bench.interactive.state import (
|
|
InteractiveState,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_full_state(**overrides: object) -> InteractiveState:
|
|
"""Build a state with all required fields set."""
|
|
values: dict[str, object] = {
|
|
"engine_url": "http://localhost:8000",
|
|
"workload": "long-doc-qa",
|
|
"tokens_per_gb_kvcache": 6553,
|
|
}
|
|
values.update(overrides)
|
|
state = InteractiveState()
|
|
for k, v in values.items():
|
|
state.set(k, v)
|
|
return state
|
|
|
|
|
|
def _make_cli_args(**overrides: object) -> argparse.Namespace:
|
|
"""Build a Namespace mimicking bench engine CLI output."""
|
|
defaults: dict[str, object] = {
|
|
"engine_url": None,
|
|
"workload": None,
|
|
"model": None,
|
|
"lmcache_url": None,
|
|
"tokens_per_gb_kvcache": None,
|
|
"kv_cache_volume": 100.0,
|
|
"seed": 42,
|
|
"output_dir": ".",
|
|
"no_csv": False,
|
|
"json": False,
|
|
"quiet": False,
|
|
"ldqa_document_length": 10000,
|
|
"ldqa_query_per_document": 2,
|
|
"ldqa_shuffle_policy": "random",
|
|
"ldqa_num_inflight_requests": 3,
|
|
"mrc_shared_prompt_length": 2000,
|
|
"mrc_chat_history_length": 10000,
|
|
"mrc_user_input_length": 50,
|
|
"mrc_output_length": 200,
|
|
"mrc_qps": 1.0,
|
|
"mrc_duration": 60.0,
|
|
"rp_request_length": 10000,
|
|
"rp_num_requests": 50,
|
|
"format": None,
|
|
"output": None,
|
|
"config": None,
|
|
"bench_target": "engine",
|
|
}
|
|
defaults.update(overrides)
|
|
return argparse.Namespace(**defaults)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Basic accessors
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBasicAccessors:
|
|
def test_empty_state(self) -> None:
|
|
state = InteractiveState()
|
|
assert not state.is_set("engine_url")
|
|
assert state.get("engine_url") is None
|
|
assert state.get("engine_url", "fallback") == "fallback"
|
|
|
|
def test_set_and_get(self) -> None:
|
|
state = InteractiveState()
|
|
state.set("engine_url", "http://localhost:8000")
|
|
assert state.is_set("engine_url")
|
|
assert state.get("engine_url") == "http://localhost:8000"
|
|
|
|
def test_values_copy(self) -> None:
|
|
state = _make_full_state()
|
|
vals = state.values
|
|
vals["engine_url"] = "mutated"
|
|
assert state.get("engine_url") == "http://localhost:8000"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Readiness
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestReadiness:
|
|
def test_empty_not_ready(self) -> None:
|
|
state = InteractiveState()
|
|
assert not state.is_ready()
|
|
|
|
def test_full_state_ready(self) -> None:
|
|
state = _make_full_state()
|
|
assert state.is_ready()
|
|
|
|
def test_missing_engine_url(self) -> None:
|
|
state = _make_full_state()
|
|
state._values.pop("engine_url")
|
|
assert not state.is_ready()
|
|
|
|
def test_missing_workload(self) -> None:
|
|
state = _make_full_state()
|
|
state._values.pop("workload")
|
|
assert not state.is_ready()
|
|
|
|
def test_tokens_per_gb_not_required_when_lmcache_set(self) -> None:
|
|
state = InteractiveState()
|
|
state.set("engine_url", "http://localhost:8000")
|
|
state.set("workload", "long-doc-qa")
|
|
state.set("lmcache_url", "http://localhost:8080")
|
|
# tokens_per_gb_kvcache not set, but lmcache_url is → condition
|
|
# for tokens_per_gb is not met → not required
|
|
assert state.is_ready()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Missing required
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestMissingRequired:
|
|
def test_all_missing(self) -> None:
|
|
state = InteractiveState()
|
|
missing = state.get_missing_required()
|
|
keys = [item.key for item in missing]
|
|
assert "engine_url" in keys
|
|
assert "workload" in keys
|
|
|
|
def test_none_missing(self) -> None:
|
|
state = _make_full_state()
|
|
assert state.get_missing_required() == []
|
|
|
|
def test_tokens_per_gb_skipped_with_lmcache(self) -> None:
|
|
state = InteractiveState()
|
|
state.set("engine_url", "http://localhost:8000")
|
|
state.set("workload", "long-doc-qa")
|
|
state.set("lmcache_url", "http://localhost:8080")
|
|
missing = state.get_missing_required()
|
|
keys = [item.key for item in missing]
|
|
assert "tokens_per_gb_kvcache" not in keys
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Workload items
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWorkloadItems:
|
|
def test_long_doc_qa(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
items = state.get_workload_items()
|
|
keys = [item.key for item in items]
|
|
assert "ldqa_document_length" in keys
|
|
assert "mrc_qps" not in keys
|
|
assert "rp_num_requests" not in keys
|
|
|
|
def test_multi_round_chat(self) -> None:
|
|
state = _make_full_state(workload="multi-round-chat")
|
|
items = state.get_workload_items()
|
|
keys = [item.key for item in items]
|
|
assert "mrc_shared_prompt_length" in keys
|
|
assert "ldqa_document_length" not in keys
|
|
|
|
def test_random_prefill(self) -> None:
|
|
state = _make_full_state(workload="random-prefill")
|
|
items = state.get_workload_items()
|
|
keys = [item.key for item in items]
|
|
assert "rp_request_length" in keys
|
|
assert "ldqa_document_length" not in keys
|
|
|
|
def test_workload_items_all_default_initially(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
assert state.workload_items_all_default()
|
|
|
|
def test_workload_items_not_default_after_set(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
state.set("ldqa_document_length", 5000)
|
|
assert not state.workload_items_all_default()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fill defaults
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFillDefaults:
|
|
def test_fills_kv_cache_volume(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
state.fill_defaults()
|
|
assert state.get("kv_cache_volume") == 100.0
|
|
|
|
def test_fills_workload_defaults(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
state.fill_defaults()
|
|
assert state.get("ldqa_document_length") == 10000
|
|
assert state.get("ldqa_query_per_document") == 2
|
|
|
|
def test_does_not_overwrite_set_values(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
state.set("ldqa_document_length", 5000)
|
|
state.fill_defaults()
|
|
assert state.get("ldqa_document_length") == 5000
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# from_cli_args
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFromCliArgs:
|
|
def test_picks_up_explicit_values(self) -> None:
|
|
args = _make_cli_args(
|
|
engine_url="http://host:8000",
|
|
workload="long-doc-qa",
|
|
tokens_per_gb_kvcache=6553,
|
|
)
|
|
state = InteractiveState.from_cli_args(args)
|
|
assert state.get("engine_url") == "http://host:8000"
|
|
assert state.get("workload") == "long-doc-qa"
|
|
assert state.get("tokens_per_gb_kvcache") == 6553
|
|
|
|
def test_none_values_not_set(self) -> None:
|
|
args = _make_cli_args() # all defaults
|
|
state = InteractiveState.from_cli_args(args)
|
|
assert not state.is_set("engine_url")
|
|
assert not state.is_set("workload")
|
|
assert not state.is_set("model")
|
|
|
|
def test_default_values_not_set(self) -> None:
|
|
# kv_cache_volume defaults to 100.0 — should not be marked as set
|
|
args = _make_cli_args(kv_cache_volume=100.0)
|
|
state = InteractiveState.from_cli_args(args)
|
|
assert not state.is_set("kv_cache_volume")
|
|
|
|
def test_non_default_values_set(self) -> None:
|
|
args = _make_cli_args(kv_cache_volume=50.0)
|
|
state = InteractiveState.from_cli_args(args)
|
|
assert state.is_set("kv_cache_volume")
|
|
assert state.get("kv_cache_volume") == 50.0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# to_namespace
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestToNamespace:
|
|
def test_round_trip(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
ns = state.to_namespace()
|
|
assert ns.engine_url == "http://localhost:8000"
|
|
assert ns.workload == "long-doc-qa"
|
|
assert ns.tokens_per_gb_kvcache == 6553
|
|
# Defaults filled
|
|
assert ns.kv_cache_volume == 100.0
|
|
assert ns.ldqa_document_length == 10000
|
|
assert ns.seed == 42
|
|
|
|
def test_has_bench_target(self) -> None:
|
|
state = _make_full_state()
|
|
ns = state.to_namespace()
|
|
assert ns.bench_target == "engine"
|
|
|
|
def test_has_output_attrs(self) -> None:
|
|
state = _make_full_state()
|
|
ns = state.to_namespace()
|
|
assert hasattr(ns, "output_dir")
|
|
assert hasattr(ns, "no_csv")
|
|
assert hasattr(ns, "quiet")
|
|
assert hasattr(ns, "format")
|
|
assert hasattr(ns, "output")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# JSON round-trip
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestJsonRoundTrip:
|
|
def test_to_json_and_back(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
state.set("ldqa_document_length", 5000)
|
|
data = state.to_json()
|
|
# engine_url excluded from export (environment-specific)
|
|
assert "engine_url" not in data
|
|
assert data["ldqa_document_length"] == 5000
|
|
restored = InteractiveState.from_json(data)
|
|
assert restored.get("ldqa_document_length") == 5000
|
|
|
|
def test_save_and_load(self, tmp_path: object) -> None:
|
|
# Standard
|
|
import pathlib
|
|
|
|
path = str(pathlib.Path(str(tmp_path)) / "test_config.json")
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
state.save_json(path)
|
|
|
|
loaded = InteractiveState.load_json(path)
|
|
assert "engine_url" not in loaded.values
|
|
assert loaded.get("workload") == "long-doc-qa"
|
|
|
|
def test_json_file_is_valid_json(self, tmp_path: object) -> None:
|
|
# Standard
|
|
import pathlib
|
|
|
|
path = str(pathlib.Path(str(tmp_path)) / "test.json")
|
|
state = _make_full_state()
|
|
state.save_json(path)
|
|
|
|
with open(path) as f:
|
|
data = json.load(f)
|
|
assert isinstance(data, dict)
|
|
assert "engine_url" not in data
|
|
assert data["workload"] == "long-doc-qa"
|
|
|
|
def test_merge_cli_args(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
args = _make_cli_args(
|
|
engine_url="http://override:9000",
|
|
)
|
|
state.merge_cli_args(args)
|
|
assert state.get("engine_url") == "http://override:9000"
|
|
# Original values preserved
|
|
assert state.get("workload") == "long-doc-qa"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Summary
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSummary:
|
|
def test_summary_lines_non_empty(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
state.fill_defaults()
|
|
lines = state.summary_lines()
|
|
assert len(lines) > 0
|
|
|
|
def test_summary_includes_engine_url(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
state.fill_defaults()
|
|
labels = [label for label, _ in state.summary_lines()]
|
|
assert "Engine URL" in labels
|
|
|
|
def test_summary_shows_auto_detect_for_empty_model(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
state.fill_defaults()
|
|
for label, value in state.summary_lines():
|
|
if label == "Model name":
|
|
assert value == "(auto-detect)"
|
|
break
|
|
|
|
def test_summary_excludes_wrong_workload_items(self) -> None:
|
|
state = _make_full_state(workload="long-doc-qa")
|
|
state.fill_defaults()
|
|
labels = [label for label, _ in state.summary_lines()]
|
|
assert "Queries per second" not in labels # mrc item
|
|
assert "Number of requests" not in labels # rp item
|