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

lmcache/cli/commands/bench/engine_bench/workloads/random_prefill.py

@@ -0,0 +1,178 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Random-prefill workload for ``lmcache bench engine``."""
+
+# Standard
+from dataclasses import dataclass
+import asyncio
+
+# First Party
+from lmcache.cli.commands.bench.engine_bench.progress import ProgressMonitor
+from lmcache.cli.commands.bench.engine_bench.request_sender import (
+    RequestSender,
+)
+from lmcache.cli.commands.bench.engine_bench.stats import StatsCollector
+from lmcache.cli.commands.bench.engine_bench.workloads.base import BaseWorkload
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class RandomPrefillConfig:
+    """Workload-specific config for the random-prefill workload."""
+
+    request_length: int = 10000
+    num_requests: int = 50
+
+    def __post_init__(self) -> None:
+        if self.request_length <= 0:
+            raise ValueError(
+                f"request_length must be positive, got {self.request_length}"
+            )
+        if self.num_requests < 1:
+            raise ValueError(f"num_requests must be >= 1, got {self.num_requests}")
+
+    @classmethod
+    def resolve(
+        cls,
+        request_length: int = 10000,
+        num_requests: int = 50,
+    ) -> "RandomPrefillConfig":
+        """Create a config from CLI args.
+
+        Unlike other workloads, random-prefill does not use the KV cache
+        budget to compute request count — the user specifies it directly.
+        """
+        return cls(
+            request_length=request_length,
+            num_requests=num_requests,
+        )
+
+
+class RandomPrefillWorkload(BaseWorkload):
+    """Workload that tests prefill speed by firing all requests at once.
+
+    Generates synthetic prompts of ``request_length`` tokens and dispatches
+    all ``num_requests`` simultaneously with ``max_tokens=1``.  There is
+    no warmup phase.
+    """
+
+    def __init__(
+        self,
+        config: RandomPrefillConfig,
+        request_sender: RequestSender,
+        stats_collector: StatsCollector,
+        progress_monitor: ProgressMonitor,
+        seed: int = 42,
+    ) -> None:
+        super().__init__(request_sender, stats_collector, progress_monitor)
+        self._config = config
+        self._seed = seed
+
+        self._prompts = self._generate_prompts()
+        self._dispatched = False
+        self._pending_tasks: set[asyncio.Task] = set()
+
+    def log_config(self) -> None:
+        """Log key workload config before the benchmark starts."""
+        c = self._config
+        B = "\033[1m"  # bold
+        C = "\033[96m"  # cyan
+        Y = "\033[93m"  # yellow
+        R = "\033[0m"  # reset
+        print(
+            f"{B}{'═' * 50}{R}\n"
+            f"{B} Workload: {C}random-prefill{R}\n"
+            f"{B}{'─' * 50}{R}\n"
+            f"  Requests:         {Y}{c.num_requests}{R}\n"
+            f"  Request length:   {Y}{c.request_length}{R} tokens\n"
+            f"  Max output:       {Y}1{R} token\n"
+            f"{B}{'═' * 50}{R}"
+        )
+
+    # ------------------------------------------------------------------
+    # Prompt generation
+    # ------------------------------------------------------------------
+
+    def _generate_prompts(self) -> list[str]:
+        """Generate synthetic prompts of approximately ``request_length`` tokens."""
+        prompts: list[str] = []
+        for i in range(self._config.num_requests):
+            prefix = f"Request {i}: "
+            body = " ".join(["hi"] * max(self._config.request_length - 10, 1))
+            prompts.append(prefix + body)
+        logger.debug(
+            "Generated %d prompts of ~%d tokens each",
+            len(prompts),
+            self._config.request_length,
+        )
+        return prompts
+
+    # ------------------------------------------------------------------
+    # Warmup
+    # ------------------------------------------------------------------
+
+    async def warmup(self) -> None:
+        """No warmup for random-prefill."""
+
+    # ------------------------------------------------------------------
+    # Benchmark dispatch
+    # ------------------------------------------------------------------
+
+    async def step(self, time_offset: float) -> float:
+        """Dispatch all requests at once on the first call.
+
+        Returns:
+            0.0 while tasks are pending, -1.0 when all done.
+        """
+        if not self._dispatched:
+            self._dispatched = True
+            for i, prompt in enumerate(self._prompts):
+                request_id = f"prefill_{i}"
+                messages = [{"role": "user", "content": prompt}]
+                self._progress_monitor.on_request_sent(request_id)
+
+                task = asyncio.create_task(
+                    self._dispatch(request_id, messages),
+                )
+                self._pending_tasks.add(task)
+                task.add_done_callback(self._on_task_done)
+
+            self._progress_monitor.log_message(
+                f"Dispatched all {self._config.num_requests} requests"
+            )
+            return 0.0
+
+        # Wait for pending tasks
+        if self._pending_tasks:
+            await asyncio.wait(
+                self._pending_tasks,
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+            return 0.0
+
+        return -1.0
+
+    async def _dispatch(
+        self,
+        request_id: str,
+        messages: list[dict[str, str]],
+    ) -> None:
+        """Send a single prefill request with max_tokens=1."""
+        await self._request_sender.send_request(
+            request_id,
+            messages,
+            max_tokens=1,
+        )
+
+    def _on_task_done(self, task: asyncio.Task) -> None:
+        """Clean up completed tasks and log unexpected errors."""
+        self._pending_tasks.discard(task)
+        if not task.cancelled():
+            exc = task.exception()
+            if exc is not None:
+                self._progress_monitor.log_message(f"Dispatch task failed: {exc}")
+
+    def on_request_finished(self, request_id: str, output: str) -> None:
+        """No-op — this workload is stateless."""
+        self._progress_monitor.log_message(f"Request {request_id} finished")

lmcache/cli/commands/describe.py

@@ -0,0 +1,310 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``lmcache describe`` — show detailed status of a running LMCache service.
+
+Usage::
+
+    lmcache describe kvcache --url http://localhost:8000
+"""
+
+# Standard
+import argparse
+import json
+import sys
+import urllib.error
+import urllib.request
+
+# First Party
+from lmcache.cli.commands.base import BaseCommand
+from lmcache.cli.metrics import Metrics
+
+# -------------------------------------------------------------------
+# Shared helpers
+# -------------------------------------------------------------------
+
+
+class DescribeError(Exception):
+    """Raised when the describe command cannot fetch or parse status data."""
+
+
+def normalize_url(url: str) -> str:
+    """Ensure *url* has an ``http://`` or ``https://`` scheme."""
+    if not url.startswith(("http://", "https://")):
+        url = f"http://{url}"
+    return url.rstrip("/")
+
+
+def fetch_json(url: str, timeout: int = 10) -> dict:
+    """GET *url* and return the parsed JSON body.
+
+    Raises:
+        DescribeError: On network/HTTP errors.
+    """
+    req = urllib.request.Request(url)
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            return json.loads(resp.read().decode())
+    except urllib.error.HTTPError as exc:
+        if exc.code == 503:
+            body = exc.read().decode()
+            try:
+                detail = json.loads(body).get("error", body)
+            except (json.JSONDecodeError, AttributeError):
+                detail = body
+            raise DescribeError(f"Server unhealthy: {detail}") from exc
+        raise DescribeError(f"HTTP {exc.code} from {url}: {exc.reason}") from exc
+    except urllib.error.URLError as exc:
+        raise DescribeError(f"Cannot connect to {url}: {exc.reason}") from exc
+    except OSError as exc:
+        raise DescribeError(f"Cannot connect to {url}: {exc}") from exc
+
+
+def fmt_bytes(n: int) -> str:
+    """Format a byte count as a human-readable string."""
+    if n >= 1024**3:
+        return f"{n / 1024**3:.2f} GB"
+    if n >= 1024**2:
+        return f"{n / 1024**2:.2f} MB"
+    if n >= 1024:
+        return f"{n / 1024:.2f} KB"
+    return f"{n} B"
+
+
+def fmt_health(is_healthy: object) -> str | None:
+    """Format a boolean health flag as ``'OK'`` / ``'UNHEALTHY'``."""
+    if is_healthy is None:
+        return None
+    return "OK" if is_healthy else "UNHEALTHY"
+
+
+def safe_get(data: dict, *keys, default=None):  # type: ignore[type-arg]
+    """Walk nested dicts by *keys*, returning *default* on any miss."""
+    cur: object = data
+    for key in keys:
+        if not isinstance(cur, dict):
+            return default
+        cur = cur.get(key)
+        if cur is None:
+            return default
+    return cur
+
+
+# -------------------------------------------------------------------
+# KVCache describer
+# -------------------------------------------------------------------
+
+
+class KVCacheDescriber:
+    """Builds the ``describe kvcache`` output from a ``/api/status`` response.
+
+    Each ``add_*`` method populates one logical section. The orchestrating
+    :meth:`describe` calls them in order and emits the result.  Adding a
+    new section is a one-method change — no other code needs to know
+    about it.
+    """
+
+    def __init__(self, metrics: Metrics, data: dict, base_url: str) -> None:
+        self.metrics = metrics
+        self.data = data
+        self.base_url = base_url
+
+    def describe(self) -> None:
+        """Run all section builders and emit."""
+        self.add_overview()
+        self.add_l1_storage()
+        self.add_models()
+        self.add_l2_adapters()
+        self.metrics.emit()
+
+    # -- sections --------------------------------------------------------
+
+    def add_overview(self) -> None:
+        """Top-level engine overview."""
+        self.metrics.add("health", "Health", fmt_health(self.data.get("is_healthy")))
+        self.metrics.add("url", "URL", self.base_url)
+        self.metrics.add("engine_type", "Engine type", self.data.get("engine_type"))
+        self.metrics.add("chunk_size", "Chunk size", self.data.get("chunk_size"))
+
+    def add_l1_storage(self) -> None:
+        """L1 cache capacity, usage, eviction, and object count."""
+        total_bytes = safe_get(
+            self.data, "storage_manager", "l1_manager", "memory_total_bytes"
+        )
+        if total_bytes is not None:
+            self.metrics.add(
+                "l1_capacity_gb",
+                "L1 capacity (GB)",
+                round(total_bytes / (1024**3), 2),
+            )
+        else:
+            self.metrics.add("l1_capacity_gb", "L1 capacity (GB)", None)
+
+        used_bytes = safe_get(
+            self.data, "storage_manager", "l1_manager", "memory_used_bytes"
+        )
+        usage_ratio = safe_get(
+            self.data, "storage_manager", "l1_manager", "memory_usage_ratio"
+        )
+        if used_bytes is not None and usage_ratio is not None:
+            gb = used_bytes / (1024**3)
+            pct = usage_ratio * 100
+            self.metrics.add("l1_used_gb", "L1 used (GB)", f"{gb:.2f} ({pct:.1f}%)")
+        else:
+            self.metrics.add("l1_used_gb", "L1 used (GB)", None)
+
+        self.metrics.add(
+            "eviction_policy",
+            "Eviction policy",
+            safe_get(
+                self.data,
+                "storage_manager",
+                "eviction_controller",
+                "eviction_policy",
+            ),
+        )
+        self.metrics.add(
+            "cached_objects",
+            "Cached objects",
+            safe_get(self.data, "storage_manager", "l1_manager", "total_object_count"),
+        )
+        self.metrics.add(
+            "active_sessions", "Active sessions", self.data.get("active_sessions")
+        )
+
+    def add_models(self) -> None:
+        """Per-model KV cache layout sections."""
+        gpu_meta = self.data.get("gpu_context_meta", {})
+        if not gpu_meta:
+            return
+
+        # Deduplicate by (model_name, world_size) — multiple GPU IDs
+        # may share the same model.
+        seen: dict[tuple[str, int], dict] = {}
+        for gpu_id, meta in gpu_meta.items():
+            key = (meta["model_name"], meta["world_size"])
+            if key not in seen:
+                seen[key] = {
+                    "gpu_ids": [],
+                    "layout": meta.get("kv_cache_layout"),
+                }
+            seen[key]["gpu_ids"].append(gpu_id)
+
+        for idx, ((model_name, world_size), info) in enumerate(seen.items()):
+            section_key = f"model_{idx}"
+            self.metrics.add_list_section("models", section_key, f"Model: {model_name}")
+            sec = self.metrics[section_key]
+            sec.add("model", "Model", model_name)
+            sec.add("world_size", "World size", world_size)
+            sec.add("gpu_ids", "GPU IDs", ", ".join(info["gpu_ids"]))
+
+            layout = info.get("layout")
+            if not layout:
+                continue
+            sec.add(
+                "attention_backend",
+                "Attention backend",
+                layout.get("attention_backend"),
+            )
+            sec.add("gpu_kv_shape", "GPU KV shape", layout.get("gpu_kv_shape"))
+            sec.add(
+                "gpu_kv_concrete_shape",
+                "GPU KV tensor shape",
+                layout.get("gpu_kv_concrete_shape"),
+            )
+            sec.add("num_layers", "Num layers", layout["num_layers"])
+            sec.add("block_size", "Block size", layout["block_size"])
+            sec.add("hidden_dim_sizes", "Hidden dim sizes", layout["hidden_dim_sizes"])
+            sec.add("dtype", "Dtype", layout["dtype"])
+            sec.add("is_mla", "MLA", layout["is_mla"])
+            sec.add("num_blocks", "Num blocks", layout["num_blocks"])
+            sec.add(
+                "cache_size_per_token",
+                "Cache size per token (bytes)",
+                layout["cache_size_per_token"],
+            )
+
+    def add_l2_adapters(self) -> None:
+        """L2 adapter sections."""
+        l2_adapters = safe_get(self.data, "storage_manager", "l2_adapters") or []
+        for idx, adapter in enumerate(l2_adapters):
+            adapter_type = adapter.get("type", "Unknown")
+            section_key = f"l2_{idx}"
+            self.metrics.add_list_section(
+                "l2_adapters", section_key, f"L2: {adapter_type}"
+            )
+            sec = self.metrics[section_key]
+            sec.add("type", "Type", adapter_type)
+            sec.add("health", "Health", fmt_health(adapter.get("is_healthy")))
+
+            if "backend" in adapter:
+                sec.add("backend", "Backend", adapter["backend"])
+            if "base_path" in adapter:
+                sec.add("base_path", "Base path", adapter["base_path"])
+            if "stored_object_count" in adapter:
+                sec.add(
+                    "stored_object_count",
+                    "Stored objects",
+                    adapter["stored_object_count"],
+                )
+
+            cap = adapter.get("max_capacity_bytes")
+            used = adapter.get("current_size_bytes")
+            if cap is not None and used is not None:
+                pct = used / cap * 100 if cap > 0 else 0.0
+                sec.add(
+                    "used",
+                    "Used",
+                    f"{fmt_bytes(used)} / {fmt_bytes(cap)} ({pct:.1f}%)",
+                )
+
+            pool_size = adapter.get("pool_size")
+            pool_free = adapter.get("pool_free_slots")
+            if pool_size is not None and pool_free is not None:
+                pool_used = pool_size - pool_free
+                pct = pool_used / pool_size * 100 if pool_size > 0 else 0.0
+                sec.add(
+                    "pool_used",
+                    "Pool used",
+                    f"{pool_used} / {pool_size} ({pct:.1f}%)",
+                )
+
+
+# -------------------------------------------------------------------
+# Command
+# -------------------------------------------------------------------
+
+
+class DescribeCommand(BaseCommand):
+    """Show detailed status of a running LMCache service."""
+
+    def name(self) -> str:
+        return "describe"
+
+    def help(self) -> str:
+        return "Show detailed status of a running LMCache service."
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        parser.add_argument(
+            "target",
+            choices=["kvcache"],
+            help="What to describe.",
+        )
+        parser.add_argument(
+            "--url",
+            help="LMCache HTTP server URL (default to http://localhost:8080).",
+            default="http://localhost:8080",
+        )
+
+    def execute(self, args: argparse.Namespace) -> None:
+        if args.target == "kvcache":
+            self._describe_kvcache(args)
+
+    def _describe_kvcache(self, args: argparse.Namespace) -> None:
+        base_url = normalize_url(args.url)
+        try:
+            data = fetch_json(f"{base_url}/api/status")
+        except DescribeError as exc:
+            print(str(exc), file=sys.stderr)
+            sys.exit(1)
+
+        metrics = self.create_metrics("LMCache KV Cache Service", args, width=50)
+        KVCacheDescriber(metrics, data, base_url).describe()

