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

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

@@ -0,0 +1,352 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Centralized config item definitions for interactive configuration.
+
+Each ``ConfigItem`` declaratively describes one configurable parameter:
+its key, display name, description, input type, default, and when it
+should be shown.  The ``ALL_ITEMS`` list is the single source of truth
+for descriptions, ordering, and defaults.
+"""
+
+# Standard
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+# ---------------------------------------------------------------------------
+# Phases
+# ---------------------------------------------------------------------------
+
+PHASE_REQUIRED = 1
+PHASE_GENERAL = 2
+PHASE_WORKLOAD = 3
+
+
+# ---------------------------------------------------------------------------
+# ConfigItem
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ConfigItem:
+    """Declarative description of a single configurable parameter.
+
+    Attributes:
+        key: State dict key (matches argparse attr name, e.g., ``"engine_url"``).
+        display_name: Heading shown in the prompt.
+        description: One-sentence explanation shown below the heading.
+        input_type: One of ``"text"``, ``"int"``, ``"float"``, ``"bool"``,
+            ``"choice"``.
+        default: Default value.  ``None`` means required (no default).
+        required: If True, this item must have a value before the benchmark
+            can start.
+        choices: For ``"choice"`` type — list of ``(value, description)`` tuples.
+        condition: Callable ``(state_dict) -> bool`` that determines whether
+            this item should be shown.  ``None`` means always shown.
+        phase: Which interactive phase this item belongs to.
+    """
+
+    key: str
+    display_name: str
+    description: str
+    input_type: str  # "text", "int", "float", "bool", "choice"
+    default: Any = None
+    required: bool = False
+    choices: list[tuple[str, str]] = field(default_factory=list)
+    condition: Callable[[dict[str, Any]], bool] | None = None
+    phase: int = PHASE_GENERAL
+
+
+# ---------------------------------------------------------------------------
+# Condition helpers
+# ---------------------------------------------------------------------------
+
+
+def _has_lmcache(state: dict[str, Any]) -> bool:
+    """Show this item only when the user said they have LMCache."""
+    return bool(state.get("has_lmcache"))
+
+
+def _no_lmcache_url(state: dict[str, Any]) -> bool:
+    """Show this item only when lmcache_url is not set."""
+    return not state.get("lmcache_url")
+
+
+def _workload_is(name: str) -> Callable[[dict[str, Any]], bool]:
+    """Return a condition that checks the workload value."""
+
+    def check(state: dict[str, Any]) -> bool:
+        return state.get("workload") == name
+
+    return check
+
+
+# ---------------------------------------------------------------------------
+# ALL_ITEMS — the centralized registry
+# ---------------------------------------------------------------------------
+
+ALL_ITEMS: list[ConfigItem] = [
+    # ── Phase 1: Required ─────────────────────────────────────────────
+    ConfigItem(
+        key="engine_url",
+        display_name="Engine URL",
+        description=(
+            "URL of the inference engine. "
+            "Set OPENAI_API_KEY env var if authentication is needed."
+        ),
+        input_type="text",
+        default="http://localhost:8000",
+        required=True,
+        phase=PHASE_REQUIRED,
+    ),
+    ConfigItem(
+        key="workload",
+        display_name="Workload",
+        description="The type of benchmark workload to run.",
+        input_type="choice",
+        default=None,
+        required=True,
+        choices=[
+            (
+                "long-doc-permutator",
+                "Query the same set of long documents with different orders",
+            ),
+            ("long-doc-qa", "Repeated Q&A over long documents (tests KV cache reuse)"),
+            ("multi-round-chat", "Multi-turn chat with stateful sessions"),
+            ("random-prefill", "Prefill-only requests fired simultaneously"),
+        ],
+        phase=PHASE_REQUIRED,
+    ),
+    ConfigItem(
+        key="has_lmcache",
+        display_name="LMCache Server",
+        description=(
+            "Do you have a running LMCache server? "
+            "It can auto-detect KV cache size information."
+        ),
+        input_type="bool",
+        default=True,
+        required=False,
+        phase=PHASE_REQUIRED,
+    ),
+    ConfigItem(
+        key="lmcache_url",
+        display_name="LMCache Server URL",
+        description="URL of the running LMCache HTTP server.",
+        input_type="text",
+        default="http://localhost:8080",
+        required=False,
+        condition=_has_lmcache,
+        phase=PHASE_REQUIRED,
+    ),
+    ConfigItem(
+        key="tokens_per_gb_kvcache",
+        display_name="Tokens per GB KV cache",
+        description=(
+            "How many tokens fit in 1 GB of KV cache for your model.\n"
+            "  If using vLLM, look for these lines in the startup log:\n"
+            '    "Available KV cache memory: XX.XX GiB"\n'
+            '    "GPU KV cache size: XXX,XXX tokens"\n'
+            "  Then compute: tokens_per_gb = "
+            "GPU_KV_cache_tokens / Available_KV_cache_GiB"
+        ),
+        input_type="int",
+        default=None,
+        required=True,
+        condition=_no_lmcache_url,
+        phase=PHASE_REQUIRED,
+    ),
+    # ── Phase 2: General ──────────────────────────────────────────────
+    ConfigItem(
+        key="model",
+        display_name="Model name",
+        description=(
+            "The model served by the engine. "
+            "Leave empty to auto-detect from the engine."
+        ),
+        input_type="text",
+        default="",
+        phase=PHASE_GENERAL,
+    ),
+    ConfigItem(
+        key="kv_cache_volume",
+        display_name="KV cache volume (GB)",
+        description="Target active KV cache size for the benchmark.",
+        input_type="float",
+        default=100.0,
+        phase=PHASE_GENERAL,
+    ),
+    # ── Phase 3: long-doc-permutator ─────────────────────────────────
+    ConfigItem(
+        key="ldp_num_contexts",
+        display_name="Number of contexts",
+        description="Number of unique context documents to generate.",
+        input_type="int",
+        default=5,
+        condition=_workload_is("long-doc-permutator"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="ldp_context_length",
+        display_name="Context length (tokens)",
+        description="Token length of each context document.",
+        input_type="int",
+        default=5000,
+        condition=_workload_is("long-doc-permutator"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="ldp_system_prompt_length",
+        display_name="System prompt length (tokens)",
+        description="Token length of the shared system prompt. Use 0 for none.",
+        input_type="int",
+        default=1000,
+        condition=_workload_is("long-doc-permutator"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="ldp_num_permutations",
+        display_name="Number of permutations",
+        description="Distinct permutations to send. Capped at N! (N = num_contexts).",
+        input_type="int",
+        default=10,
+        condition=_workload_is("long-doc-permutator"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="ldp_num_inflight_requests",
+        display_name="Max inflight requests",
+        description="Maximum concurrent in-flight requests.",
+        input_type="int",
+        default=1,
+        condition=_workload_is("long-doc-permutator"),
+        phase=PHASE_WORKLOAD,
+    ),
+    # ── Phase 3: long-doc-qa ──────────────────────────────────────────
+    ConfigItem(
+        key="ldqa_document_length",
+        display_name="Document length (tokens)",
+        description="Token length of each synthetic document.",
+        input_type="int",
+        default=10000,
+        condition=_workload_is("long-doc-qa"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="ldqa_query_per_document",
+        display_name="Queries per document",
+        description="Number of questions asked per document.",
+        input_type="int",
+        default=2,
+        condition=_workload_is("long-doc-qa"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="ldqa_shuffle_policy",
+        display_name="Shuffle policy",
+        description="How benchmark requests are ordered.",
+        input_type="choice",
+        default="random",
+        choices=[
+            ("random", "Shuffle all (doc, query) pairs randomly"),
+            ("tile", "Process queries round by round across all documents"),
+        ],
+        condition=_workload_is("long-doc-qa"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="ldqa_num_inflight_requests",
+        display_name="Max inflight requests",
+        description="Maximum concurrent in-flight requests.",
+        input_type="int",
+        default=3,
+        condition=_workload_is("long-doc-qa"),
+        phase=PHASE_WORKLOAD,
+    ),
+    # ── Phase 3: multi-round-chat ─────────────────────────────────────
+    ConfigItem(
+        key="mrc_shared_prompt_length",
+        display_name="System prompt length (tokens)",
+        description="Token length of the system prompt per session.",
+        input_type="int",
+        default=2000,
+        condition=_workload_is("multi-round-chat"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="mrc_chat_history_length",
+        display_name="Chat history length (tokens)",
+        description="Token length of pre-filled conversation history.",
+        input_type="int",
+        default=10000,
+        condition=_workload_is("multi-round-chat"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="mrc_user_input_length",
+        display_name="User input length (tokens)",
+        description="Tokens per user query in each round.",
+        input_type="int",
+        default=50,
+        condition=_workload_is("multi-round-chat"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="mrc_output_length",
+        display_name="Output length (tokens)",
+        description="Max tokens to generate per response.",
+        input_type="int",
+        default=200,
+        condition=_workload_is("multi-round-chat"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="mrc_qps",
+        display_name="Queries per second",
+        description="Target request dispatch rate.",
+        input_type="float",
+        default=1.0,
+        condition=_workload_is("multi-round-chat"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="mrc_duration",
+        display_name="Duration (seconds)",
+        description="How long the benchmark runs.",
+        input_type="float",
+        default=60.0,
+        condition=_workload_is("multi-round-chat"),
+        phase=PHASE_WORKLOAD,
+    ),
+    # ── Phase 3: random-prefill ───────────────────────────────────────
+    ConfigItem(
+        key="rp_request_length",
+        display_name="Request length (tokens)",
+        description="Token length of each prefill request.",
+        input_type="int",
+        default=10000,
+        condition=_workload_is("random-prefill"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="rp_num_requests",
+        display_name="Number of requests",
+        description="Total prefill requests to fire simultaneously.",
+        input_type="int",
+        default=50,
+        condition=_workload_is("random-prefill"),
+        phase=PHASE_WORKLOAD,
+    ),
+]
+
+
+def get_items_by_phase(phase: int) -> list[ConfigItem]:
+    """Return all items belonging to a given phase."""
+    return [item for item in ALL_ITEMS if item.phase == phase]
+
+
+def get_item(key: str) -> ConfigItem:
+    """Look up a ConfigItem by key.  Raises KeyError if not found."""
+    for item in ALL_ITEMS:
+        if item.key == key:
+            return item
+    raise KeyError(f"No ConfigItem with key {key!r}")

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

@@ -0,0 +1,327 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Intermediate state tracker for interactive configuration.
+
+``InteractiveState`` holds the partially-configured benchmark parameters,
+can be initialized from CLI args or a saved JSON file, and can be
+converted to the ``argparse.Namespace`` that ``_bench_engine()`` expects.
+"""
+
+# Standard
+from typing import Any
+import argparse
+import json
+
+# First Party
+from lmcache.cli.commands.bench.engine_bench.interactive.schema import (
+    ALL_ITEMS,
+    PHASE_GENERAL,
+    PHASE_REQUIRED,
+    PHASE_WORKLOAD,
+    ConfigItem,
+)
+
+# Keys that exist on argparse.Namespace but are NOT part of the interactive
+# config item registry (operational flags, handled separately).
+_OUTPUT_KEYS = ("output_dir", "seed", "no_csv", "export_csv", "json", "quiet")
+
+# Keys used only during the interactive flow, never serialized or
+# passed to the orchestrator.
+_INTERACTIVE_ONLY_KEYS = {"has_lmcache"}
+
+# Keys excluded from exported JSON configs.  These are either
+# environment-specific (engine_url, lmcache_url) or interactive-only.
+_EXPORT_EXCLUDED_KEYS = _INTERACTIVE_ONLY_KEYS | {"engine_url", "lmcache_url"}
+
+# Mapping from ConfigItem.key to the argparse attribute name when they differ.
+# Most keys match directly; these are the exceptions.
+_KEY_TO_ATTR: dict[str, str] = {
+    "kv_cache_volume": "kv_cache_volume",
+    "tokens_per_gb_kvcache": "tokens_per_gb_kvcache",
+}
+
+# argparse attribute names where the CLI default is None (meaning "not set"),
+# versus attributes where a non-None argparse default is the real default
+# (e.g., kv_cache_volume defaults to 100.0).
+_ARGPARSE_NONE_MEANS_UNSET = {
+    "engine_url",
+    "workload",
+    "model",
+    "lmcache_url",
+    "tokens_per_gb_kvcache",
+}
+
+
+class InteractiveState:
+    """Tracks which config items have been set and their values.
+
+    Keys present in ``_values`` are considered "set".  Missing keys are
+    "unset" and will either be prompted for or filled with defaults.
+    """
+
+    def __init__(self) -> None:
+        self._values: dict[str, Any] = {}
+
+    # ------------------------------------------------------------------
+    # Basic accessors
+    # ------------------------------------------------------------------
+
+    def is_set(self, key: str) -> bool:
+        return key in self._values
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return self._values.get(key, default)
+
+    def set(self, key: str, value: Any) -> None:
+        self._values[key] = value
+
+    @property
+    def values(self) -> dict[str, Any]:
+        return dict(self._values)
+
+    # ------------------------------------------------------------------
+    # Readiness checks
+    # ------------------------------------------------------------------
+
+    def is_ready(self) -> bool:
+        """True when all required items (whose conditions are met) have values."""
+        for item in ALL_ITEMS:
+            if not item.required:
+                continue
+            if not self._condition_met(item):
+                continue
+            if not self.is_set(item.key):
+                return False
+        return True
+
+    def get_missing_required(self) -> list[ConfigItem]:
+        """Return phase-1 items that still need user input.
+
+        Includes both required items that are unset, and non-required
+        phase-1 items (like ``lmcache_url``) that are relevant because
+        a downstream required item (``tokens_per_gb_kvcache``) is unset.
+        """
+        missing: list[ConfigItem] = []
+        for item in ALL_ITEMS:
+            if item.phase != PHASE_REQUIRED:
+                continue
+            if not self._condition_met(item):
+                continue
+            if self.is_set(item.key):
+                continue
+            if item.required:
+                missing.append(item)
+            elif item.key in ("has_lmcache", "lmcache_url") and not self.is_set(
+                "tokens_per_gb_kvcache"
+            ):
+                # Only ask about LMCache when tokens_per_gb is needed
+                missing.append(item)
+        return missing
+
+    def get_general_items(self) -> list[ConfigItem]:
+        """Return phase-2 (general) items that are not yet set."""
+        return [
+            item
+            for item in ALL_ITEMS
+            if item.phase == PHASE_GENERAL
+            and not self.is_set(item.key)
+            and self._condition_met(item)
+        ]
+
+    def get_workload_items(self) -> list[ConfigItem]:
+        """Return phase-3 (workload-specific) items whose conditions are met."""
+        return [
+            item
+            for item in ALL_ITEMS
+            if item.phase == PHASE_WORKLOAD and self._condition_met(item)
+        ]
+
+    def has_unconfigured_general(self) -> bool:
+        """True if there are general items the user hasn't explicitly set."""
+        return len(self.get_general_items()) > 0
+
+    def has_workload_items(self) -> bool:
+        """True if there are workload-specific items to configure."""
+        return len(self.get_workload_items()) > 0
+
+    def workload_items_all_default(self) -> bool:
+        """True if no workload-specific items have been explicitly set."""
+        for item in ALL_ITEMS:
+            if item.phase != PHASE_WORKLOAD:
+                continue
+            if not self._condition_met(item):
+                continue
+            if self.is_set(item.key):
+                return False
+        return True
+
+    # ------------------------------------------------------------------
+    # Defaults
+    # ------------------------------------------------------------------
+
+    def fill_defaults(self) -> None:
+        """Set all unset items (whose conditions are met) to their defaults."""
+        for item in ALL_ITEMS:
+            if self.is_set(item.key):
+                continue
+            if not self._condition_met(item):
+                continue
+            if item.default is not None:
+                self._values[item.key] = item.default
+
+    # ------------------------------------------------------------------
+    # Conversion: CLI args ↔ InteractiveState
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "InteractiveState":
+        """Build state from parsed CLI arguments.
+
+        Only sets values that the user explicitly provided (i.e., not the
+        argparse default).  This lets us distinguish "user set
+        ``--kv-cache-volume 100``" from "user didn't touch it".
+        """
+        state = cls()
+        for item in ALL_ITEMS:
+            if item.key in _INTERACTIVE_ONLY_KEYS:
+                continue
+            attr = _KEY_TO_ATTR.get(item.key, item.key)
+            value = getattr(args, attr, None)
+            if value is None:
+                continue
+            # For keys where argparse default is None, any non-None value
+            # means the user set it.
+            if attr in _ARGPARSE_NONE_MEANS_UNSET:
+                state._values[item.key] = value
+                continue
+            # For keys with real argparse defaults, we can't easily tell
+            # if the user typed --kv-cache-volume 100 vs it being the
+            # default.  We mark it as "set" only if it differs from the
+            # schema default.  This is imperfect but good enough — the
+            # worst case is we re-prompt for a value the user explicitly
+            # set to the default.
+            if item.default is not None and value == item.default:
+                continue
+            state._values[item.key] = value
+
+        # Derive has_lmcache from lmcache_url if provided via CLI
+        if state.is_set("lmcache_url"):
+            state._values["has_lmcache"] = True
+
+        return state
+
+    def to_namespace(self) -> argparse.Namespace:
+        """Convert to an ``argparse.Namespace`` compatible with ``_bench_engine``.
+
+        Fills defaults for any unset items, then builds the namespace
+        with the attribute names that ``parse_args_to_config()`` and
+        ``create_workload()`` expect.
+        """
+        self.fill_defaults()
+        ns = argparse.Namespace()
+
+        # Map state keys to namespace attributes
+        for item in ALL_ITEMS:
+            if item.key in _INTERACTIVE_ONLY_KEYS:
+                continue
+            attr = _KEY_TO_ATTR.get(item.key, item.key)
+            # Only fall back to schema default if the item's condition is met.
+            # This prevents lmcache_url's default from leaking when
+            # has_lmcache is not set.
+            if item.key in self._values:
+                value = self._values[item.key]
+            elif self._condition_met(item):
+                value = item.default
+            else:
+                value = None
+            setattr(ns, attr, value)
+
+        # Output settings (not in the interactive registry)
+        ns.output_dir = self._values.get("output_dir", ".")
+        ns.seed = self._values.get("seed", 42)
+        ns.no_csv = self._values.get("no_csv", False)
+        ns.json = self._values.get("export_json", False)
+        ns.quiet = self._values.get("quiet", False)
+        ns.bench_target = "engine"
+
+        # Ensure format/output attrs exist for create_metrics
+        if not hasattr(ns, "format"):
+            ns.format = None
+        if not hasattr(ns, "output"):
+            ns.output = None
+
+        return ns
+
+    # ------------------------------------------------------------------
+    # Conversion: JSON ↔ InteractiveState
+    # ------------------------------------------------------------------
+
+    def to_json(self) -> dict[str, Any]:
+        """Serialize to a JSON-compatible dict for config export.
+
+        Excludes environment-specific keys (``engine_url``,
+        ``lmcache_url``) and interactive-only keys so the exported
+        config is portable and works without an LMCache server.
+        """
+        self.fill_defaults()
+        return {k: v for k, v in self._values.items() if k not in _EXPORT_EXCLUDED_KEYS}
+
+    @classmethod
+    def from_json(cls, data: dict[str, Any]) -> "InteractiveState":
+        """Load from a saved config JSON dict."""
+        state = cls()
+        for key, value in data.items():
+            state._values[key] = value
+        return state
+
+    def save_json(self, path: str) -> None:
+        """Export the current state to a JSON file."""
+        with open(path, "w") as f:
+            json.dump(self.to_json(), f, indent=2)
+            f.write("\n")
+
+    @classmethod
+    def load_json(cls, path: str) -> "InteractiveState":
+        """Load state from a JSON config file."""
+        with open(path) as f:
+            data = json.load(f)
+        return cls.from_json(data)
+
+    def merge_cli_args(self, args: argparse.Namespace) -> None:
+        """Merge CLI args on top of existing state (CLI args win)."""
+        cli_state = InteractiveState.from_cli_args(args)
+        self._values.update(cli_state._values)
+
+    # ------------------------------------------------------------------
+    # Summary
+    # ------------------------------------------------------------------
+
+    def summary_lines(self) -> list[tuple[str, str]]:
+        """Return ``(label, value_str)`` pairs for the config summary."""
+        lines: list[tuple[str, str]] = []
+        for item in ALL_ITEMS:
+            if item.key in _INTERACTIVE_ONLY_KEYS:
+                continue
+            if not self._condition_met(item):
+                continue
+            value = self._values.get(item.key, item.default)
+            if value is None or value == "":
+                if item.key == "model":
+                    display = "(auto-detect)"
+                elif item.key == "lmcache_url":
+                    continue  # skip empty lmcache_url
+                else:
+                    display = "(not set)"
+            else:
+                display = str(value)
+            lines.append((item.display_name, display))
+        return lines
+
+    # ------------------------------------------------------------------
+    # Internal
+    # ------------------------------------------------------------------
+
+    def _condition_met(self, item: ConfigItem) -> bool:
+        """Check whether an item's condition is satisfied."""
+        if item.condition is None:
+            return True
+        return item.condition(self._values)