lmcache-cli 0.4.5.dev0__py3-none-any.whl

lmcache/cli/commands/bench/engine_bench/config.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Configuration dataclasses and CLI-arg parsing for ``lmcache bench engine``."""
+
+# Standard
+from dataclasses import dataclass
+import argparse
+import json
+import os
+import urllib.error
+import urllib.request
+
+# Third Party
+from openai import OpenAI
+
+# First Party
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+_GB = 1024**3
+
+
+@dataclass
+class EngineBenchConfig:
+    """Top-level config produced from CLI args, interactive mode, or saved config.
+
+    Contains only general benchmark parameters. Workload-specific configs
+    (e.g., ``LongDocQAConfig``) live in their respective workload modules
+    and are resolved by the workload factory.
+    """
+
+    engine_url: str
+    model: str
+    workload: str
+    kv_cache_volume_gb: float
+    tokens_per_gb_kvcache: int
+    seed: int
+    output_dir: str
+    export_csv: bool
+    export_json: bool
+    quiet: bool
+
+    def __post_init__(self) -> None:
+        if not self.engine_url:
+            raise ValueError("engine_url must be non-empty")
+        if self.kv_cache_volume_gb <= 0:
+            raise ValueError(
+                f"kv_cache_volume_gb must be positive, got {self.kv_cache_volume_gb}"
+            )
+        if self.tokens_per_gb_kvcache <= 0:
+            raise ValueError(
+                f"tokens_per_gb_kvcache must be positive, "
+                f"got {self.tokens_per_gb_kvcache}"
+            )
+
+
+def auto_detect_model(engine_url: str) -> str:
+    """Fetch the first model ID from the engine's ``/v1/models`` endpoint.
+
+    Args:
+        engine_url: Base URL of the inference engine (e.g.,
+            ``http://localhost:8000``).
+
+    Returns:
+        The model ID string.
+
+    Raises:
+        RuntimeError: If the engine is unreachable or returns no models.
+    """
+    base_url = engine_url.rstrip("/")
+    if not base_url.startswith(("http://", "https://")):
+        base_url = f"http://{base_url}"
+    if not base_url.endswith("/v1"):
+        base_url += "/v1"
+
+    api_key = os.getenv("OPENAI_API_KEY", "sk-dummy")
+    logger.debug("Auto-detecting model from %s/models", base_url)
+
+    try:
+        client = OpenAI(base_url=base_url, api_key=api_key)
+        models = client.models.list()
+    except Exception as e:
+        raise RuntimeError(f"Failed to fetch models from {base_url}/models: {e}") from e
+
+    if not models.data:
+        raise RuntimeError(
+            f"No models returned by {base_url}/models; pass --model explicitly."
+        )
+
+    model_id = models.data[0].id
+    logger.debug("Auto-detected model: %s", model_id)
+    return model_id
+
+
+def _fetch_lmcache_status(lmcache_url: str) -> dict:
+    """Fetch ``/api/status`` from the LMCache HTTP server.
+
+    Returns:
+        Parsed JSON response.
+
+    Raises:
+        RuntimeError: If the server is unreachable.
+    """
+    url = lmcache_url.rstrip("/")
+    if not url.startswith(("http://", "https://")):
+        url = f"http://{url}"
+    status_url = f"{url}/api/status"
+
+    logger.debug("Fetching LMCache status from %s", status_url)
+
+    try:
+        req = urllib.request.Request(status_url)
+        with urllib.request.urlopen(req, timeout=10) as resp:
+            return json.loads(resp.read().decode())
+    except (urllib.error.URLError, OSError) as e:
+        raise RuntimeError(
+            f"Cannot connect to LMCache server at {status_url}: {e}"
+        ) from e
+
+
+def _find_model_meta(
+    gpu_meta: dict,
+    model_name: str,
+) -> dict:
+    """Find the GPU metadata entry matching *model_name*.
+
+    Args:
+        gpu_meta: The ``gpu_context_meta`` dict from ``/api/status``.
+        model_name: Model name to match.
+
+    Returns:
+        The matching GPU metadata dict.
+
+    Raises:
+        RuntimeError: If no entry matches *model_name*.
+    """
+    for meta in gpu_meta.values():
+        if meta.get("model_name") == model_name:
+            return meta
+
+    available = sorted({m.get("model_name", "?") for m in gpu_meta.values()})
+    raise RuntimeError(
+        f"Model {model_name!r} not found on LMCache server. "
+        f"Available: {', '.join(available)}"
+    )
+
+
+def resolve_tokens_per_gb(lmcache_url: str, model_name: str) -> int:
+    """Query the LMCache server and compute tokens per GB of KV cache.
+
+    Fetches ``/api/status``, finds the model entry matching
+    *model_name*, and computes::
+
+        global_bytes_per_token = cache_size_per_token * world_size
+        tokens_per_gb = (1024**3) // global_bytes_per_token
+
+    ``cache_size_per_token`` is rank-local, so it must be multiplied
+    by ``world_size`` for tensor-parallel models.
+
+    Args:
+        lmcache_url: URL of the LMCache HTTP server.
+        model_name: Model name to look up (must match a model served
+            by the LMCache server).
+
+    Returns:
+        tokens_per_gb_kvcache value.
+
+    Raises:
+        RuntimeError: If the server is unreachable, the model is not
+            found, or the layout is missing required fields.
+    """
+    data = _fetch_lmcache_status(lmcache_url)
+
+    gpu_meta = data.get("gpu_context_meta", {})
+    if not gpu_meta:
+        raise RuntimeError(
+            "No model info returned by LMCache server; "
+            "is the server running with a model loaded?"
+        )
+
+    meta = _find_model_meta(gpu_meta, model_name)
+    layout = meta.get("kv_cache_layout")
+    if not layout:
+        raise RuntimeError(f"No kv_cache_layout for model {model_name!r}")
+
+    cache_size_per_token = layout.get("cache_size_per_token")
+    if cache_size_per_token is None:
+        raise RuntimeError(
+            f"cache_size_per_token not available for model "
+            f"{model_name!r}; is the LMCache server up to date?"
+        )
+
+    world_size = meta.get("world_size", 1)
+    global_bytes_per_token = cache_size_per_token * world_size
+    tokens_per_gb = _GB // global_bytes_per_token
+
+    logger.info(
+        "Resolved from LMCache: model=%s, "
+        "cache_size_per_token=%d bytes (rank-local), "
+        "world_size=%d -> %d bytes/token (global) -> %d tokens/GB",
+        model_name,
+        cache_size_per_token,
+        world_size,
+        global_bytes_per_token,
+        tokens_per_gb,
+    )
+    return tokens_per_gb
+
+
+def parse_args_to_config(args: argparse.Namespace) -> EngineBenchConfig:
+    """Convert parsed CLI arguments into a fully-resolved EngineBenchConfig.
+
+    Handles model auto-detection and tokens-per-GB resolution from the
+    LMCache server when ``--lmcache-url`` is provided.
+
+    Args:
+        args: Parsed argparse Namespace from the bench engine subcommand.
+
+    Returns:
+        A fully-resolved EngineBenchConfig.
+    """
+    model = args.model if args.model else auto_detect_model(args.engine_url)
+
+    tokens_per_gb = args.tokens_per_gb_kvcache
+    if tokens_per_gb is None:
+        lmcache_url = getattr(args, "lmcache_url", None)
+        if lmcache_url is not None:
+            tokens_per_gb = resolve_tokens_per_gb(lmcache_url, model)
+        else:
+            raise ValueError(
+                "--tokens-per-gb-kvcache is required when --lmcache-url is not set"
+            )
+
+    return EngineBenchConfig(
+        engine_url=args.engine_url,
+        model=model,
+        workload=args.workload,
+        kv_cache_volume_gb=args.kv_cache_volume,
+        tokens_per_gb_kvcache=tokens_per_gb,
+        seed=args.seed,
+        output_dir=args.output_dir,
+        export_csv=not args.no_csv,
+        export_json=args.json,
+        quiet=args.quiet,
+    )

lmcache/cli/commands/bench/engine_bench/interactive/__init__.py

@@ -0,0 +1,274 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Interactive configuration flow for ``lmcache bench engine``.
+
+Entry point: ``run_interactive(args)`` — walks the user through
+missing configuration, offers gates for optional settings, and
+returns a complete ``argparse.Namespace`` ready for the orchestrator.
+"""
+
+# Standard
+import argparse
+import sys
+
+# First Party
+from lmcache.cli.commands.bench.engine_bench.interactive.schema import (
+    ConfigItem,
+)
+from lmcache.cli.commands.bench.engine_bench.interactive.state import (
+    InteractiveState,
+)
+from lmcache.cli.commands.bench.engine_bench.interactive.terminal import (
+    BOLD,
+    CYAN,
+    RESET,
+    YELLOW,
+    prompt_bool,
+    prompt_choice,
+    prompt_number,
+    prompt_text,
+)
+
+__all__ = ["run_interactive"]
+
+
+# ---------------------------------------------------------------------------
+# Prompt dispatcher
+# ---------------------------------------------------------------------------
+
+
+def _prompt_for_item(item: ConfigItem) -> object:
+    """Prompt the user for a single config item based on its type."""
+    if item.input_type == "text":
+        return prompt_text(
+            item.display_name,
+            item.description,
+            default=item.default if item.default is not None else "",
+        )
+    if item.input_type == "int":
+        return prompt_number(
+            item.display_name,
+            item.description,
+            default=item.default,
+            number_type=int,
+        )
+    if item.input_type == "float":
+        return prompt_number(
+            item.display_name,
+            item.description,
+            default=item.default,
+            number_type=float,
+        )
+    if item.input_type == "bool":
+        return prompt_bool(
+            item.display_name,
+            item.description,
+            default=bool(item.default) if item.default is not None else True,
+        )
+    if item.input_type == "choice":
+        return prompt_choice(
+            item.display_name,
+            item.description,
+            choices=item.choices,
+            default=item.default if item.default is not None else "",
+        )
+    raise ValueError(f"Unknown input_type {item.input_type!r} for {item.key}")
+
+
+# ---------------------------------------------------------------------------
+# Gate prompt
+# ---------------------------------------------------------------------------
+
+
+def _prompt_gate(section_name: str, detail: str) -> bool:
+    """Ask the user whether to configure a section or skip with defaults.
+
+    Returns True if the user wants to configure.
+    """
+    return (
+        prompt_choice(
+            section_name,
+            f"Would you like to configure {detail}?\n"
+            f"  Defaults will be used if you skip.",
+            choices=[
+                ("use defaults", "Skip, use defaults"),
+                ("configure", "Yes, configure"),
+            ],
+            default="use defaults",
+        )
+        == "configure"
+    )
+
+
+# ---------------------------------------------------------------------------
+# Summary
+# ---------------------------------------------------------------------------
+
+
+def _print_summary(state: InteractiveState) -> None:
+    """Print a formatted configuration summary."""
+    print()
+    print(f"{BOLD}{'─' * 50}{RESET}")
+    print(f"{BOLD} Configuration Summary{RESET}")
+    print(f"{BOLD}{'─' * 50}{RESET}")
+    for label, value in state.summary_lines():
+        padding = max(0, 26 - len(label))
+        print(f"  {label}:{' ' * padding}{CYAN}{value}{RESET}")
+    print(f"{BOLD}{'─' * 50}{RESET}")
+
+
+# ---------------------------------------------------------------------------
+# Action prompt
+# ---------------------------------------------------------------------------
+
+
+def _prompt_action() -> str:
+    """Ask the user to start the benchmark or export config.
+
+    Returns ``"start"`` or ``"export"``.
+    """
+    return prompt_choice(
+        "What would you like to do?",
+        "",
+        choices=[
+            ("start", "Start benchmark"),
+            ("export", "Export configuration for later use and exit"),
+        ],
+        default="start",
+    )
+
+
+def _resolve_before_export(state: InteractiveState) -> None:
+    """Resolve tokens_per_gb_kvcache and model before exporting.
+
+    If the user provided an LMCache URL, query the server to get
+    ``tokens_per_gb_kvcache`` so the exported config is standalone.
+    If the model is empty and an engine URL is available, auto-detect it.
+    """
+    # First Party
+    from lmcache.cli.commands.bench.engine_bench.config import (
+        auto_detect_model,
+        resolve_tokens_per_gb,
+    )
+
+    engine_url = state.get("engine_url", "")
+    model = state.get("model", "")
+
+    # Auto-detect model if empty
+    if not model and engine_url:
+        try:
+            model = auto_detect_model(engine_url)
+            state.set("model", model)
+        except RuntimeError as e:
+            print(f"  {YELLOW}Warning: could not auto-detect model: {e}{RESET}")
+
+    # Resolve tokens_per_gb from LMCache if needed
+    lmcache_url = state.get("lmcache_url", "")
+    if lmcache_url and not state.is_set("tokens_per_gb_kvcache"):
+        try:
+            tokens = resolve_tokens_per_gb(lmcache_url, model)
+            state.set("tokens_per_gb_kvcache", tokens)
+        except RuntimeError as e:
+            print(
+                f"  {YELLOW}Warning: could not resolve "
+                f"tokens_per_gb from LMCache: {e}{RESET}"
+            )
+
+
+def _handle_export(state: InteractiveState) -> None:
+    """Prompt for filename, resolve values, save JSON, and exit."""
+    _resolve_before_export(state)
+    filename = prompt_text(
+        "Export filename",
+        "",
+        default="bench_config.json",
+    )
+    state.save_json(filename)
+    print()
+    print(f"  {CYAN}Saved to {filename}{RESET}")
+    print(
+        f"  {BOLD}Replay with:{RESET} "
+        f"{CYAN}lmcache bench engine "
+        f"--engine-url <URL> --config {filename}{RESET}"
+    )
+    print()
+    sys.exit(0)
+
+
+# ---------------------------------------------------------------------------
+# Main entry point
+# ---------------------------------------------------------------------------
+
+
+def run_interactive(args: argparse.Namespace) -> argparse.Namespace:
+    """Run the interactive configuration flow.
+
+    Walks the user through missing required items, offers gates for
+    general and workload-specific settings, shows a summary, and
+    returns a complete ``argparse.Namespace``.
+
+    Args:
+        args: Partially-populated CLI args (some may be None).
+
+    Returns:
+        A fully-populated ``argparse.Namespace`` ready for the
+        benchmark orchestrator.
+    """
+    state = InteractiveState.from_cli_args(args)
+
+    print()
+    print(f"{BOLD}{'═' * 50}{RESET}")
+    print(f"{BOLD} lmcache bench engine — Interactive Setup{RESET}")
+    print(f"{BOLD}{'═' * 50}{RESET}")
+
+    # ── Phase 1: Required items ───────────────────────────────────────
+    # Walk through missing required items one by one.  Re-evaluate the
+    # list after each prompt because setting one value (e.g., has_lmcache)
+    # can make new items eligible (e.g., lmcache_url) or skip others
+    # (e.g., tokens_per_gb_kvcache).
+    while True:
+        missing = state.get_missing_required()
+        if not missing:
+            break
+        item = missing[0]
+        value = _prompt_for_item(item)
+        state.set(item.key, value)
+
+    # ── Phase 2: General settings gate ────────────────────────────────
+    if state.has_unconfigured_general():
+        if _prompt_gate(
+            "General settings",
+            "general settings (model, KV cache volume, etc.)",
+        ):
+            for item in state.get_general_items():
+                value = _prompt_for_item(item)
+                state.set(item.key, value)
+
+    # ── Phase 3: Workload settings gate (always shown) ────────────────
+    if state.has_workload_items():
+        workload = state.get("workload", "workload")
+        if _prompt_gate(
+            f"Workload settings ({workload})",
+            "workload-specific settings",
+        ):
+            for item in state.get_workload_items():
+                value = _prompt_for_item(item)
+                state.set(item.key, value)
+
+    # Fill all remaining defaults
+    state.fill_defaults()
+
+    # ── Phase 4: Summary + action ─────────────────────────────────────
+    _print_summary(state)
+    action = _prompt_action()
+
+    if action == "export":
+        _handle_export(state)
+
+    # Carry over output settings from original CLI args
+    ns = state.to_namespace()
+    for attr in ("output_dir", "seed", "no_csv", "json", "quiet", "format", "output"):
+        cli_val = getattr(args, attr, None)
+        if cli_val is not None:
+            setattr(ns, attr, cli_val)
+
+    return ns

lmcache/cli/commands/bench/engine_bench/interactive/config.json

@@ -0,0 +1,10 @@
+{
+  "model": "Qwen/Qwen3-14B",
+  "workload": "long-doc-qa",
+  "kv_cache_volume": 100.0,
+  "tokens_per_gb_kvcache": 6553,
+  "ldqa_document_length": 10000,
+  "ldqa_query_per_document": 2,
+  "ldqa_shuffle_policy": "random",
+  "ldqa_num_inflight_requests": 3
+}