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

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

@@ -0,0 +1,291 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Terminal UI primitives for interactive configuration.
+
+Provides simple prompt functions for collecting user input:
+
+- ``prompt_text`` — free-form text with optional default
+- ``prompt_number`` — numeric input with type validation
+- ``prompt_bool`` — Y/N confirmation
+- ``prompt_choice`` — arrow-key selection (falls back to numbered list on non-TTY)
+"""
+
+# Standard
+import sys
+
+# ---------------------------------------------------------------------------
+# ANSI helpers
+# ---------------------------------------------------------------------------
+
+BOLD = "\033[1m"
+CYAN = "\033[96m"
+GREEN = "\033[92m"
+YELLOW = "\033[93m"
+DIM = "\033[2m"
+RESET = "\033[0m"
+CLEAR_LINE = "\033[2K"
+CURSOR_UP = "\033[A"
+
+
+def _is_tty() -> bool:
+    return hasattr(sys.stdin, "fileno") and sys.stdin.isatty()
+
+
+# ---------------------------------------------------------------------------
+# prompt_text
+# ---------------------------------------------------------------------------
+
+
+def prompt_text(
+    label: str,
+    description: str = "",
+    default: str = "",
+) -> str:
+    """Prompt for free-form text input.
+
+    Args:
+        label: The config item name shown as a heading.
+        description: One-line explanation shown below the label.
+        default: Default value; shown in brackets, returned on empty Enter.
+
+    Returns:
+        The user's input, or the default if they pressed Enter.
+    """
+    print()
+    print(f"{BOLD}{label}{RESET}")
+    if description:
+        print(f"  {description}")
+    if default:
+        prompt = f"  {DIM}[default: {default}]{RESET} {GREEN}>{RESET} "
+    else:
+        prompt = f"  {GREEN}>{RESET} "
+    value = input(prompt).strip()
+    return value if value else default
+
+
+# ---------------------------------------------------------------------------
+# prompt_number
+# ---------------------------------------------------------------------------
+
+
+def prompt_number(
+    label: str,
+    description: str = "",
+    default: float | int | None = None,
+    number_type: type = int,
+) -> float | int:
+    """Prompt for a numeric value with validation.
+
+    Args:
+        label: The config item name.
+        description: One-line explanation.
+        default: Default value; returned on empty Enter.  None means required.
+        number_type: ``int`` or ``float``.
+
+    Returns:
+        The parsed number.
+    """
+    print()
+    print(f"{BOLD}{label}{RESET}")
+    if description:
+        print(f"  {description}")
+
+    while True:
+        if default is not None:
+            prompt = f"  {DIM}[default: {default}]{RESET} {GREEN}>{RESET} "
+        else:
+            prompt = f"  {GREEN}>{RESET} "
+        raw = input(prompt).strip()
+        if not raw and default is not None:
+            return default
+        try:
+            return number_type(raw)
+        except (ValueError, TypeError):
+            type_name = "integer" if number_type is int else "number"
+            print(f"  {YELLOW}Please enter a valid {type_name}.{RESET}")
+
+
+# ---------------------------------------------------------------------------
+# prompt_bool
+# ---------------------------------------------------------------------------
+
+
+def prompt_bool(
+    label: str,
+    description: str = "",
+    default: bool = True,
+) -> bool:
+    """Prompt for a yes/no confirmation.
+
+    Shows ``[Y/n]`` or ``[y/N]`` depending on the default.
+    Enter alone accepts the default.
+
+    Args:
+        label: The config item name.
+        description: One-line explanation.
+        default: Default value when user presses Enter.
+
+    Returns:
+        True for yes, False for no.
+    """
+    print()
+    print(f"{BOLD}{label}{RESET}")
+    if description:
+        print(f"  {description}")
+
+    hint = "[default: Y] (Y/n)" if default else "[default: N] (y/N)"
+    while True:
+        raw = input(f"  {DIM}{hint}{RESET} {GREEN}>{RESET} ").strip().lower()
+        if not raw:
+            return default
+        if raw in ("y", "yes"):
+            return True
+        if raw in ("n", "no"):
+            return False
+        print(f"  {YELLOW}Please enter Y or N.{RESET}")
+
+
+# ---------------------------------------------------------------------------
+# prompt_choice — arrow-key selection
+# ---------------------------------------------------------------------------
+
+
+def _read_key() -> str:
+    """Read a single keypress from stdin in raw mode.
+
+    Returns:
+        ``"up"``, ``"down"``, ``"enter"``, or the literal character.
+    """
+    # Standard
+    import termios
+    import tty
+
+    fd = sys.stdin.fileno()
+    old_settings = termios.tcgetattr(fd)
+    try:
+        tty.setraw(fd)
+        ch = sys.stdin.read(1)
+        if ch == "\r" or ch == "\n":
+            return "enter"
+        if ch == "\x1b":
+            seq = sys.stdin.read(2)
+            if seq == "[A":
+                return "up"
+            if seq == "[B":
+                return "down"
+            return "escape"
+        return ch
+    finally:
+        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
+
+
+def _render_choices(
+    choices: list[tuple[str, str]],
+    selected: int,
+) -> list[str]:
+    """Build display lines for the choice selector."""
+    lines: list[str] = []
+    for i, (value, desc) in enumerate(choices):
+        if i == selected:
+            marker = f"{CYAN}● {value}{RESET}"
+        else:
+            marker = f"  {value}"
+        # Pad value to align descriptions
+        padding = max(0, 18 - len(value))
+        line = f"  {marker}{' ' * padding} {DIM}{desc}{RESET}"
+        lines.append(line)
+    return lines
+
+
+def prompt_choice(
+    label: str,
+    description: str = "",
+    choices: list[tuple[str, str]] = [],  # noqa: B006
+    default: str = "",
+) -> str:
+    """Prompt user to select from a list using arrow keys.
+
+    Each choice is a ``(value, description)`` tuple.  The description is
+    shown next to the value as a brief explanation.
+
+    Falls back to numbered input on non-TTY stdin.
+
+    Args:
+        label: The config item name.
+        description: One-line explanation.
+        choices: List of ``(value, one_line_description)`` tuples.
+        default: Pre-selected value.  Defaults to the first choice.
+
+    Returns:
+        The selected value string.
+    """
+    if not choices:
+        raise ValueError("choices must not be empty")
+
+    # Find default index
+    selected = 0
+    for i, (val, _desc) in enumerate(choices):
+        if val == default:
+            selected = i
+            break
+
+    print()
+    print(f"{BOLD}{label}{RESET}")
+    if description:
+        print(f"  {description}")
+
+    # Non-TTY fallback: numbered list
+    if not _is_tty():
+        return _prompt_choice_fallback(choices, selected)
+
+    print(f"  {DIM}Use ↑↓ to navigate, Enter to select.{RESET}")
+    print()
+
+    # Initial render
+    lines = _render_choices(choices, selected)
+    for line in lines:
+        print(line)
+
+    while True:
+        key = _read_key()
+        if key == "up":
+            selected = (selected - 1) % len(choices)
+        elif key == "down":
+            selected = (selected + 1) % len(choices)
+        elif key == "enter":
+            # Clear and re-render final state
+            num_lines = len(choices)
+            sys.stdout.write(f"\r{CURSOR_UP * num_lines}")
+            lines = _render_choices(choices, selected)
+            for line in lines:
+                print(f"{CLEAR_LINE}{line}")
+            return choices[selected][0]
+        else:
+            continue
+
+        # Redraw
+        num_lines = len(choices)
+        sys.stdout.write(f"\r{CURSOR_UP * num_lines}")
+        lines = _render_choices(choices, selected)
+        for line in lines:
+            print(f"{CLEAR_LINE}{line}")
+
+
+def _prompt_choice_fallback(
+    choices: list[tuple[str, str]],
+    default_index: int,
+) -> str:
+    """Numbered fallback for non-TTY environments."""
+    for i, (val, desc) in enumerate(choices):
+        marker = "*" if i == default_index else " "
+        print(f"  {marker} {i + 1}) {val}  — {desc}")
+    while True:
+        raw = input(f"  [default: {default_index + 1}] {GREEN}>{RESET} ").strip()
+        if not raw:
+            return choices[default_index][0]
+        try:
+            idx = int(raw) - 1
+            if 0 <= idx < len(choices):
+                return choices[idx][0]
+        except ValueError:
+            pass
+        print(f"  {YELLOW}Enter a number 1-{len(choices)}.{RESET}")

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

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Real-time terminal display for ``lmcache bench engine``."""
+
+# Standard
+import threading
+import time
+
+# First Party
+from lmcache.cli.commands.bench.engine_bench.stats import StatsCollector
+
+# ANSI escape codes
+CLEAR_LINE = "\033[2K"
+CURSOR_UP = "\033[A"
+BOLD = "\033[1m"
+GREEN = "\033[92m"
+YELLOW = "\033[93m"
+CYAN = "\033[96m"
+RED = "\033[91m"
+RESET = "\033[0m"
+
+
+class ProgressMonitor:
+    """Real-time terminal display for benchmark progress.
+
+    Runs a daemon thread that redraws stats every second using ANSI
+    cursor control for in-place updates.  Reads aggregated metrics
+    from a ``StatsCollector`` and tracks in-flight count internally.
+    """
+
+    LOG_LINES = 5
+    DISPLAY_LINES = 9 + LOG_LINES  # stats block + log lines
+
+    def __init__(
+        self,
+        stats_collector: StatsCollector,
+        quiet: bool = False,
+    ) -> None:
+        self._stats_collector = stats_collector
+        self._quiet = quiet
+
+        self._lock = threading.Lock()
+        self._in_flight: int = 0
+        self._log_messages: list[str] = []
+
+        self._running: bool = False
+        self._thread: threading.Thread = None  # type: ignore[assignment]
+        self._first_draw: bool = True
+        self._start_time: float = time.monotonic()
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def start(self) -> None:
+        """Start the display daemon thread."""
+        if self._quiet:
+            return
+        self._start_time = time.monotonic()
+        self._running = True
+        self._first_draw = True
+        self._thread = threading.Thread(target=self._display_loop, daemon=True)
+        self._thread.start()
+
+    def stop(self) -> None:
+        """Stop the display thread and print final state."""
+        self._running = False
+        if self._thread is not None:
+            self._thread.join(timeout=3)
+        if not self._quiet:
+            self._draw()
+            print()  # final newline
+
+    def on_request_sent(self, request_id: str) -> None:
+        """Called when a request is dispatched. Increments in-flight."""
+        with self._lock:
+            self._in_flight += 1
+
+    def on_request_finished(self, request_id: str, successful: bool) -> None:
+        """Called when a request completes. Decrements in-flight."""
+        with self._lock:
+            self._in_flight = max(self._in_flight - 1, 0)
+
+    def log_message(self, message: str) -> None:
+        """Add a log line to the display. Thread-safe."""
+        with self._lock:
+            self._log_messages.append(message)
+            if len(self._log_messages) > self.LOG_LINES:
+                self._log_messages = self._log_messages[-self.LOG_LINES :]
+
+    # ------------------------------------------------------------------
+    # Internal
+    # ------------------------------------------------------------------
+
+    def _display_loop(self) -> None:
+        """Background loop: redraw every second."""
+        while self._running:
+            self._draw()
+            time.sleep(1.0)
+
+    def _draw(self) -> None:
+        """Render the real-time stats display."""
+        # Read stats from collector (already thread-safe)
+        stats = self._stats_collector.get_current_stats()
+
+        # Read local state under lock
+        with self._lock:
+            in_flight = self._in_flight
+            log_messages = list(self._log_messages)
+
+        elapsed = time.monotonic() - self._start_time
+
+        # Move cursor up to overwrite previous display
+        if not self._first_draw:
+            print(f"\r{CURSOR_UP * self.DISPLAY_LINES}", end="")
+        self._first_draw = False
+
+        lines = [
+            f"{BOLD}{'─' * 50}{RESET}",
+            (f"{BOLD} Engine Benchmark{RESET}   elapsed: {CYAN}{elapsed:.0f}s{RESET}"),
+            (
+                f"  {GREEN}{stats.successful_requests}{RESET} successful"
+                f"   {CYAN}{in_flight}{RESET} in-flight"
+                f"   {RED}{stats.failed_requests}{RESET} failed"
+            ),
+            f"  Avg TTFT:       {YELLOW}{stats.mean_ttft_ms:.1f} ms{RESET}",
+            (f"  Avg decode:     {YELLOW}{stats.mean_decode_speed:.1f} tok/s{RESET}"),
+            f"  Input tokens:   {stats.total_input_tokens:,}",
+            f"  Output tokens:  {stats.total_output_tokens:,}",
+            (
+                f"  Throughput:     "
+                f"{stats.input_throughput:.1f} in "
+                f"/ {stats.output_throughput:.1f} out tok/s"
+            ),
+            f"{BOLD}{'─' * 50}{RESET}",
+        ]
+
+        # Pad log lines to fixed count
+        for i in range(self.LOG_LINES):
+            if i < len(log_messages):
+                lines.append(f"  {YELLOW}[log] {log_messages[i]}{RESET}")
+            else:
+                lines.append("")
+
+        for line in lines:
+            print(f"{CLEAR_LINE}{line}")

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