lmcache/cli/commands/kvcache.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``lmcache kvcache`` — KV cache management.
+
+Sub-commands:
+    clear         Clear all cached KV data
+"""
+
+# Standard
+from typing import Any, Optional
+import argparse
+import json
+import sys
+import urllib.error
+import urllib.request
+
+# First Party
+from lmcache.cli.commands.base import BaseCommand
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+def _http_request(
+    method: str, url: str, data: Optional[dict[str, Any]] = None
+) -> dict[str, Any]:
+    """Send an HTTP request and return the parsed JSON response.
+
+    Args:
+        method: HTTP method (GET, POST, PUT, DELETE).
+        url: Full URL to request.
+        data: Optional JSON body to send.
+
+    Returns:
+        Parsed JSON response as a dict.
+
+    Raises:
+        SystemExit: On connection error or non-2xx HTTP response.
+    """
+    body = None
+    headers: dict[str, str] = {}
+    if data is not None:
+        body = json.dumps(data).encode()
+        headers["Content-Type"] = "application/json"
+
+    req = urllib.request.Request(url, data=body, headers=headers, method=method)
+    try:
+        with urllib.request.urlopen(req, timeout=10) as resp:
+            return json.loads(resp.read().decode())
+    except urllib.error.HTTPError as e:
+        try:
+            error_body = json.loads(e.read().decode())
+            msg = error_body.get("message") or error_body.get("error") or str(e)
+        except (json.JSONDecodeError, ValueError, OSError):
+            msg = str(e)
+        logger.error("Server error: %s", msg)
+        sys.exit(1)
+    except urllib.error.URLError as e:
+        logger.error("Cannot reach %s — is the server running? (%s)", url, e.reason)
+        sys.exit(1)
+
+
+class KVCacheCommand(BaseCommand):
+    """Manage KV cache state.
+
+    This command provides sub-commands for managing KV cache data on the
+    MP HTTP server. Currently supports clearing all L1 cache. Future
+    sub-commands (pin, compress, info) are defined in the design doc.
+    """
+
+    def name(self) -> str:
+        """Return the subcommand name.
+
+        Returns:
+            The string ``"kvcache"``.
+        """
+        return "kvcache"
+
+    def help(self) -> str:
+        """Return short help text.
+
+        Returns:
+            Help string shown by ``lmcache -h``.
+        """
+        return "Manage KV cache state."
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        """Add kvcache-specific sub-commands and their arguments.
+
+        Args:
+            parser: The ``ArgumentParser`` for this subcommand.
+        """
+        sub = parser.add_subparsers(dest="action")
+
+        # clear
+        p_clear = sub.add_parser("clear", help="Clear all cached KV data in L1 (CPU).")
+        p_clear.add_argument(
+            "--url",
+            type=str,
+            required=True,
+            help="Target MP HTTP endpoint (e.g. http://localhost:8000).",
+        )
+
+    def execute(self, args: argparse.Namespace) -> None:
+        """Dispatch to the appropriate sub-command handler.
+
+        Args:
+            args: Parsed CLI arguments. Must contain ``action`` indicating
+                which sub-command was invoked.
+        """
+        action = getattr(args, "action", None)
+        if action is None:
+            logger.error("No sub-command specified. Run: lmcache kvcache -h")
+            sys.exit(1)
+
+        dispatch = {
+            "clear": self._clear,
+        }
+        dispatch[action](args)
+
+    def _clear(self, args: argparse.Namespace) -> None:
+        """Clear all cached KV data via the MP HTTP server."""
+        url = args.url.rstrip("/")
+
+        # MP HTTP server endpoint: POST /api/clear-cache
+        _http_request("POST", f"{url}/api/clear-cache")
+
+        quiet = getattr(args, "quiet", False)
+        if quiet:
+            return
+
+        metrics = self.create_metrics("KV Cache Clear", args)
+        metrics.add("status", "Status", "OK")
+        metrics.emit()

lmcache/cli/commands/mock.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``lmcache mock`` — example command that demonstrates the CLI framework.
+
+This command does not connect to any server.  It generates fake metrics
+to exercise argument parsing, the ``Metrics`` logger, and both terminal
+and JSON output paths.
+"""
+
+# Standard
+import argparse
+
+# First Party
+from lmcache.cli.commands.base import BaseCommand
+
+
+class MockCommand(BaseCommand):
+    """Mock command used as a reference implementation for new commands."""
+
+    def name(self) -> str:
+        """Return the subcommand name.
+
+        Returns:
+            The string ``"mock"``.
+        """
+        return "mock"
+
+    def help(self) -> str:
+        """Return short help text.
+
+        Returns:
+            Help string shown by ``lmcache -h``.
+        """
+        return "Run a mock command (example/test)."
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        """Add mock-specific arguments.
+
+        Args:
+            parser: The ``ArgumentParser`` for this subcommand.
+        """
+        parser.add_argument(
+            "--name",
+            type=str,
+            default="default",
+            help="Name tag for this mock run.",
+        )
+        parser.add_argument(
+            "--num-items",
+            type=int,
+            default=10,
+            help="Number of fake items to process.",
+        )
+
+    def execute(self, args: argparse.Namespace) -> None:
+        """Execute the mock command.
+
+        Args:
+            args: Parsed CLI arguments containing ``name``, ``num_items``,
+                and optionally ``output``.
+        """
+        metrics = self.create_metrics("Mock Result", args, width=40)
+
+        metrics.add_section("input", "Input Parameters")
+        metrics["input"].add("name", "Name", args.name)
+        metrics["input"].add("num_items", "Num items", args.num_items)
+
+        metrics.add_section("mock", "Mock Metrics")
+        metrics["mock"].add("items_processed", "Items processed", 42)
+        metrics["mock"].add("total_time_ms", "Total time (ms)", 12.34)
+        metrics["mock"].add("throughput", "Throughput (items/s)", 3403.73)
+
+        metrics.add_section("validation", "Validation")
+        metrics["validation"].add("status", "Status", "OK")
+
+        metrics.emit()