wehub-resource-sync
341 lines
12 KiB
Python
341 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for bench engine config module."""
|
|
|
|
# Standard
|
|
from unittest.mock import MagicMock, patch
|
|
import argparse
|
|
|
|
# Third Party
|
|
import pytest
|
|
|
|
# First Party
|
|
from lmcache.cli.commands.bench.engine_bench.config import (
|
|
EngineBenchConfig,
|
|
_find_model_meta,
|
|
auto_detect_model,
|
|
parse_args_to_config,
|
|
resolve_tokens_per_gb,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def base_namespace() -> argparse.Namespace:
|
|
"""Namespace with all required fields for parse_args_to_config."""
|
|
return argparse.Namespace(
|
|
engine_url="http://localhost:8000",
|
|
lmcache_url=None,
|
|
model="test-model",
|
|
workload="long-doc-qa",
|
|
kv_cache_volume=100.0,
|
|
tokens_per_gb_kvcache=50000,
|
|
seed=42,
|
|
output_dir=".",
|
|
no_csv=False,
|
|
json=False,
|
|
quiet=False,
|
|
ignore_eos=False,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# EngineBenchConfig
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestEngineBenchConfig:
|
|
def _make_config(self, **overrides) -> EngineBenchConfig:
|
|
defaults = dict(
|
|
engine_url="http://localhost:8000",
|
|
model="test-model",
|
|
workload="long-doc-qa",
|
|
kv_cache_volume_gb=100.0,
|
|
tokens_per_gb_kvcache=50000,
|
|
seed=42,
|
|
output_dir=".",
|
|
export_csv=True,
|
|
export_json=False,
|
|
quiet=False,
|
|
)
|
|
defaults.update(overrides)
|
|
return EngineBenchConfig(**defaults) # type: ignore[arg-type]
|
|
|
|
def test_valid_construction(self) -> None:
|
|
cfg = self._make_config()
|
|
assert cfg.engine_url == "http://localhost:8000"
|
|
assert cfg.model == "test-model"
|
|
assert cfg.workload == "long-doc-qa"
|
|
assert cfg.tokens_per_gb_kvcache == 50000
|
|
|
|
def test_ignore_eos_defaults_false_and_overridable(self) -> None:
|
|
assert self._make_config().ignore_eos is False
|
|
assert self._make_config(ignore_eos=True).ignore_eos is True
|
|
|
|
def test_empty_engine_url(self) -> None:
|
|
with pytest.raises(ValueError, match="engine_url must be non-empty"):
|
|
self._make_config(engine_url="")
|
|
|
|
def test_invalid_kv_cache_volume(self) -> None:
|
|
with pytest.raises(ValueError, match="kv_cache_volume_gb must be positive"):
|
|
self._make_config(kv_cache_volume_gb=0)
|
|
|
|
def test_invalid_tokens_per_gb(self) -> None:
|
|
with pytest.raises(ValueError, match="tokens_per_gb_kvcache must be positive"):
|
|
self._make_config(tokens_per_gb_kvcache=0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# auto_detect_model
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAutoDetectModel:
|
|
@patch(
|
|
"lmcache.cli.commands.bench.engine_bench.config.OpenAI",
|
|
)
|
|
def test_returns_first_model_id(self, mock_openai_cls) -> None:
|
|
mock_model = MagicMock()
|
|
mock_model.id = "Qwen/Qwen3-14B"
|
|
mock_client = MagicMock()
|
|
mock_client.models.list.return_value = MagicMock(data=[mock_model])
|
|
mock_openai_cls.return_value = mock_client
|
|
|
|
result = auto_detect_model("http://localhost:8000")
|
|
assert result == "Qwen/Qwen3-14B"
|
|
|
|
@patch(
|
|
"lmcache.cli.commands.bench.engine_bench.config.OpenAI",
|
|
)
|
|
def test_empty_data_raises(self, mock_openai_cls) -> None:
|
|
mock_client = MagicMock()
|
|
mock_client.models.list.return_value = MagicMock(data=[])
|
|
mock_openai_cls.return_value = mock_client
|
|
|
|
with pytest.raises(RuntimeError, match="No models returned"):
|
|
auto_detect_model("http://localhost:8000")
|
|
|
|
@patch(
|
|
"lmcache.cli.commands.bench.engine_bench.config.OpenAI",
|
|
)
|
|
def test_connection_error_raises(self, mock_openai_cls) -> None:
|
|
mock_openai_cls.side_effect = ConnectionError("refused")
|
|
|
|
with pytest.raises(RuntimeError, match="Failed to fetch models"):
|
|
auto_detect_model("http://localhost:8000")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# parse_args_to_config
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestParseArgsToConfig:
|
|
def test_basic_parsing(self, base_namespace) -> None:
|
|
cfg = parse_args_to_config(base_namespace)
|
|
assert cfg.engine_url == "http://localhost:8000"
|
|
assert cfg.model == "test-model"
|
|
assert cfg.workload == "long-doc-qa"
|
|
assert cfg.kv_cache_volume_gb == 100.0
|
|
assert cfg.seed == 42
|
|
assert cfg.export_csv is True
|
|
assert cfg.export_json is False
|
|
|
|
@patch(
|
|
"lmcache.cli.commands.bench.engine_bench.config.auto_detect_model",
|
|
return_value="auto-detected-model",
|
|
)
|
|
def test_model_auto_detect(self, mock_auto_detect, base_namespace) -> None:
|
|
base_namespace.model = None
|
|
cfg = parse_args_to_config(base_namespace)
|
|
assert cfg.model == "auto-detected-model"
|
|
mock_auto_detect.assert_called_once_with("http://localhost:8000")
|
|
|
|
def test_no_tokens_per_gb_no_lmcache_url_raises(
|
|
self,
|
|
base_namespace,
|
|
) -> None:
|
|
base_namespace.tokens_per_gb_kvcache = None
|
|
base_namespace.lmcache_url = None
|
|
with pytest.raises(ValueError, match="--tokens-per-gb-kvcache"):
|
|
parse_args_to_config(base_namespace)
|
|
|
|
@patch(
|
|
"lmcache.cli.commands.bench.engine_bench.config.resolve_tokens_per_gb",
|
|
return_value=6553,
|
|
)
|
|
def test_lmcache_url_resolves_tokens_per_gb(
|
|
self,
|
|
mock_resolve,
|
|
base_namespace,
|
|
) -> None:
|
|
base_namespace.tokens_per_gb_kvcache = None
|
|
base_namespace.lmcache_url = "http://localhost:8080"
|
|
cfg = parse_args_to_config(base_namespace)
|
|
assert cfg.tokens_per_gb_kvcache == 6553
|
|
mock_resolve.assert_called_once_with(
|
|
"http://localhost:8080",
|
|
"test-model",
|
|
)
|
|
|
|
def test_export_flags(self, base_namespace) -> None:
|
|
base_namespace.no_csv = True
|
|
base_namespace.json = True
|
|
cfg = parse_args_to_config(base_namespace)
|
|
assert cfg.export_csv is False
|
|
assert cfg.export_json is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _find_model_meta
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFindModelMeta:
|
|
def _gpu_meta(self) -> dict:
|
|
return {
|
|
"gpu_0": {
|
|
"model_name": "Qwen/Qwen3-14B",
|
|
"world_size": 1,
|
|
"kv_cache_layout": {"cache_size_per_token": 163840},
|
|
},
|
|
"gpu_1": {
|
|
"model_name": "meta-llama/Llama-3.1-70B",
|
|
"world_size": 4,
|
|
"kv_cache_layout": {"cache_size_per_token": 327680},
|
|
},
|
|
}
|
|
|
|
def test_finds_matching_model(self) -> None:
|
|
meta = _find_model_meta(self._gpu_meta(), "Qwen/Qwen3-14B")
|
|
assert meta["model_name"] == "Qwen/Qwen3-14B"
|
|
|
|
def test_finds_second_model(self) -> None:
|
|
meta = _find_model_meta(
|
|
self._gpu_meta(),
|
|
"meta-llama/Llama-3.1-70B",
|
|
)
|
|
assert meta["world_size"] == 4
|
|
|
|
def test_missing_model_raises(self) -> None:
|
|
with pytest.raises(RuntimeError, match="not found"):
|
|
_find_model_meta(self._gpu_meta(), "nonexistent-model")
|
|
|
|
def test_error_lists_available(self) -> None:
|
|
with pytest.raises(RuntimeError, match="Qwen/Qwen3-14B"):
|
|
_find_model_meta(self._gpu_meta(), "nonexistent")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# resolve_tokens_per_gb
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestResolveTokensPerGb:
|
|
def _status_response(
|
|
self,
|
|
cache_size_per_token: int = 163840,
|
|
world_size: int = 1,
|
|
model_name: str = "Qwen/Qwen3-14B",
|
|
) -> dict:
|
|
return {
|
|
"cache_context_meta": {
|
|
"gpu_0": {
|
|
"model_name": model_name,
|
|
"world_size": world_size,
|
|
"kv_cache_layout": {
|
|
"num_layers": 40,
|
|
"hidden_dim_size": 1024,
|
|
"dtype": "torch.bfloat16",
|
|
"cache_size_per_token": cache_size_per_token,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
@patch(
|
|
"lmcache.cli.commands.bench.engine_bench.config._fetch_lmcache_status",
|
|
)
|
|
def test_basic_calculation(self, mock_fetch) -> None:
|
|
# 163840 bytes/token, world_size=1
|
|
# 1 GB = 1073741824 bytes
|
|
# 1073741824 // 163840 = 6553
|
|
mock_fetch.return_value = self._status_response(
|
|
cache_size_per_token=163840,
|
|
world_size=1,
|
|
)
|
|
result = resolve_tokens_per_gb(
|
|
"http://localhost:8080",
|
|
"Qwen/Qwen3-14B",
|
|
)
|
|
assert result == 6553
|
|
|
|
@patch(
|
|
"lmcache.cli.commands.bench.engine_bench.config._fetch_lmcache_status",
|
|
)
|
|
def test_world_size_multiplier(self, mock_fetch) -> None:
|
|
# 163840 bytes/token (rank-local), world_size=4
|
|
# global = 163840 * 4 = 655360
|
|
# 1073741824 // 655360 = 1638
|
|
mock_fetch.return_value = self._status_response(
|
|
cache_size_per_token=163840,
|
|
world_size=4,
|
|
)
|
|
result = resolve_tokens_per_gb(
|
|
"http://localhost:8080",
|
|
"Qwen/Qwen3-14B",
|
|
)
|
|
assert result == 1638
|
|
|
|
@patch(
|
|
"lmcache.cli.commands.bench.engine_bench.config._fetch_lmcache_status",
|
|
)
|
|
def test_no_gpu_meta_raises(self, mock_fetch) -> None:
|
|
mock_fetch.return_value = {"cache_context_meta": {}}
|
|
with pytest.raises(RuntimeError, match="No model info"):
|
|
resolve_tokens_per_gb(
|
|
"http://localhost:8080",
|
|
"Qwen/Qwen3-14B",
|
|
)
|
|
|
|
@patch(
|
|
"lmcache.cli.commands.bench.engine_bench.config._fetch_lmcache_status",
|
|
)
|
|
def test_model_not_found_raises(self, mock_fetch) -> None:
|
|
mock_fetch.return_value = self._status_response()
|
|
with pytest.raises(RuntimeError, match="not found"):
|
|
resolve_tokens_per_gb(
|
|
"http://localhost:8080",
|
|
"nonexistent-model",
|
|
)
|
|
|
|
@patch(
|
|
"lmcache.cli.commands.bench.engine_bench.config._fetch_lmcache_status",
|
|
)
|
|
def test_no_cache_size_raises(self, mock_fetch) -> None:
|
|
data = self._status_response()
|
|
del data["cache_context_meta"]["gpu_0"]["kv_cache_layout"][
|
|
"cache_size_per_token"
|
|
]
|
|
mock_fetch.return_value = data
|
|
with pytest.raises(RuntimeError, match="cache_size_per_token"):
|
|
resolve_tokens_per_gb(
|
|
"http://localhost:8080",
|
|
"Qwen/Qwen3-14B",
|
|
)
|
|
|
|
@patch(
|
|
"lmcache.cli.commands.bench.engine_bench.config._fetch_lmcache_status",
|
|
)
|
|
def test_no_layout_raises(self, mock_fetch) -> None:
|
|
data = self._status_response()
|
|
del data["cache_context_meta"]["gpu_0"]["kv_cache_layout"]
|
|
mock_fetch.return_value = data
|
|
with pytest.raises(RuntimeError, match="kv_cache_layout"):
|
|
resolve_tokens_per_gb(
|
|
"http://localhost:8080",
|
|
"Qwen/Qwen3-14B",
|
|
)
|