@@ -0,0 +1,232 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Async streaming request sender for ``lmcache bench engine``."""
+
+# Standard
+from collections.abc import Callable
+import collections.abc
+import os
+import time
+
+# Third Party
+from openai import AsyncOpenAI
+
+# First Party
+from lmcache.cli.commands.bench.engine_bench.stats import RequestResult
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+# Callback signature: (result, response_text) -> None
+OnFinishedCallback = Callable[[RequestResult, str], None]
+
+
+def _normalize_url(engine_url: str) -> str:
+    """Ensure *engine_url* has a scheme and ends with ``/v1``."""
+    url = engine_url.rstrip("/")
+    if not url.startswith(("http://", "https://")):
+        url = f"http://{url}"
+    if not url.endswith("/v1"):
+        url += "/v1"
+    return url
+
+
+def _extract_content(chunk: object, completions_mode: bool) -> str:
+    """Return text content from a streaming chunk, or ``""`` if none.
+
+    Ported from Tensormesh-Benchmark ``streaming_utils.py``.
+    """
+    choices = getattr(chunk, "choices", None)
+    if not choices:
+        return ""
+
+    choice = choices[0]
+
+    if completions_mode:
+        text = getattr(choice, "text", None)
+        return text if text is not None else ""
+
+    # Chat mode: delta.content, with fallback for reasoning_content
+    delta = getattr(choice, "delta", None)
+    if delta is None:
+        return ""
+    content = getattr(delta, "content", None)
+    if content is not None:
+        return content
+    # Fallback for reasoning models
+    for attr in ("reasoning_content", "reasoning"):
+        fallback = getattr(delta, attr, None)
+        if fallback is not None:
+            return fallback
+    return ""
+
+
+class RequestSender:
+    """Async streaming request sender for inference engines.
+
+    Each ``send_request`` call is a self-contained coroutine.
+    Concurrency is controlled externally by the workload module.
+    """
+
+    def __init__(
+        self,
+        engine_url: str,
+        model: str,
+        completions_mode: bool = False,
+        on_finished: list[OnFinishedCallback] = [],  # noqa: B006
+    ) -> None:
+        self._model = model
+        self._completions_mode = completions_mode
+        self._on_finished = list(on_finished)
+
+        base_url = _normalize_url(engine_url)
+        api_key = os.getenv("OPENAI_API_KEY", "")
+        if not api_key:
+            api_key = "sk-dummy"
+            logger.debug("API key source: default dummy key")
+        else:
+            logger.debug("API key source: OPENAI_API_KEY env var")
+
+        self._client = AsyncOpenAI(
+            base_url=base_url,
+            api_key=api_key,
+            timeout=None,
+        )
+
+    def add_on_finished_callback(self, callback: OnFinishedCallback) -> None:
+        """Register a callback to be invoked when a request finishes."""
+        self._on_finished.append(callback)
+
+    async def send_request(
+        self,
+        request_id: str,
+        messages: list[dict[str, str]],
+        max_tokens: int = 128,
+    ) -> RequestResult:
+        """Send a single streaming request and return the result.
+
+        Streams the response via SSE, measures TTFT, decode speed, and
+        total latency.  Extracts token counts from server usage reports.
+        After collecting the result, invokes all registered
+        ``on_finished`` callbacks.
+        """
+        submit_time = time.time()
+        first_token_time = 0.0
+        tokens: list[str] = []
+        num_input_tokens = 0
+        num_output_tokens = 0
+
+        try:
+            response = await self._create_stream(messages, max_tokens)
+
+            async for chunk in response:
+                # Extract usage from final chunk
+                usage = getattr(chunk, "usage", None)
+                if usage is not None:
+                    pt = getattr(usage, "prompt_tokens", 0)
+                    ct = getattr(usage, "completion_tokens", 0)
+                    if pt:
+                        num_input_tokens = pt
+                    if ct:
+                        num_output_tokens = ct
+
+                content = _extract_content(chunk, self._completions_mode)
+                if content:
+                    if not first_token_time:
+                        first_token_time = time.time()
+                    tokens.append(content)
+
+            finish_time = time.time()
+            successful = first_token_time > 0.0
+            ttft = (first_token_time - submit_time) if successful else -1.0
+            request_latency = finish_time - submit_time
+            decode_time = (finish_time - first_token_time) if successful else 0.0
+            num_output = num_output_tokens if num_output_tokens > 0 else len(tokens)
+            decode_speed = (num_output / decode_time) if decode_time > 0 else 0.0
+
+            result = RequestResult(
+                request_id=request_id,
+                successful=successful,
+                ttft=ttft,
+                request_latency=request_latency,
+                num_input_tokens=num_input_tokens,
+                num_output_tokens=num_output,
+                decode_speed=decode_speed,
+                submit_time=submit_time,
+                first_token_time=first_token_time,
+                finish_time=finish_time,
+                error="",
+            )
+            response_text = "".join(tokens)
+
+        except Exception as e:
+            finish_time = time.time()
+            result = RequestResult(
+                request_id=request_id,
+                successful=False,
+                ttft=-1.0,
+                request_latency=finish_time - submit_time,
+                num_input_tokens=0,
+                num_output_tokens=0,
+                decode_speed=0.0,
+                submit_time=submit_time,
+                first_token_time=0.0,
+                finish_time=finish_time,
+                error=str(e),
+            )
+            response_text = ""
+            logger.debug(
+                "Request %s failed: %s",
+                request_id,
+                e,
+            )
+
+        for cb in self._on_finished:
+            cb(result, response_text)
+
+        return result
+
+    async def send_warmup_request(
+        self,
+        request_id: str,
+        messages: list[dict[str, str]],
+        max_tokens: int = 1,
+    ) -> RequestResult:
+        """Send a warmup request (``max_tokens=1`` by default)."""
+        return await self.send_request(
+            request_id,
+            messages,
+            max_tokens=max_tokens,
+        )
+
+    async def close(self) -> None:
+        """Close the underlying HTTP client."""
+        await self._client.close()
+
+    # ------------------------------------------------------------------
+    # Internal
+    # ------------------------------------------------------------------
+
+    async def _create_stream(
+        self,
+        messages: list[dict[str, str]],
+        max_tokens: int,
+    ) -> collections.abc.AsyncIterator:
+        """Dispatch the streaming API call (chat or completions)."""
+        if self._completions_mode:
+            prompt = messages[0]["content"] if messages else ""
+            return await self._client.completions.create(
+                model=self._model,
+                prompt=prompt,
+                stream=True,
+                max_tokens=max_tokens,
+                temperature=0.0,
+                stream_options={"include_usage": True},
+            )
+        return await self._client.chat.completions.create(
+            model=self._model,
+            messages=messages,
+            stream=True,
+            max_tokens=max_tokens,
+            temperature=0.0,
+            stream_options={"include_usage": True},
+        )